In [ ]:
import os
import math
import matplotlib.pyplot as plt
from keras.models import Sequential
from keras.layers import Dense, Dropout, Flatten
from keras.layers import Conv2D, MaxPooling2D
from keras.optimizers import Adam, SGD
from keras.callbacks import ModelCheckpoint
from keras.preprocessing.image import load_img, img_to_array, ImageDataGenerator
from PIL import Image
In [ ]:
DATA_DIR = '../data/preprocessed/'
BATCH_SIZE = 32
BATCH_SIZE_TEST = 16
GRAYSCALE = True
INPUT_DIM = (64, 64, 1 if GRAYSCALE else 3)
AUGMENTATION_FACTOR = 3
EPOCHS = 100
In [ ]:
    train_datagen = ImageDataGenerator(
        rotation_range=10,
        rescale=1./255,
        shear_range=0.2,
        zoom_range=0.2,
        horizontal_flip=True)

test_datagen = ImageDataGenerator(rescale=1./255)

train_generator = train_datagen.flow_from_directory(
        DATA_DIR + 'train',
        target_size=INPUT_DIM[:2],
        batch_size=BATCH_SIZE,
        class_mode='categorical',
        color_mode='grayscale' if GRAYSCALE else 'rgb')

validation_generator = test_datagen.flow_from_directory(
        DATA_DIR + 'validation',
        target_size=INPUT_DIM[:2],
        batch_size=BATCH_SIZE_TEST,
        class_mode='categorical',
        color_mode='grayscale' if GRAYSCALE else 'rgb')

test_generator = test_datagen.flow_from_directory(
        DATA_DIR + 'test',
        target_size=INPUT_DIM[:2],
        batch_size=BATCH_SIZE_TEST,
        class_mode='categorical',
        color_mode='grayscale' if GRAYSCALE else 'rgb')

n_train = train_generator.n
n_validation = validation_generator.n
n_test = test_generator.n
In [ ]:
def get_model():
    model = Sequential()
    model.add(Conv2D(32, (3, 3), activation='relu', input_shape=INPUT_DIM))
    model.add(MaxPooling2D(pool_size=(2, 2)))

    model.add(Conv2D(32, (3, 3), activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))

    model.add(Conv2D(32, (3, 3), activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    
    model.add(Flatten())
    model.add(Dense(64, activation='relu'))
    model.add(Dropout(0.5))
    model.add(Dense(4, activation='softmax'))
    
    model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['acc'])
    return model
In [ ]:
model = get_model()
name_prefix = 'gray' if GRAYSCALE else 'rgb'
callbacks = [ModelCheckpoint(name_prefix + '-{epoch:02d}-{val_acc:.2f}.hdf5', monitor='val_acc', verbose=1, save_best_only=False, mode='max')]
In [ ]:
history = model.fit_generator(
        train_generator,
        steps_per_epoch= (n_train // BATCH_SIZE) * AUGMENTATION_FACTOR,
        epochs=EPOCHS,
        validation_data=validation_generator,
        validation_steps=n_validation // BATCH_SIZE_TEST,
        callbacks=callbacks)
In [ ]:
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()
In [ ]:
modelfiles = [f for f in os.listdir('.') if f.endswith('.hdf5') and f.startswith('rgb' if not GRAYSCALE else 'gray')]
for f in modelfiles:
    model.load_weights(f)
    result = model.evaluate_generator(
        test_generator,
        steps=n_test // BATCH_SIZE
    )
    print(f'{f}: {result[1]}')
In [ ]:
print(model.predict(img_to_array(examples_cs[0].convert('L')).reshape(1, *INPUT_DIM), batch_size=1))
print(model.predict(img_to_array(examples_econ[0].convert('L')).reshape(1, *INPUT_DIM), batch_size=1))
print(model.predict(img_to_array(examples_german[0].convert('L')).reshape(1, *INPUT_DIM), batch_size=1))
print(model.predict(img_to_array(examples_mechanical[0].convert('L')).reshape(1, *INPUT_DIM), batch_size=1))
In [ ]:
test_generator.class_indices

Some results

  • 64x64x3, 100 epochs, 3x augmentation, simple net: 0.571
  • 64x64x1, 100 epochs, 3x augmentation, simple net: 0.500
In [ ]:
get_model().summary()