Merge pull request 'garbage_recognition' (#27) from garbage_recognition into master

Reviewed-on: #27

NeuralNetwork/NeuralNetwork.py

@@ -0,0 +1,29 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+class NeuralNetwork(nn.Module):
+    def __init__(self):
+        super(NeuralNetwork, self).__init__()
+        
+        #  Warstwy konwolucyjnej sieci neuronowej, filtr 5×5, 3 kanały dla RGB
+        self.convolutional_nn_1 = nn.Conv2d(3, 6, 5)
+        self.convolutional_nn_2 = nn.Conv2d(6, 16, 5)
+
+        # Wyciaganie "najwazniejszej" informacji z obrazu
+        self.pool = nn.MaxPool2d(2, 2)
+
+        self.full_connection_layer_1 = nn.Linear(16 * 71 * 71, 120)
+        self.full_connection_layer_2 = nn.Linear(120, 84)
+        self.full_connection_layer_3 = nn.Linear(84, 4)
+
+    # Forward określa przepływ inputu przez warstwy
+    def forward(self, x):
+        x = self.pool(F.relu(self.convolutional_nn_1(x)))
+        x = self.pool(F.relu(self.convolutional_nn_2(x)))
+
+        # 16 kanałów o rozmiarach 71x71
+        x = x.view(x.size(0), 16 * 71 * 71)
+        x = F.relu(self.full_connection_layer_1(x))
+        x = F.relu(self.full_connection_layer_2(x))
+        x = self.full_connection_layer_3(x)
+        return x

NeuralNetwork/prediction.py

@@ -0,0 +1,20 @@
+import torch
+import torchvision.transforms as transforms
+
+from PIL import Image
+from NeuralNetwork import NeuralNetwork
+
+def getPrediction(img_path):
+
+    # Inicjacja sieci neuronowej
+    neural_net = NeuralNetwork()
+    PATH = './trained_nn.pth'
+    img = Image.open(img_path)
+    transform_tensor = transforms.ToTensor()(img).unsqueeze_(0)
+    classes = ['glass', 'metal', 'paper', 'plastic']
+    neural_net.load_state_dict(torch.load(PATH))
+    neural_net.eval()
+    outputs = neural_net(transform_tensor)
+
+    # Wyciągnięcie największej wagi co przekłada się na rozpoznanie klasy, w tym przypadku rodzju odpadu
+    return classes[torch.max(outputs, 1)[1]]

NeuralNetwork/train_nn.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn as nn
+import torchvision.transforms as transforms
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torchvision.datasets import ImageFolder
+from NeuralNetwork import NeuralNetwork
+
+# import matplotlib.pyplot as plt
+# import numpy as np
+# import cv2
+
+def trainNeuralNetwork():
+    neural_net = NeuralNetwork()
+    train_set = ImageFolder(root='./resources/trash_dataset/train', transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]))
+    trainloader = DataLoader(
+        train_set, batch_size=2, shuffle=True, num_workers=2)
+
+    # potrzebne do wyświetlania loss w każdej iteracji
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.SGD(neural_net.parameters(), lr=0.001, momentum=0.9)
+
+    epoch_num = 4 # najlepiej 10, dla lepszej wiarygodności
+    for epoch in range(epoch_num):
+        measure_loss = 0.0
+        for i, data in enumerate(trainloader, 0):
+            inputs, labels = data
+            # czyszczenie gradientu f-cji
+            optimizer.zero_grad()
+            outputs = neural_net(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+            measure_loss += loss.item()
+            if i:
+                print('[%d, %5d] loss: %.3f' %
+                      (epoch + 1, i + 1, measure_loss))
+                measure_loss = 0.0
+
+    print('Finished.')
+    PATH = './trained_nn.pth'
+    torch.save(neural_net.state_dict(), PATH)
+
+def main():
+    trainNeuralNetwork()
+
+if __name__ == '__main__':
+    main()
+

decision_tree/tree_as_txt.txt

@@ -16,28 +16,6 @@
 |   |   |   |   |   |--- feature_1 <= 2.50
 |   |   |   |   |   |   |--- feature_0 <= 2.50
 |   |   |   |   |   |   |   |--- feature_1 <= 1.50
-|   |   |   |   |   |   |   |   |--- feature_4 <= 2.50
-|   |   |   |   |   |   |   |   |   |--- class: 1
-|   |   |   |   |   |   |   |   |--- feature_4 >  2.50
-|   |   |   |   |   |   |   |   |   |--- feature_2 <= 2.00
-|   |   |   |   |   |   |   |   |   |   |--- class: 1
-|   |   |   |   |   |   |   |   |   |--- feature_2 >  2.00
-|   |   |   |   |   |   |   |   |   |   |--- class: 0
-|   |   |   |   |   |   |   |--- feature_1 >  1.50
-|   |   |   |   |   |   |   |   |--- class: 0
-|   |   |   |   |   |   |--- feature_0 >  2.50
-|   |   |   |   |   |   |   |--- feature_2 <= 2.50
-|   |   |   |   |   |   |   |   |--- class: 1
-|   |   |   |   |   |   |   |--- feature_2 >  2.50
-|   |   |   |   |   |   |   |   |--- feature_4 <= 2.50
-|   |   |   |   |   |   |   |   |   |--- class: 1
-|   |   |   |   |   |   |   |   |--- feature_4 >  2.50
-|   |   |   |   |   |   |   |   |   |--- class: 0
-|   |   |   |   |   |--- feature_1 >  2.50
-|   |   |   |   |   |   |--- feature_0 <= 3.50
-|   |   |   |   |   |   |   |--- class: 0
-|   |   |   |   |   |   |--- feature_0 >  3.50
-|   |   |   |   |   |   |   |--- feature_1 <= 3.50
 |   |   |   |   |   |   |   |   |--- feature_2 <= 2.50
 |   |   |   |   |   |   |   |   |   |--- class: 1
 |   |   |   |   |   |   |   |   |--- feature_2 >  2.50
@@ -45,6 +23,28 @@
 |   |   |   |   |   |   |   |   |   |   |--- class: 1
 |   |   |   |   |   |   |   |   |   |--- feature_4 >  2.00
 |   |   |   |   |   |   |   |   |   |   |--- class: 0
+|   |   |   |   |   |   |   |--- feature_1 >  1.50
+|   |   |   |   |   |   |   |   |--- class: 0
+|   |   |   |   |   |   |--- feature_0 >  2.50
+|   |   |   |   |   |   |   |--- feature_4 <= 2.50
+|   |   |   |   |   |   |   |   |--- class: 1
+|   |   |   |   |   |   |   |--- feature_4 >  2.50
+|   |   |   |   |   |   |   |   |--- feature_2 <= 2.50
+|   |   |   |   |   |   |   |   |   |--- class: 1
+|   |   |   |   |   |   |   |   |--- feature_2 >  2.50
+|   |   |   |   |   |   |   |   |   |--- class: 0
+|   |   |   |   |   |--- feature_1 >  2.50
+|   |   |   |   |   |   |--- feature_0 <= 3.50
+|   |   |   |   |   |   |   |--- class: 0
+|   |   |   |   |   |   |--- feature_0 >  3.50
+|   |   |   |   |   |   |   |--- feature_1 <= 3.50
+|   |   |   |   |   |   |   |   |--- feature_4 <= 2.50
+|   |   |   |   |   |   |   |   |   |--- class: 1
+|   |   |   |   |   |   |   |   |--- feature_4 >  2.50
+|   |   |   |   |   |   |   |   |   |--- feature_2 <= 2.00
+|   |   |   |   |   |   |   |   |   |   |--- class: 1
+|   |   |   |   |   |   |   |   |   |--- feature_2 >  2.00
+|   |   |   |   |   |   |   |   |   |   |--- class: 0
 |   |   |   |   |   |   |   |--- feature_1 >  3.50
 |   |   |   |   |   |   |   |   |--- class: 0
 |   |   |   |   |--- feature_3 >  4.50
@@ -59,13 +59,13 @@
 |   |   |   |   |   |   |--- class: 1
 |   |   |   |   |--- feature_3 >  3.50
 |   |   |   |   |   |--- feature_1 <= 2.50
-|   |   |   |   |   |   |--- feature_3 <= 4.50
-|   |   |   |   |   |   |   |--- feature_0 <= 2.50
-|   |   |   |   |   |   |   |   |--- class: 0
-|   |   |   |   |   |   |   |--- feature_0 >  2.50
-|   |   |   |   |   |   |   |   |--- class: 1
-|   |   |   |   |   |   |--- feature_3 >  4.50
+|   |   |   |   |   |   |--- feature_0 <= 2.50
 |   |   |   |   |   |   |   |--- class: 0
+|   |   |   |   |   |   |--- feature_0 >  2.50
+|   |   |   |   |   |   |   |--- feature_3 <= 4.50
+|   |   |   |   |   |   |   |   |--- class: 1
+|   |   |   |   |   |   |   |--- feature_3 >  4.50
+|   |   |   |   |   |   |   |   |--- class: 0
 |   |   |   |   |   |--- feature_1 >  2.50
 |   |   |   |   |   |   |--- class: 0
 |   |   |   |--- feature_4 >  4.50

main.py

@@ -118,7 +118,7 @@ class Game():
         #     print('----')
 
         print('positive actions')
-        print(len(self.positive_actions))
+        print(len(self.positive_decision))
         for i in self.positive_decision:
             # print(i.get_coords())
             trash_x, trash_y = i.get_coords()