First neurNet prototype

NewralNetwork/neural_network_learning.py

@@ -0,0 +1,60 @@
+import glob
+from src.torchvision_resize_dataset import combined_dataset, images_path, classes
+import src.data_model
+from torch.optim import Adam
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+train_loader = DataLoader(
+    combined_dataset,   #dataset of images
+    batch_size=256,     # accuracy
+    shuffle=True    # rand order
+)
+
+model = src.data_model.DataModel(num_objects=2).to(device)
+
+#optimizer
+optimizer = Adam(model.parameters(), lr=0.001, weight_decay=0.0001)
+#loss function
+criterion = nn.CrossEntropyLoss()
+
+num_epochs = 20
+# train_size = len(glob.glob(images_path+'*.jpg'))
+train_size = 2000
+
+go_to_accuracy = 0.0
+for epoch in range(num_epochs):
+    #training on dataset
+    model.train()
+    train_accuracy = 0.0
+    train_loss = 0.0
+    for i, (images, labels) in enumerate(train_loader):
+        if torch.cuda.is_available():
+            images = torch.Variable(images.cuda())
+            labels = torch.Variable(labels.cuda())
+        # clearing the optimizer gradients
+        optimizer.zero_grad()
+
+        outputs = model(images) # prediction
+        loss = criterion(outputs, labels)   #loss calculation
+        loss.backward()
+        optimizer.step()
+
+        train_loss += loss.cpu().data*images.size(0)
+        _, prediction = torch.max(outputs.data, 1)
+
+        train_accuracy += int(torch.sum(prediction == labels.data))
+
+    train_accuracy = train_accuracy/train_size
+    train_loss = train_loss/train_size
+
+    model.eval()
+
+    print('Epoch:  '+ str(epoch+1) +' Train Loss: '+ str(int(train_loss)) +' Train Accuracy: '+ str(train_accuracy))
+
+    if train_accuracy > go_to_accuracy:
+        go_to_accuracy= train_accuracy
+        torch.save(model.state_dict(), "best_model.pth")

NewralNetwork/prediction.py

@@ -0,0 +1,146 @@
+import torch
+import torch.nn as nn
+from torchvision.transforms import transforms
+import numpy as np
+from torch.autograd import Variable
+from torchvision.models import squeezenet1_1
+import torch.functional as F
+from io import open
+import os
+from PIL import Image
+import pathlib
+import glob
+from tkinter import Tk, Label
+from PIL import Image, ImageTk
+
+train_path = os.path.abspath('src/train_images')
+pred_path = os.path.abspath('src/test_images')
+
+root = pathlib.Path(train_path)
+classes = sorted([j.name.split('/')[-1] for j in root.iterdir()])
+
+
+class DataModel(nn.Module):
+    def __init__(self, num_classes):
+        super(DataModel, self).__init__()
+        # input (batch=256, nr of channels rgb=3 , size=244x244)
+
+        # convolution
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=12, kernel_size=3, stride=1, padding=1)
+        # shape (256, 12, 224x224)
+
+        # batch normalization
+        self.bn1 = nn.BatchNorm2d(num_features=12)
+        # shape (256, 12, 224x224)
+        self.reul1 = nn.ReLU()
+
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+        # reduce image size by factor 2
+        # pooling window moves by 2 pixels at a time instead of 1
+        # shape (256, 12, 112x112)
+
+        self.conv2 = nn.Conv2d(in_channels=12, out_channels=24, kernel_size=3, stride=1, padding=1)
+        self.bn2 = nn.BatchNorm2d(num_features=24)
+        self.reul2 = nn.ReLU()
+        # shape (256, 24, 112x112)
+
+        self.conv3 = nn.Conv2d(in_channels=24, out_channels=48, kernel_size=3, stride=1, padding=1)
+        # shape (256, 48, 112x112)
+        self.bn3 = nn.BatchNorm2d(num_features=48)
+        # shape (256, 48, 112x112)
+        self.reul3 = nn.ReLU()
+
+        # connected layer
+        self.fc = nn.Linear(in_features=48 * 112 * 112, out_features=num_classes)
+
+    def forward(self, input):
+        output = self.conv1(input)
+        output = self.bn1(output)
+        output = self.reul1(output)
+
+        output = self.pool(output)
+        output = self.conv2(output)
+        output = self.bn2(output)
+        output = self.reul2(output)
+
+        output = self.conv3(output)
+        output = self.bn3(output)
+        output = self.reul3(output)
+
+        # output shape matrix (256, 48, 112x112)
+        # print(output.shape)
+        # print(self.fc.weight.shape)
+
+        output = output.view(-1, 48 * 112 * 112)
+        output = self.fc(output)
+
+        return output
+
+
+script_dir = os.path.dirname(os.path.abspath(__file__))
+file_path = os.path.join(script_dir, 'best_model.pth')
+checkpoint = torch.load(file_path)
+model = DataModel(num_classes=2)
+model.load_state_dict(checkpoint)
+model.eval()
+
+transformer = transforms.Compose([
+    transforms.Resize((224, 224)),  # Resize images to (224, 224)
+    transforms.ToTensor(),  # Convert images to tensors, 0-255 to 0-1
+    # transforms.RandomHorizontalFlip(),  # 0.5 chance to flip the image
+    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+])
+
+
+def prediction(img_path, transformer):
+    image = Image.open(img_path)
+
+    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 prediction_keys():
+    # funkcja zwracajaca sciezki do kazdego pliku w folderze w postaci listy
+
+    images_path = glob.glob(pred_path + '/*.jpg')
+
+    pred_list = []
+
+    for i in images_path:
+        pred_list.append(i)
+
+    return pred_list
+
+
+def predict_one(path):
+    # wyswietlanie obrazka po kazdym podniesieniu itemu
+    root = Tk()
+    root.title("Okno z obrazkiem")
+
+    image = Image.open(path)
+    photo = ImageTk.PhotoImage(image)
+    label = Label(root, image = photo)
+    label.pack()
+
+    root.mainloop()
+
+    # uruchamia sie funkcja spr czy obrazek to paczka czy list
+    pred_print = prediction(path, transformer)
+    print('Zdjecie jest: ' + pred_print)
+    return pred_print
+
+predict_one('src/test_images/1.jpg')

