From c70e745a76ae1b767be0ba43dc1dbf19e43c3b9e Mon Sep 17 00:00:00 2001 From: dardwo Date: Sat, 3 Jun 2023 22:11:48 +0200 Subject: [PATCH] neural network --- main.py | 16 ++--- src/cnn_model.py | 132 ----------------------------------------- src/neural_networks.py | 74 +++++++++++++++++++++++ 3 files changed, 82 insertions(+), 140 deletions(-) delete mode 100644 src/cnn_model.py create mode 100644 src/neural_networks.py diff --git a/main.py b/main.py index 09cbbf5..f2f43e7 100644 --- a/main.py +++ b/main.py @@ -11,11 +11,11 @@ import pickle import os from src.ID3 import make_decision import torch -import cnn_model +import neural_networks def recognize_plants(plants_array): checkpoint = torch.load(f'plants.model') - model = cnn_model.Net(num_classes=3) + model = neural_networks.Net(num_classes=3) model.load_state_dict(checkpoint) model.eval() img = '' @@ -29,15 +29,15 @@ def recognize_plants(plants_array): for j in range(11): if plants_array[j][i] == 'carrot': img = 'assets/learning/test/carrot/' + str(random.randint(1, 200)) + '.jpg' - pred = cnn_model.prediction(img, model) + pred = neural_networks.prediction(img, model) #show_plant_img(img) elif plants_array[j][i] == 'potato': img = 'assets/learning/test/potato/' + str(random.randint(1, 200)) + '.jpg' - pred = cnn_model.prediction(img, model) + pred = neural_networks.prediction(img, model) # show_plant_img(img) elif plants_array[j][i] == 'wheat': img = 'assets/learning/test/wheat/' + str(random.randint(1, 200)) + '.jpg' - pred = cnn_model.prediction(img, model) + pred = neural_networks.prediction(img, model) # show_plant_img(img) else: pred = 'none' @@ -50,15 +50,15 @@ def recognize_plants(plants_array): for j in range(10,-1,-1): if plants_array[j][i] == 'carrot': img = 'assets/learning/test/carrot/' + str(random.randint(1, 200)) + '.jpg' - pred = cnn_model.prediction(img, model) + pred = neural_networks.prediction(img, model) # show_plant_img(img) elif plants_array[j][i] == 'potato': img = 'assets/learning/test/potato/' + str(random.randint(1, 200)) + '.jpg' - pred = cnn_model.prediction(img, model) + pred = neural_networks.prediction(img, model) # show_plant_img(img) elif plants_array[j][i] == 'wheat': img = 'assets/learning/test/wheat/' + str(random.randint(1, 200)) + '.jpg' - pred = cnn_model.prediction(img, model) + pred = neural_networks.prediction(img, model) # show_plant_img(img) else: pred = 'none' diff --git a/src/cnn_model.py b/src/cnn_model.py deleted file mode 100644 index d910cd5..0000000 --- a/src/cnn_model.py +++ /dev/null @@ -1,132 +0,0 @@ -import torch -import torchvision -from PIL import Image -import torch.nn as nn -from torch.optim import Adam -from torch.autograd import Variable -from torch.utils.data import DataLoader - - -device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') -classes = ['carrot', 'potato', 'wheat'] - -train_path = ('assets/learning/train') - -test_path = ('assets/learning/test') - -transformer = torchvision.transforms.Compose([ - torchvision.transforms.Resize((150, 150)), - torchvision.transforms.ToTensor(), - torchvision.transforms.Normalize([0.5, 0.5, 0.5], - [0.5, 0.5, 0.5]) -]) - -class Net(nn.Module): - def __init__(self, num_classes=6): - super(Net, self).__init__() - self.conv1 = nn.Conv2d(in_channels=3, out_channels=12, kernel_size=3, stride=1, padding=1) - self.bn1 = nn.BatchNorm2d(num_features=12) - self.relu1 = nn.ReLU() - - self.pool = nn.MaxPool2d(kernel_size=2) - - self.conv2 = nn.Conv2d(in_channels=12, out_channels=20, kernel_size=3, stride=1, padding=1) - self.relu2 = nn.ReLU() - - self.conv3 = nn.Conv2d(in_channels=20, out_channels=32, kernel_size=3, stride=1, padding=1) - self.bn3 = nn.BatchNorm2d(num_features=32) - self.relu3 = nn.ReLU() - - self.fc = nn.Linear(in_features=75 * 75 * 32, out_features=num_classes) - - def forward(self, input): - output = self.conv1(input) - output = self.bn1(output) - output = self.relu1(output) - - output = self.pool(output) - - output = self.conv2(output) - output = self.relu2(output) - - output = self.conv3(output) - output = self.bn3(output) - output = self.relu3(output) - - output = output.view(-1, 32 * 75 * 75) - output = self.fc(output) - - return output - - -def train(dataloader, model: Net, optimizer: Adam, loss_fn: nn.CrossEntropyLoss): - size = len(dataloader.dataset) - for batch, (X, y) in enumerate(dataloader): - X, y = X.to(device), y.to(device) - - pred = model(X.float()) - loss = loss_fn(pred, y) - - optimizer.zero_grad() - loss.backward() - optimizer.step() - - if batch % 5 == 0: - loss, current = loss.item(), batch * len(X) - print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]") - - -def test(dataloader, model: Net, loss_fn: nn.CrossEntropyLoss): - size = len(dataloader.dataset) - model.eval() - test_loss, correct = 0, 0 - - with torch.no_grad(): - for X, y in dataloader: - X, y = X.to(device), y.to(device) - pred = model(X.float()) - test_loss += loss_fn(pred, y).item() - correct += (pred.argmax(1) == y).type(torch.float).sum().item() - - test_loss /= size - correct /= size - print(f"Test Error: \n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \n") - - -def prediction(img_path, model: Net): - image = Image.open(img_path).convert('RGB') - image_tensor = transformer(image).float() - image_tensor = image_tensor.unsqueeze_(0) - - if torch.cuda.is_available(): - image_tensor.cuda() - input = Variable(image_tensor) - output = model(input) - - index = output.data.numpy().argmax() - pred = classes[index] - return pred - - -def learn(): - num_epochs = 50 - - train_loader = DataLoader( - torchvision.datasets.ImageFolder(train_path, transform=transformer), - batch_size=64, shuffle=True - ) - test_loader = DataLoader( - torchvision.datasets.ImageFolder(test_path, transform=transformer), - batch_size=32, shuffle=True - ) - - model = Net(3).to(device) - optimizer = Adam(model.parameters(), lr=1e-3, weight_decay=0.0001) - loss_fn = nn.CrossEntropyLoss() - - for t in range(num_epochs): - print(f"Epoch {t + 1}\n-------------------------------") - train(train_loader, model, optimizer, loss_fn) - test(test_loader, model, loss_fn) - print("Done!") - torch.save(model.state_dict(), f'plants.model') diff --git a/src/neural_networks.py b/src/neural_networks.py new file mode 100644 index 0000000..fb32276 --- /dev/null +++ b/src/neural_networks.py @@ -0,0 +1,74 @@ +import torch +import torchvision +from PIL import Image +import torch.nn as nn +from torch.optim import Adam +from torch.autograd import Variable +from torch.utils.data import DataLoader + +device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +classes = ['carrot', 'potato', 'wheat'] + +train_path = 'assets/learning/train' +test_path = 'assets/learning/test' + +transformer = torchvision.transforms.Compose([ + torchvision.transforms.Resize((150, 150)), + torchvision.transforms.ToTensor(), + torchvision.transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]) +]) + +class Net(nn.Module): + def __init__(self, num_classes=3): + super(Net, self).__init__() + self.features = nn.Sequential( + nn.Conv2d(3, 12, kernel_size=3, stride=1, padding=1), + nn.BatchNorm2d(12), + nn.ReLU(), + nn.MaxPool2d(kernel_size=2), + nn.Conv2d(12, 20, kernel_size=3, stride=1, padding=1), + nn.ReLU(), + nn.Conv2d(20, 32, kernel_size=3, stride=1, padding=1), + nn.BatchNorm2d(32), + nn.ReLU() + ) + self.classifier = nn.Linear(32 * 75 * 75, num_classes) + + def forward(self, x): + x = self.features(x) + x = x.view(x.size(0), -1) + x = self.classifier(x) + return x + +def train(dataloader, model, optimizer, loss_fn): + model.train() + size = len(dataloader.dataset) + for batch, (X, y) in enumerate(dataloader): + X, y = X.to(device), y.to(device) + + optimizer.zero_grad() + pred = model(X.float()) + loss = loss_fn(pred, y) + loss.backward() + optimizer.step() + + if batch % 5 == 0: + current = batch * len(X) + print(f"loss: {loss.item():>7f} [{current:>5d}/{size:>5d}]") + +def test(dataloader, model, loss_fn): + model.eval() + size = len(dataloader.dataset) + test_loss, correct = 0, 0 + + with torch.no_grad(): + for X, y in dataloader: + X, y = X.to(device), y.to(device) + pred = model(X.float()) + test_loss += loss_fn(pred, y).item() + correct += (pred.argmax(1) == y).sum().item() + + test_loss /= size + accuracy = 100.0 * correct / size + print(f"Test Error:\n Accuracy: {accuracy:.1f}%, Avg loss: {test_loss:.8f}\n") +