powerofaugmentation/LettersModel.py at master · liorshk/powerofaugmentation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
import numpy as np
import cv2

from PIL import ImageFont
from PIL import Image
from PIL import ImageDraw
from PIL import ImageFilter

from random import randint
import os
import os.path
from matplotlib import pyplot as plt
from os import listdir
from keras.models import Sequential
from keras.layers.core import Dense, Dropout, Flatten
from keras.layers.convolutional import Convolution2D, MaxPooling2D
from keras.utils import np_utils
from sklearn.cross_validation import train_test_split
from keras.preprocessing.image import ImageDataGenerator


class LettersModel:

    model = None

    def __init__(self,letters,weights_path="letters.h5",letters_folder="letters"):
        np.random.seed(123)
        self.img_rows = 32
        self.img_cols = 32
        self.batch_size = 500
        self.nb_epoch = 20
        self.letters = letters
        self.weights_path=weights_path
        self.letters_folder = letters_folder
        self.nb_classes = len(self.letters)

    def load_model(self):
        self.model = Sequential()
        self.model.add(Convolution2D(32, 4, 4, border_mode='same', activation='relu',input_shape=(1,self.img_rows,self.img_cols)))
        self.model.add(Convolution2D(32, 4, 4, border_mode='same', activation='relu'))
        self.model.add(MaxPooling2D(pool_size=(3, 3)))
        self.model.add(Dropout(0.25))
        self.model.add(Convolution2D(64, 4, 4, border_mode='same', activation='relu'))
        self.model.add(Convolution2D(64, 4, 4, border_mode='same', activation='relu'))
        self.model.add(MaxPooling2D(pool_size=(2, 2)))
        self.model.add(Dropout(0.25))
        self.model.add(Flatten())
        self.model.add(Dense(512, input_dim=64*5*5, activation='relu'))
        self.model.add(Dropout(0.5))
        self.model.add(Dense(self.nb_classes, input_dim=512, activation='softmax'))

        if os.path.isfile(self.weights_path):
            self.model.load_weights(self.weights_path)

        self.model.compile(loss='categorical_crossentropy', optimizer='adam',metrics=['accuracy'])


    def train(self,save_model_to_file = True,rotation_range = 20,width_shift_range=0.5,height_shift_range=0.2):
        """ Trains the model using the dataset in letters_folder """

        # Read the data
        data = []
        labels = []
        for imgName in listdir(self.letters_folder):
            img = cv2.imread(self.letters_folder+"/"+imgName, cv2.IMREAD_GRAYSCALE)
            data.append(img)
            # Get the label from the image path and then get the index from the letters list
            labels.append(self.letters.index(imgName.split('_')[0]))

        data = np.array(data)
        labels = np.array(labels)

        # Split train and test
        X_train, X_test, y_train, y_test = train_test_split(
                 data, labels, test_size=0.33, random_state=42)

        X_train = X_train.reshape(X_train.shape[0], 1, self.img_rows, self.img_cols)
        X_test = X_test.reshape(X_test.shape[0], 1, self.img_rows, self.img_cols)
        X_train = X_train.astype('float32')
        X_test = X_test.astype('float32')
        X_train /= 255
        X_test /= 255

        # convert class vectors to binary class matrices
        Y_train = np_utils.to_categorical(y_train, self.nb_classes)
        Y_test = np_utils.to_categorical(y_test, self.nb_classes)

        # this will do preprocessing and realtime data augmentation
        datagen = ImageDataGenerator(rotation_range=rotation_range,  # randomly rotate images in the range (degrees, 0 to 180)
            width_shift_range=width_shift_range,  # randomly shift images horizontally (fraction of total width)
            height_shift_range=height_shift_range)# randomly shift images vertically (fraction of total height))

        datagen.fit(X_train)

        # fit the model on the batches generated by datagen.flow()
        history = self.model.fit_generator(datagen.flow(X_train, Y_train, batch_size=self.batch_size),
                            samples_per_epoch=X_train.shape[0],
                            nb_epoch=self.nb_epoch,
                            validation_data=(X_test, Y_test))


        # Plot History
        plt.figure(figsize=(10,10))
        plt.plot(history.history['acc'])
        plt.plot(history.history['val_acc'])
        plt.title('model accuracy')
        plt.ylabel('accuracy')
        plt.xlabel('epoch')
        plt.legend(['train', 'test'], loc='upper left')
        plt.show()

        if save_model_to_file:
            self.model.save_weights(self.weights_path,overwrite=True)


    def getModel(self):
        return self.model


    def getLetters(self):
        return self.letters

    def generate_dataset(self,generateNoise = True):
        """ Loads the fonts in the machine and creates a 32x32 image based on the letters in self.letters """

        fonts = []

        invalidFonts = ["Kacst","opens_","Saab","Lohit","lklug","Umpush","Waree"]

        for dirpath, dirnames, filenames in os.walk("/usr/share/fonts/truetype"):
            for filename in [f for f in filenames if f.endswith(".ttf")]:
                fontPath = os.path.join(dirpath, filename)

                isvalidfont = True
                for invalidFont in invalidFonts:
                    if invalidFont in filename:
                        isvalidfont = False

                if isvalidfont:
                    fonts.append(fontPath)

        if not os.path.exists(self.letters_folder):
            os.makedirs(self.letters_folder)

        for fontPath in fonts:
            fontName = fontPath.split('/')[-1][:-4]
            font = ImageFont.truetype(fontPath, 20)

            for l in self.letters:
                if generateNoise:
                    for j in range(1,20):
                        img = Image.new("RGBA", (32,32),"black")
                        draw = ImageDraw.Draw(img)
                        # Draw the letter
                        draw.text((5,5), l.upper(),"white", font=font)

                        for i in range(1,100):
                            # Draw a random line
                            xL = randint(0,32)
                            yL = randint(0,32)
                            draw.point((xL,yL), fill=255)

                        img = img.filter(ImageFilter.GaussianBlur(radius=1))
                        img.save(self.letters_folder+"/"+l+"_"+fontName+str(j)+".png")
                else:
                    img = Image.new("RGBA", (32,32),"black")
                    draw = ImageDraw.Draw(img)
                    # Draw the letter
                    draw.text((5,5), l.upper(),"white", font=font)
                    img.save(self.letters_folder+"/"+l+"_"+fontName+".png")