NewralNetwork/src/data_model.py

@@ -0,0 +1,59 @@
+import torch.nn as nn
+import torch
+
+
+class DataModel(nn.Module):
+    def __init__(self, num_objects):
+        super(DataModel, self).__init__()
+        # input (batch=256, nr of channels rgb=3 , size=244x244)
+
+        # convolution
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=12, kernel_size=3, stride=1, padding=1)
+        # shape (256, 12, 224x224)
+
+        # batch normalization
+        self.bn1 = nn.BatchNorm2d(num_features=12)
+        # shape (256, 12, 224x224)
+        self.reul1 = nn.ReLU()
+
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+        # reduce image size by factor 2
+        # pooling window moves by 2 pixels at a time instead of 1
+        # shape (256, 12, 112x112)
+
+        self.conv2 = nn.Conv2d(in_channels=12, out_channels=24, kernel_size=3, stride=1, padding=1)
+        self.bn2 = nn.BatchNorm2d(num_features=24)
+        self.reul2 = nn.ReLU()
+        # shape (256, 24, 112x112)
+
+        self.conv3 = nn.Conv2d(in_channels=24, out_channels=48, kernel_size=3, stride=1, padding=1)
+        # shape (256, 48, 112x112)
+        self.bn3 = nn.BatchNorm2d(num_features=48)
+        # shape (256, 48, 112x112)
+        self.reul3 = nn.ReLU()
+
+        # connected layer
+        self.fc = nn.Linear(in_features=48 * 112 * 112, out_features=num_objects)
+
+    def forward(self, input):
+        output = self.conv1(input)
+        output = self.bn1(output)
+        output = self.reul1(output)
+
+        output = self.pool(output)
+        output = self.conv2(output)
+        output = self.bn2(output)
+        output = self.reul2(output)
+
+        output = self.conv3(output)
+        output = self.bn3(output)
+        output = self.reul3(output)
+
+        # output shape matrix (256, 48, 112x112)
+        # print(output.shape)
+        # print(self.fc.weight.shape)
+
+        output = output.view(-1, 48 * 112 * 112)
+        output = self.fc(output)
+
+        return output

NewralNetwork/src/torchvision_resize_dataset.py

@@ -0,0 +1,31 @@
+import glob
+import pathlib
+import torchvision.transforms as transforms
+from torchvision.datasets import ImageFolder
+from torch.utils.data import ConcatDataset
+
+# images have to be the same size for the algorithm to work
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),  # Resize images to (224, 224)
+    transforms.ToTensor(),  # Convert images to tensors, 0-255 to 0-1
+    # transforms.RandomHorizontalFlip(),  # 0.5 chance to flip the image
+    transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])
+])
+
+carrot_path = '/train_images/carrot'
+potato_path = '/train_images/potato'
+images_path = '/train_images'
+
+# # Load images from folders
+# letter_folder = ImageFolder(letters_path, transform=transform)
+# package_folder = ImageFolder(package_path, transform=transform)
+
+# Combine the both datasets into a single dataset
+#combined_dataset = ConcatDataset([letter_folder, package_folder])
+combined_dataset = ImageFolder(images_path, transform=transform)
+
+#image classes
+path=pathlib.Path(images_path)
+classes = sorted([i.name.split("/")[-1] for i in path.iterdir()])
+
+# print(classes)