neural network model + datasetconverter

NeuralNetwork/neural_network_learning.py

@@ -0,0 +1,17 @@
+from src.torchvision_resize_dataset import combined_dataset
+import torch 
+import torch.nn as nn 
+from torch.utils.data import DataLoader 
+from torchvision import datasets 
+from torchvision.transforms import Compose, Lambda, ToTensor 
+import matplotlib.pyplot as plt
+
+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
+)
+classes = ["package", "list"]
+

NeuralNetwork/src/data_model.py

@@ -0,0 +1,58 @@
+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, 244x244)
+
+        # batch normalization
+        self.bn1 = nn.BatchNorm2d(num_features=12)
+        #shape (256, 12, 244x244)
+        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, 122x122)
+
+
+
+        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, 122x122)
+
+        self.conv3 = nn.Conv2d(in_channels=24, out_channels=48, kernel_size=3, stride=1, padding=1)
+        #shape (256, 48, 122x122)
+        self.bn3 = nn.BatchNorm2d(num_features=48)
+        #shape (256, 48, 122x122)
+        self.reul3 = nn.ReLU()
+
+        # connected layer
+        self.fc = nn.Linear(in_features=48*122*122, 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, 122x122)
+            output = output.view(-1, 48*122*122)
+            output = self.fc(output)
+            
+            return output

NeuralNetwork/src/torchvision_resize_dataset.py

@@ -0,0 +1,22 @@
+import torchvision.transforms as transforms
+from torchvision.datasets import ImageFolder
+from torch.utils.data import ConcatDataset
+import os
+
+# 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]) 
+])
+
+letters_path = './letters'
+package_path = './package'
+
+# # 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])