diff --git a/src/machine_learning/neural_network/best_checkpoint.model b/src/machine_learning/neural_network/best_checkpoint.model deleted file mode 100644 index a4ba027..0000000 Binary files a/src/machine_learning/neural_network/best_checkpoint.model and /dev/null differ diff --git a/src/machine_learning/neural_network/learning.py b/src/machine_learning/neural_network/learning.py index 692c6d4..ad75ff4 100644 --- a/src/machine_learning/neural_network/learning.py +++ b/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.Normalize([0.5,0.5,0.5], # 0-1 to [-1,1] , formula (x-mean)/std + transforms.ToTensor(), + 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) -loss_fn=nn.CrossEntropyLoss() +optimizer = Adam(model.parameters(),lr=1e-3,weight_decay=0.0001) +loss_fn = nn.CrossEntropyLoss() -num_epochs=10 +num_epochs = 10 -train_count=len(glob.glob(train_path+'/**/*.*')) -test_count=len(glob.glob(test_path+'/**/*.*')) +train_count = len(glob.glob(train_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 \ No newline at end of file diff --git a/src/machine_learning/neural_network/mine_recognizer.model b/src/machine_learning/neural_network/mine_recognizer.model new file mode 100644 index 0000000..9b6228f Binary files /dev/null and b/src/machine_learning/neural_network/mine_recognizer.model differ