Add neuralnetwork.py

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')