arnaucube
arnaucube/objectImageIdentifierAIPublic Self-hosted repository
objectImageIdentifierAI/nnTrain/train.py
nnTrain/train.pyRaw
 1import matplotlib.pyplot as plt
 2import numpy as np
 3from random import randint
 4import pickle
 5from sknn.mlp import Classifier, Layer, Convolution
 6
 7def datasetToTrainAndTestData(dataset, numtest):
 8    np.random.shuffle(dataset)
 9    print "length total data:" + str(len(dataset))
10
11    traindata = np.copy(dataset)
12    testdata = []
13    for i in range(numtest):
14        #get random integer between 0 and the total amount of images in the dataset
15        n = randint(0, len(traindata))
16        testdata.append(dataset[n])
17
18        #delete the n image (dataset[n]) of the traindata
19        traindata = np.delete(traindata, n, axis=0)
20    testdataNP = np.array(testdata)
21    return traindata, testdataNP
22
23
24#read the dataset made with the 'imagesToDataset' repository
25dataset = np.load('dataset.npy')
26
27traindata, testdata = datasetToTrainAndTestData(dataset, 10)
28print "length traindata: " + str(len(traindata))
29print "length testdata: " + str(len(testdata))
30
31#traindataAttributes contains all the pixels of each image
32traindataAttributes = traindata[:,0]
33traindataAttributes = np.array([[row] for row in traindataAttributes])
34
35#traindataLabels contains each label of each image
36traindataLabels = traindata[:,1]
37traindataLabels = traindataLabels.astype('int')
38
39#testdataAttributes contains the pixels of the test images
40testdataAttributes = testdata[:,0]
41testdataAttributes = np.array([[row] for row in testdataAttributes])
42
43#testdataLabels contains each label of each image
44testdataLabels = testdata[:,1]
45testdataLabels = testdataLabels.astype('int')
46
47#default: units=100, learning_rate=0.001, n_iter=25
48nn = Classifier(
49    layers=[
50        Layer("Sigmoid", units=10),
51        Layer("Softmax")],
52    learning_rate=0.001,
53    n_iter=20,
54    verbose=True)
55
56nn.fit(traindataAttributes, traindataLabels)
57
58print('\nTRAIN SCORE', nn.score(traindataAttributes, traindataLabels))
59print('TEST SCORE', nn.score(testdataAttributes, testdataLabels))
60
61#save the neural network configuration
62pickle.dump(nn, open('nn.pkl', 'wb'))