train.py

# This module trains the CNN based on the labels provided in ./data/CNN
# Note that data must be first split into train, validation, and test data
# by running split_data.py.
# Reference:
# https://towardsdatascience.com/a-single-function-to-streamline-image-classification-with-keras-bd04f5cfe6df

from matplotlib import pyplot as plt
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense
from tensorflow.keras.optimizers import RMSprop
import cv2
import json
import os


NUM_EPOCHS = 10
BATCH_SIZE = 16
DATA_FOLDER = './data/CNN/'


def create_generators(folderpath=DATA_FOLDER):
    '''Creates flow generators to supply images one by one during 
    training/validation phases. Useful when working with large datasets 
    that can't be directly loaded into the memory.'''
    # All images will be rescaled by 1./255
    train_datagen = ImageDataGenerator(rescale=1/255)
    # Flow training images in batches of 128 using train_datagen generator
    train_generator = train_datagen.flow_from_directory(
        folderpath+'train',  # This is the source directory for training images
        target_size=(300, 150),  # All images will be resized to 300 x 150
        batch_size=BATCH_SIZE,
        # Specify the classes explicitly
        classes=['Bishop_Black', 'Bishop_White', 'Empty', 'King_Black', 'King_White', 'Knight_Black',
                 'Knight_White', 'Pawn_Black', 'Pawn_White', 'Queen_Black', 'Queen_White', 'Rook_Black', 'Rook_White'],
        # Since we use categorical_crossentropy loss, we need categorical labels
        class_mode='categorical')
    # Follow the same steps for validation generator
    validation_datagen = ImageDataGenerator(rescale=1/255)
    validation_generator = validation_datagen.flow_from_directory(
        folderpath+'validation',
        target_size=(300, 150),
        batch_size=BATCH_SIZE,
        class_mode='categorical')
    return (train_generator, validation_generator)


def create_model(optimizer=RMSprop(learning_rate=0.001)):
    '''Creates a CNN architecture and compiles it.'''
    model = Sequential([
        # Note the input shape is the desired size of the image 300 x 150 with 3 bytes color
        # The first convolution
        Conv2D(16, (3, 3), activation='relu', input_shape=(300, 150, 3)),
        MaxPooling2D(2, 2),
        # The second convolution
        Conv2D(32, (3, 3), activation='relu'),
        MaxPooling2D(2, 2),
        # The third convolution
        Conv2D(64, (3, 3), activation='relu'),
        MaxPooling2D(2, 2),
        # The fourth convolution
        Conv2D(64, (3, 3), activation='relu'),
        MaxPooling2D(2, 2),
        # The fifth convolution
        Conv2D(64, (3, 3), activation='relu'),
        MaxPooling2D(2, 2),
        # Flatten the results to feed into a dense layer
        Flatten(),
        # 128 neuron in the fully-connected layer
        Dense(128, activation='relu'),
        # 13 output neurons for 13 classes with the softmax activation
        Dense(13, activation='softmax')
    ])

    model.compile(loss='categorical_crossentropy',
                  optimizer=optimizer,
                  metrics=['acc'])

    return model


def fit_model(model, train_generator, validation_generator, callbacks=[], save=False, filename=""):
    '''Given the model and generators, trains the model and saves weights if
    needed. Callbacks can be provided to save intermediate results.
    Returns a history of model's performance (for plotting purpose).'''

    total_sample = train_generator.n

    history = model.fit(
        train_generator,
        steps_per_epoch=int(total_sample/BATCH_SIZE),
        epochs=NUM_EPOCHS,
        verbose=1,
        validation_data=validation_generator,
        callbacks=callbacks)

    if save:
        model.save_weights(filename)

    return history


def plot_accuracy(history):
    '''Given training history, plots accuracy of a model.'''
    plt.figure(figsize=(7, 4))
    plt.plot([i+1 for i in range(NUM_EPOCHS)],
             history.history['acc'], '-o', c='k', lw=2, markersize=9)
    plt.grid(True)
    plt.title("Training accuracy with epochs\n", fontsize=18)
    plt.xlabel("Training epochs", fontsize=15)
    plt.ylabel("Training accuracy", fontsize=15)
    plt.xticks(fontsize=15)
    plt.yticks(fontsize=15)
    plt.show()


def plot_loss(history):
    '''Given training history, plots loss of a model.'''
    plt.figure(figsize=(7, 4))
    plt.plot([i+1 for i in range(NUM_EPOCHS)],
             history.history['loss'], '-o', c='k', lw=2, markersize=9)
    plt.grid(True)
    plt.title("Training loss with epochs\n", fontsize=18)
    plt.xlabel("Training epochs", fontsize=15)
    plt.ylabel("Training loss", fontsize=15)
    plt.xticks(fontsize=15)
    plt.yticks(fontsize=15)
    plt.show()


def save_history(history, filename="./history.json"):
    '''Saves the given training history as a .json file.'''
    # Get the dictionary containing each metric and the loss for each epoch
    history_dict = history.history
    # Save it under the form of a json file
    json.dump(history_dict, open(filename, 'w'))


def load_history(filename="./history.json"):
    '''Loads training history from the path to a .json file. Returns a dict.'''
    with open(filename) as json_file:
        data = json.load(json_file)
    return data


def test_model(model):
    '''Tests the given model on the test set and prints its accuracy.
    Does not return anything.'''
    testdir = DATA_FOLDER + 'test'

    # IMPORTANT: Ordre EXACT comme dans le générateur (alphabétique)
    pieces = ['Bishop_Black', 'Bishop_White', 'Empty', 'King_Black', 'King_White', 'Knight_Black',
              'Knight_White', 'Pawn_Black', 'Pawn_White', 'Queen_Black', 'Queen_White', 'Rook_Black', 'Rook_White']
    score = 0
    total_size = 0
    for subdir, dirs, files in os.walk(testdir):
        for file in files:
            if file == ".DS_Store":
                continue
            piece = subdir.split('/')[-1]
            path = os.path.join(subdir, file)
            img = cv2.imread(path).reshape(1, 300, 150, 3) / 255.0
            y_prob = model.predict(img, verbose=0)
            y_pred = y_prob.argmax()
            if y_pred < 0 or y_pred >= len(pieces):
                print(y_pred, y_prob)
            if piece == pieces[y_pred]:
                score += 1
            total_size += 1
    print("TEST SET ACCURACY:", score/total_size)


if __name__ == '__main__':
    train_generator, validation_generator = create_generators(DATA_FOLDER)
    model = create_model()
    history = fit_model(model, train_generator,
                        validation_generator, save=False)
    save_history(history, "./history.json")
    # plot_accuracy(history)
    # plot_loss(history)
    test_model(model)
    model.save_weights('./model_weights.weights.h5')