Add Model

.gitignore

@@ -1,3 +1,7 @@
+# Dataset
+Dataset/
+*.jpeg
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

Main.py

@@ -0,0 +1,26 @@
+from PIL import Image
+import torch
+import torchvision.transforms as transforms
+from Model import CNN
+
+transform = transforms.Compose([
+    transforms.ToTensor()
+])
+
+classes = ("Circle", "Square", "Triangle")
+
+# Load image
+image = Image.open("./Input.jpeg")
+image = transform(image)
+
+# Load model
+model = CNN()
+
+model.load_state_dict(torch.load("Model.pth"))
+model.eval()
+
+with torch.no_grad():
+    outputs = model(image)
+    _, predicted = torch.max(outputs.data, 1)
+
+print(classes[predicted[0]])

Model.py

@@ -0,0 +1,28 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+class CNN(nn.Module):
+    def __init__(self):
+        super(CNN, self).__init__()
+
+        self.layer1 = nn.Sequential(
+            nn.Conv2d(3, 32, 3, 1, 1),
+            nn.ReLU(True),
+            nn.MaxPool2d(2, 2)
+        )
+        self.layer2 = nn.Sequential(
+            nn.Conv2d(32, 64, 3, 1, 1),
+            nn.ReLU(True),
+            nn.MaxPool2d(2, 2)
+        )
+        self.func1 = nn.Linear(64 * 7 * 7, 128)
+        self.func2 = nn.Linear(128, 3)
+
+    def forward(self, x):
+        out = self.layer1(x)
+        out = self.layer2(out)
+        out = out.view(-1, 64 * 7 * 7)
+        out = F.relu(self.func1(out))
+        out = self.func2(out)
+
+        return out

Test.py

@@ -0,0 +1,41 @@
+import torch
+import torchvision
+import torchvision.transforms as transforms
+from Model import CNN
+
+batchSize = 4
+transform = transforms.Compose([
+    transforms.ToTensor()
+])
+
+# Load test dataset
+testSet = torchvision.datasets.ImageFolder(
+    root = "./Dataset/Test/",
+    transform = transform
+)
+testLoader = torch.utils.data.DataLoader(
+    testSet,
+    batch_size = batchSize,
+    shuffle = True,
+    num_workers = 0
+)
+
+# Load model
+model = CNN()
+
+model.load_state_dict(torch.load("Model.pth"))
+model.eval()
+
+correct = 0
+total = 0
+
+with torch.no_grad():
+    for inputs, labels in testLoader:
+        outputs = model(inputs)
+        _, predicted = torch.max(outputs.data, 1)
+        total += labels.size(0)
+        correct += (predicted == labels).sum().item()
+
+accuracy = correct / total * 100
+
+print(f"Accuracy: {accuracy:.2f} %")

Train.py

@@ -0,0 +1,48 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torchvision
+import torchvision.transforms as transforms
+from Model import CNN
+
+batchSize = 4
+transform = transforms.Compose([
+    transforms.ToTensor()
+])
+
+# Load train dataset
+trainSet = torchvision.datasets.ImageFolder(
+    root = "./Dataset/Train/",
+    transform = transform
+)
+trainLoader = torch.utils.data.DataLoader(
+    trainSet,
+    batch_size = batchSize,
+    shuffle = True,
+    num_workers = 0
+)
+
+# Model training
+epochs = 20
+learningRate = 1e-4
+model = CNN()
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr = learningRate)
+
+model.train()
+
+for epoch in range(epochs):
+    for images, labels in trainLoader:
+        optimizer.zero_grad()
+
+        outputs = model(images)
+        loss = criterion(outputs, labels)
+
+        loss.backward()
+        optimizer.step()
+
+    print(f"[Epoch: {epoch + 1:5d}/{epochs}] Loss: {loss.item()}")
+
+torch.save(model.state_dict(), "Model.pth")
+
+print("Finished Training!")