clean code

src/machine_learning/neural_network/learning.py

@@ -1,5 +1,4 @@
 import os
-import numpy as np
 import torch
 import glob
 import torch.nn as nn
@@ -22,49 +21,42 @@ test_dir = r'images\learning\test'
 train_dir = os.path.join(temp_path, train_dir)
 test_dir = os.path.join(temp_path, test_dir)
 
-device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-#Transforms
+transformer = transforms.Compose([
-transformer=transforms.Compose([
     transforms.Resize((150,150)),
     transforms.RandomHorizontalFlip(),
-    transforms.ToTensor(),  #0-255 to 0-1, numpy to tensors
+    transforms.ToTensor(),
-    transforms.Normalize([0.5,0.5,0.5], # 0-1 to [-1,1] , formula (x-mean)/std
+    transforms.Normalize([0.5,0.5,0.5],
                         [0.5,0.5,0.5])
 ])
 
-#Dataloader
+train_path = r'C:\Users\User\PycharmProjects\Super-Saper222\images\learning\training\training'
+test_path = r'C:\Users\User\PycharmProjects\Super-Saper222\images\learning\test\test'
+pred_path = r'C:\Users\User\PycharmProjects\Super-Saper222\images\learning\prediction\prediction'
 
-#Path for training and testing directory
-train_path = r'D:\Documents\Studia\Semestr-4\sztuczna-inteligencja\super-saper\images\learning\training\training'
-test_path = r'D:\Documents\Studia\Semestr-4\sztuczna-inteligencja\super-saper\images\learning\test\test'
-pred_path = r'D:\Documents\Studia\Semestr-4\sztuczna-inteligencja\super-saper\images\learning\prediction\prediction'
 
-print(train_path)
+train_loader = DataLoader(
-
-train_loader=DataLoader(
     torchvision.datasets.ImageFolder(train_path, transform=transformer),
     batch_size=64, shuffle=True
 )
-test_loader=DataLoader(
+test_loader = DataLoader(
     torchvision.datasets.ImageFolder(test_path, transform=transformer),
     batch_size=32, shuffle=True
 )
 
-#categories
 root=pathlib.Path(train_path)
-classes=sorted([j.name.split('/')[-1] for j in root.iterdir()])
+classes = sorted([j.name.split('/')[-1] for j in root.iterdir()])
 
-model=Net(num_classes=6).to(device)
+model = Net(num_classes=6).to(device)
 
-#Optmizer and loss function
+optimizer = Adam(model.parameters(),lr=1e-3,weight_decay=0.0001)
-optimizer=Adam(model.parameters(),lr=1e-3,weight_decay=0.0001)
+loss_fn = nn.CrossEntropyLoss()
-loss_fn=nn.CrossEntropyLoss()
 
-num_epochs=10
+num_epochs = 10
 
-train_count=len(glob.glob(train_path+'/**/*.*'))
+train_count = len(glob.glob(train_path+'/**/*.*'))
-test_count=len(glob.glob(test_path+'/**/*.*'))
+test_count = len(glob.glob(test_path+'/**/*.*'))
 
 print(train_count,test_count)
 
@@ -75,10 +67,9 @@ def train(dataloader, model, loss_fn, optimizer):
     for batch, (X, y) in enumerate(dataloader):
         X, y = X.to(device), y.to(device)
 
-        # Compute prediction error
         pred = model(X.float())
         loss = loss_fn(pred, y)
-        # Backpropagation
+        
         optimizer.zero_grad()
         loss.backward()
         optimizer.step()
@@ -92,12 +83,14 @@ def test(dataloader, model):
     size = len(dataloader.dataset)
     model.eval()
     test_loss, correct = 0, 0
+    
     with torch.no_grad():
         for X, y in dataloader:
             X, y = X.to(device), y.to(device)
             pred = model(X.float())
             test_loss += loss_fn(pred, y).item()
             correct += (pred.argmax(1) == y).type(torch.float).sum().item()
+            
     test_loss /= size
     correct /= size
     print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
@@ -107,29 +100,40 @@ def prediction1(classes, img_path, model, transformer):
     image = Image.open(img_path).convert('RGB')
     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
 
-transformer1 = transforms.Compose([transforms.Resize((150, 150)), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+transformer1 = transforms.Compose([transforms.Resize((150, 150)),
+                                   transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+
+#creating new model
 
 # for t in range(9):
 #     print(f"Epoch {t+1}\n-------------------------------")
 #     train(train_loader, model, loss_fn, optimizer)
 #     test(test_loader, model)
 # print("Done!")
+# torch.save(model.state_dict(), 'mine_recognizer.model')
 
-checkpoint = torch.load(os.path.join('.', 'best_checkpoint.model'))
+
+#checking work of new model
+
+checkpoint = torch.load(os.path.join('.', 'mine_recognizer.model'))
 model = Net(num_classes=6)
 model.load_state_dict(checkpoint)
 model.eval()
+
 transformer1 = transforms.Compose([transforms.Resize((150, 150)), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
 images_path = glob.glob(pred_path+'/*.*')
 pred_dict = {}
+
 for i in images_path:
     pred_dict[i[i.rfind('/') + 1:]] = prediction1(classes, i, model, transformer1)
 print(pred_dict)
@@ -144,52 +148,3 @@ for i in images_path:
     pred_dict[i[i.rfind('/') + 1:]] = prediction1(classes, i, model, transformer1)
 print(pred_dict)
 
-# for epoch in range(num_epochs):
-#
-#     # Evaluation and training on training 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 = Variable(images.cuda())
-#             labels = Variable(labels.cuda())
-#
-#         optimizer.zero_grad()
-#
-#         outputs = model(images)
-#         loss = loss_function(outputs, labels)
-#         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_count
-#     train_loss = train_loss / train_count
-#
-#     # Evaluation on testing dataset
-#     model.eval()
-#
-#     test_accuracy = 0.0
-#     for i, (images, labels) in enumerate(test_loader):
-#         if torch.cuda.is_available():
-#             images = Variable(images.cuda())
-#             labels = Variable(labels.cuda())
-#
-#         outputs = model(images)
-#         _, prediction = torch.max(outputs.data, 1)
-#         test_accuracy += int(torch.sum(prediction == labels.data))
-#
-#     test_accuracy = test_accuracy / test_count
-#
-#     print('Epoch: ' + str(epoch) + ' Train Loss: ' + str(train_loss) + ' Train Accuracy: ' + str(
-#         train_accuracy) + ' Test Accuracy: ' + str(test_accuracy))
-#
-#     # Save the best model
-#     if test_accuracy > best_accuracy:
-#         torch.save(model.state_dict(), 'best_checkpoint.model')
-#         best_accuracy = test_accuracy