From 28b120769d3f34b16c722a31bc4a2ac93d21abfb Mon Sep 17 00:00:00 2001 From: Woj Date: Sun, 4 Jun 2023 00:09:57 +0200 Subject: [PATCH] neural network model + datasetconverter --- NeuralNetwork/neural_network_learning.py | 17 ++++++ NeuralNetwork/src/data_model.py | 58 +++++++++++++++++++ .../src/torchvision_resize_dataset.py | 22 +++++++ 3 files changed, 97 insertions(+) create mode 100644 NeuralNetwork/neural_network_learning.py create mode 100644 NeuralNetwork/src/data_model.py create mode 100644 NeuralNetwork/src/torchvision_resize_dataset.py diff --git a/NeuralNetwork/neural_network_learning.py b/NeuralNetwork/neural_network_learning.py new file mode 100644 index 0000000..a0c75e6 --- /dev/null +++ b/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"] + diff --git a/NeuralNetwork/src/data_model.py b/NeuralNetwork/src/data_model.py new file mode 100644 index 0000000..74f40b1 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/NeuralNetwork/src/torchvision_resize_dataset.py b/NeuralNetwork/src/torchvision_resize_dataset.py new file mode 100644 index 0000000..7684909 --- /dev/null +++ b/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])