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'

src/cnn_model.py

@@ -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')

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")
+