134 lines
6.5 KiB
Python
134 lines
6.5 KiB
Python
from PIL import Image
|
|
from torch.autograd import Variable
|
|
from torch.optim import Adam
|
|
from torch.utils.data import DataLoader
|
|
from torchvision.transforms import transforms
|
|
import glob
|
|
import os
|
|
import pathlib
|
|
import torch
|
|
import torch.nn as nn
|
|
import torchvision
|
|
transformer1 = transforms.Compose([transforms.Resize((150, 150)), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
|
|
class ConvNet(nn.Module):
|
|
def __init__(self, num_classes=6):
|
|
super(ConvNet, self).__init__()
|
|
self.conv1 = nn.Conv2d(in_channels=3, out_channels=12, kernel_size=3, stride=1, padding=1)
|
|
self.bn1 = nn.BatchNorm2d(num_features=12)
|
|
self.relu1 = nn.ReLU()
|
|
self.pool = nn.MaxPool2d(kernel_size=2)
|
|
self.conv2 = nn.Conv2d(in_channels=12, out_channels=20, kernel_size=3, stride=1, padding=1)
|
|
self.relu2 = nn.ReLU()
|
|
self.conv3 = nn.Conv2d(in_channels=20, out_channels=32, kernel_size=3, stride=1, padding=1)
|
|
self.bn3 = nn.BatchNorm2d(num_features=32)
|
|
self.relu3 = nn.ReLU()
|
|
self.fc = nn.Linear(in_features=75 * 75 * 32, out_features=num_classes)
|
|
def forward(self, input):
|
|
output = self.conv1(input)
|
|
output = self.bn1(output)
|
|
output = self.relu1(output)
|
|
output = self.pool(output)
|
|
output = self.conv2(output)
|
|
output = self.relu2(output)
|
|
output = self.conv3(output)
|
|
output = self.bn3(output)
|
|
output = self.relu3(output)
|
|
output = output.view(-1, 32 * 75 * 75)
|
|
output = self.fc(output)
|
|
return output
|
|
def create_neural_network(): #tworzenie sieci neuronowej
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') #użyj cuda jeśli możliwe
|
|
transformer = transforms.Compose([transforms.Resize((150, 150)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
|
|
train_path = os.path.join('resources/neural_network/train/') #ścieżka do obrazków do treningu
|
|
test_path = os.path.join('resources/neural_network/test/') #ścieżka do obrazków do testu
|
|
train_loader = DataLoader(torchvision.datasets.ImageFolder(train_path, transform=transformer), batch_size=64, shuffle=True)
|
|
test_loader = DataLoader(torchvision.datasets.ImageFolder(test_path, transform=transformer), batch_size=32, shuffle=True)
|
|
root = pathlib.Path(train_path)
|
|
classes = sorted([j.name.split('/')[-1] for j in root.iterdir()])
|
|
if os.path.exists("resources/neural_network/checkpoint.model"): #jeżeli istnieje model to wczytaj
|
|
checkpoint = torch.load(os.path.join('resources/neural_network', 'checkpoint.model'))
|
|
model = ConvNet(num_classes=6)
|
|
model.load_state_dict(checkpoint)
|
|
model.eval()
|
|
else: #w przeciwnym razie utwórz nowy model
|
|
model = ConvNet(num_classes=6).to(device)
|
|
optimizer = Adam(model.parameters(), lr=0.001, weight_decay=0.0001)
|
|
loss_function = nn.CrossEntropyLoss()
|
|
num_epochs = 10
|
|
train_count = len(glob.glob(train_path + '/**/*.png')) #liczba obrazków treningowych
|
|
test_count = len(glob.glob(test_path + '/**/*.png')) #liczba obrazków testowych
|
|
best_accuracy = 0.0
|
|
for epoch in range(num_epochs):
|
|
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
|
|
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 + 1) + ' Train Loss: ' + str(train_loss) + ' Train Accuracy: ' + str(train_accuracy) + ' Test Accuracy: ' + str(test_accuracy))
|
|
if test_accuracy > best_accuracy:
|
|
torch.save(model.state_dict(), 'resources/neural_network/checkpoint.model')
|
|
best_accuracy = test_accuracy
|
|
checkpoint = torch.load(os.path.join('resources/neural_network', 'checkpoint.model'))
|
|
model = ConvNet(num_classes=6)
|
|
model.load_state_dict(checkpoint)
|
|
model.eval()
|
|
return classes, model
|
|
def predfield(classes, model): #zwraca miejsce pola z wyrośniętą rośliną na podstawie wykrywania obrazu
|
|
pred_path = os.path.join('resources/neural_network/sliced/') #ścieżka do obrazków do sprawdzenia
|
|
pred_dict = {}
|
|
images_path = glob.glob(pred_path + '/*.png')
|
|
x = None #x'owa pola
|
|
y = None#y'kowa pola
|
|
for i in images_path: #dodajemy pocięte obrazki do listy i ustawiamy im przewidywaną metkę
|
|
pred_dict[i[i.rfind('/') + 1:]] = prediction1(classes, i, model, transformer1)
|
|
for img_name, field in pred_dict.items():
|
|
if field != "random": #jeżeli metka nie jest 'random' to przypisz do x'a i y'a miejsce wyrośniętej rośliny
|
|
x = img_name[15]
|
|
y = img_name[18]
|
|
x = int(x)
|
|
y = int(y)
|
|
if x == 0:
|
|
x = 9
|
|
else:
|
|
x = x - 1
|
|
if y == 0:
|
|
y = 9
|
|
else:
|
|
y = y - 1
|
|
if x == None and y == None: #jeżeli nie ma wyrośniętej rośliny to zwróć False
|
|
return False
|
|
else:
|
|
return y, x
|
|
def prediction1(classes, img_path, model, transformer): #zwraca predykcję dla danego obrazka
|
|
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 |