clean code

This commit is contained in:
piowyg 2021-05-31 20:10:40 +02:00
parent e77c581dda
commit 1999432920
3 changed files with 32 additions and 77 deletions

View File

@ -1,5 +1,4 @@
import os import os
import numpy as np
import torch import torch
import glob import glob
import torch.nn as nn import torch.nn as nn
@ -22,49 +21,42 @@ test_dir = r'images\learning\test'
train_dir = os.path.join(temp_path, train_dir) train_dir = os.path.join(temp_path, train_dir)
test_dir = os.path.join(temp_path, test_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.Resize((150,150)),
transforms.RandomHorizontalFlip(), 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]) [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), torchvision.datasets.ImageFolder(train_path, transform=transformer),
batch_size=64, shuffle=True batch_size=64, shuffle=True
) )
test_loader=DataLoader( test_loader = DataLoader(
torchvision.datasets.ImageFolder(test_path, transform=transformer), torchvision.datasets.ImageFolder(test_path, transform=transformer),
batch_size=32, shuffle=True batch_size=32, shuffle=True
) )
#categories
root=pathlib.Path(train_path) 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) print(train_count,test_count)
@ -75,10 +67,9 @@ def train(dataloader, model, loss_fn, optimizer):
for batch, (X, y) in enumerate(dataloader): for batch, (X, y) in enumerate(dataloader):
X, y = X.to(device), y.to(device) X, y = X.to(device), y.to(device)
# Compute prediction error
pred = model(X.float()) pred = model(X.float())
loss = loss_fn(pred, y) loss = loss_fn(pred, y)
# Backpropagation
optimizer.zero_grad() optimizer.zero_grad()
loss.backward() loss.backward()
optimizer.step() optimizer.step()
@ -92,12 +83,14 @@ def test(dataloader, model):
size = len(dataloader.dataset) size = len(dataloader.dataset)
model.eval() model.eval()
test_loss, correct = 0, 0 test_loss, correct = 0, 0
with torch.no_grad(): with torch.no_grad():
for X, y in dataloader: for X, y in dataloader:
X, y = X.to(device), y.to(device) X, y = X.to(device), y.to(device)
pred = model(X.float()) pred = model(X.float())
test_loss += loss_fn(pred, y).item() test_loss += loss_fn(pred, y).item()
correct += (pred.argmax(1) == y).type(torch.float).sum().item() correct += (pred.argmax(1) == y).type(torch.float).sum().item()
test_loss /= size test_loss /= size
correct /= size correct /= size
print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n") 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 = Image.open(img_path).convert('RGB')
image_tensor = transformer(image).float() image_tensor = transformer(image).float()
image_tensor = image_tensor.unsqueeze_(0) image_tensor = image_tensor.unsqueeze_(0)
if torch.cuda.is_available(): if torch.cuda.is_available():
image_tensor.cuda() image_tensor.cuda()
input = Variable(image_tensor) input = Variable(image_tensor)
output = model(input) output = model(input)
index = output.data.numpy().argmax() index = output.data.numpy().argmax()
pred = classes[index] pred = classes[index]
return pred 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): # for t in range(9):
# print(f"Epoch {t+1}\n-------------------------------") # print(f"Epoch {t+1}\n-------------------------------")
# train(train_loader, model, loss_fn, optimizer) # train(train_loader, model, loss_fn, optimizer)
# test(test_loader, model) # test(test_loader, model)
# print("Done!") # 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 = Net(num_classes=6)
model.load_state_dict(checkpoint) model.load_state_dict(checkpoint)
model.eval() 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])]) 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+'/*.*') images_path = glob.glob(pred_path+'/*.*')
pred_dict = {} pred_dict = {}
for i in images_path: for i in images_path:
pred_dict[i[i.rfind('/') + 1:]] = prediction1(classes, i, model, transformer1) pred_dict[i[i.rfind('/') + 1:]] = prediction1(classes, i, model, transformer1)
print(pred_dict) print(pred_dict)
@ -144,52 +148,3 @@ for i in images_path:
pred_dict[i[i.rfind('/') + 1:]] = prediction1(classes, i, model, transformer1) pred_dict[i[i.rfind('/') + 1:]] = prediction1(classes, i, model, transformer1)
print(pred_dict) 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