diff --git a/neuralnetwork.py b/neuralnetwork.py new file mode 100644 index 000000000..5e8b5d4aa --- /dev/null +++ b/neuralnetwork.py @@ -0,0 +1,87 @@ +import torch +import torch.nn as nn +from torch.utils.data import DataLoader +from torchvision import datasets, transforms +from PIL import Image + +class SimpleNN(nn.Module): + def __init__(self): + super(SimpleNN, self).__init__() + self.fc1 = nn.Linear(64 * 64, 128) + self.fc2 = nn.Linear(128, 64) + self.fc3 = nn.Linear(64, 4) + self.relu = nn.ReLU() + self.log_softmax = nn.LogSoftmax(dim=1) + + def forward(self, x): + x = x.view(x.size(0), -1) # Spłaszczenie obrazów + x = self.relu(self.fc1(x)) + x = self.relu(self.fc2(x)) + x = self.log_softmax(self.fc3(x)) + return x + +def train(model, train_loader, n_iter=100): + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + criterion = nn.NLLLoss() + model.train() + for epoch in range(n_iter): + running_loss = 0.0 + for images, targets in train_loader: + images, targets = images.to(device), targets.to(device) + optimizer.zero_grad() + out = model(images) + loss = criterion(out, targets) + loss.backward() + optimizer.step() + running_loss += loss.item() + if epoch % 10 == 0: + print(f'Epoch: {epoch:3d}, Loss: {running_loss/len(train_loader):.4f}') + +def predict_image(image_path, model): + image_path = "warzywa/" + str(image_path) + ".png" + image = Image.open(image_path) + image = transform(image).unsqueeze(0).to(device) + class_names = ["marchewka","ziemniak","rzodkiewka","pomidor"] + with torch.no_grad(): + output = model(image) + _, predicted = torch.max(output, 1) + return class_names[predicted.item()] + +def accuracy(model, dataset): + model.eval() + correct = 0 + total = 0 + with torch.no_grad(): + for images, targets in DataLoader(dataset, batch_size=256): + images, targets = images.to(device), targets.to(device) + outputs = model(images) + _, predicted = torch.max(outputs, 1) + correct += (predicted == targets).sum().item() + total += targets.size(0) + return correct / total + +def load_model(model_path): + model = SimpleNN() + model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'))) + model.eval() + return model + +device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') +transform = transforms.Compose([ + transforms.Resize((64, 64)), + transforms.Grayscale(num_output_channels=1), + transforms.ToTensor(), +]) +train_data_path = 'train' +train_dataset = datasets.ImageFolder(root=train_data_path, transform=transform) +train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) + +test_data_path = 'test' +test_dataset = datasets.ImageFolder(root=train_data_path, transform=transform) + +model = SimpleNN().to(device) + + +if __name__ == "__main__": + train(model, train_loader, n_iter=100) + torch.save(model.state_dict(), 'model.pth')