125 lines
4.0 KiB
Python
125 lines
4.0 KiB
Python
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\\Konrad\\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\\Konrad\\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()
|