import numpy as np import torch from numpy import asarray from torch import nn from torch import optim from torch.autograd import Variable from torch.utils.data import Dataset import os import pandas as pd from PIL import Image from torchvision.io import read_image from torchvision.transforms import Resize, Lambda, transforms, ToTensor from mnist import MNIST from Constants import * from torch.utils.data import DataLoader class Neurons: def __init__(self): dataset = MNIST('C:\\Users\\kratu\\PycharmProjects\\Projekt_AI-Automatyczny_saper\\Engine', gz=True, return_type='numpy') images, labels = dataset.load_training() test_images, test_labels = dataset.load_testing() images = images / 255 test_images = test_images / 255 self.train_images = [torch.tensor(image, dtype=torch.float32) for image in images] self.train_labels = [torch.tensor(label, dtype=torch.long) for label in labels] test_images = [torch.tensor(image, dtype=torch.float32) for image in test_images] test_labels = [torch.tensor(label, dtype=torch.long) for label in test_labels] input_dim = 28 * 28 output_dim = 10 self.model = nn.Sequential( nn.Linear(input_dim, output_dim), nn.LogSoftmax() ) self.model.eval() torch.save(self.train(self.model, 100), 'model.pth') imsize = 28 loader = transforms.Compose([transforms.Resize((imsize, imsize)), transforms.ToTensor()]) image = Image.open("C:\\Users\\kratu\\PycharmProjects\\JPG-PNG-to-MNIST-NN-Format\\test-images\\1\\claymore (1).png") image = image.convert("1") image = loader(image).float() image = Variable(image, requires_grad=True) predicted = self.model(image) print(predicted) def train(self,model, n_iter): criterion = nn.NLLLoss() optimizer = optim.SGD(model.parameters(), lr=0.001) for epoch in range(n_iter): for image, label in zip(self.train_images, self.train_labels): optimizer.zero_grad() output = model(image) loss = criterion(output.unsqueeze(0), label.unsqueeze(0)) loss.backward() optimizer.step() print(f'epoch: {epoch:03}') return model def accuracy(self,expected, predicted): return len([int for e, p in zip(expected, predicted) if e == p])/len(expected) def loadImages(self): col_dir = 'D:\\fotyProjekt\\clowns' images = [] labels = [] imgs = os.listdir(col_dir) for i in imgs: images.append(i) print("Load clowns") print(imgs) for i in imgs: labels.append(self.getLabelsForImages(DECOY)) col_dir = 'D:\\fotyProjekt\\chemical' imgs = os.listdir(col_dir) for i in imgs: images.append(i) print("Load chemical") for i in imgs: labels.append(self.getLabelsForImages(CHEMICAL_BOMB)) col_dir = 'D:\\fotyProjekt\\atomic' imgs = os.listdir(col_dir) for i in imgs: images.append(i) print("Load atomic") for i in imgs: labels.append(self.getLabelsForImages(ATOMIC_BOMB)) col_dir = 'D:\\fotyProjekt\\landmine' imgs = os.listdir(col_dir) for i in imgs: images.append(i) print("Load landmine") for i in imgs: labels.append(self.getLabelsForImages(LAND_MINE)) col_dir = 'D:\\fotyProjekt\\claymore' imgs = os.listdir(col_dir) for i in imgs: images.append(i) print("Load claymore") for i in imgs: labels.append(self.getLabelsForImages(CLAYMORE)) f = open('labels.csv', 'w') while images: img = images.pop() lb = labels.pop() f.write(os.path.basename(img) + ", " + lb.__str__() + "\n") f.close() def getLabelsForImages(self,imageType): return self.d.get(imageType) if __name__ == '__main__': Neurons()