diff --git a/neural_network.py b/neural_network.py new file mode 100644 index 0000000..b94dbeb --- /dev/null +++ b/neural_network.py @@ -0,0 +1,103 @@ +import torch +import torchvision +import torchvision.transforms as transforms +import torch.nn as nn +import torch.nn.functional as f +import torch.optim as optim +import numpy as np +from matplotlib.pyplot import imshow +import os +import PIL +import numpy as np +from matplotlib.pyplot import imshow + +def to_negative(img): + img = PIL.ImageOps.invert(img) + return img + +class Negative(object): + def __init__(self): + pass + + def __call__(self, img): + return to_negative(img) + +def plotdigit(image): + img = np.reshape(image, (-1, 100)) + imshow(img, cmap='Greys') + +transform = transforms.Compose([Negative(), transforms.ToTensor()]) +train_set = torchvision.datasets.ImageFolder(root='train', transform=transform) +classes = ("apple", "potato") + +BATCH_SIZE = 2 +train_loader = torch.utils.data.DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True, num_workers=0) + +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + +class Net(nn.Module): + def __init__(self): + super(Net, self).__init__() + self.flatten = nn.Flatten() + self.linear_relu_stack = nn.Sequential( + nn.Linear(3*100*100, 512), + nn.ReLU(), + nn.Linear(512, 512), + nn.ReLU(), + nn.Linear(512, 2), + nn.ReLU() + ) + self.linear_relu_stack = self.linear_relu_stack.to(device) + + def forward(self, x): + x = self.flatten(x).to(device) + logits = self.linear_relu_stack(x).to(device) + return logits + +def training_network(): + net = Net() + net = net.to(device) + + criterion = nn.CrossEntropyLoss() + optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) + + for epoch in range(4): + running_loss = 0.0 + for i, data in enumerate(train_loader, 0): + inputs, labels = data[0].to(device), data[1].to(device) + optimizer.zero_grad() + outputs = net(inputs.to(device)) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + if i % 2000 == 1999: + print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss)) + running_loss = 0.0 + + print("Finished training") + save_network_to_file(net) + + +def result_from_network(net, loaded_image): + image = PIL.Image.open(loaded_image) + pil_to_tensor = transforms.ToTensor()(image.convert("RGB")).unsqueeze_(0) + outputs = net(pil_to_tensor.to(device)) + + return classes[torch.max(outputs, 1)[1]] + + +def save_network_to_file(network): + torch.save(network.state_dict(), 'network_model.pth') + print("Network saved to file") + + +def load_network_from_structure(network): + network.load_state_dict(torch.load('network_model.pth')) + + +if __name__ == "__main__": + print(torch.cuda.is_available()) + training_network() + \ No newline at end of file