A Convolutional neural Network in Python using Keras on TensorFlow

ai.py

# importing libraries 
from keras.preprocessing.image import ImageDataGenerator 
from keras.models import Sequential 
from keras.layers import Conv2D, MaxPooling2D 
from keras.layers import Activation, Dropout, Flatten, Dense
from keras.preprocessing import image
from keras import backend as K
from keras.models import load_model

import numpy as np
import sys
  
  
img_width, img_height = 224, 224
  
train_data_dir = 'v_data/train'
validation_data_dir = 'v_data/test'
nb_train_samples = 400
nb_validation_samples = 100
epochs = 10
batch_size = 16
  
if K.image_data_format() == 'channels_first': 
    input_shape = (3, img_width, img_height) 
else: 
    input_shape = (img_width, img_height, 3) 



if (len(sys.argv) > 2 and sys.argv[1] == '--load_model'):
    # LOAD EXISTING MODEL
    model_file = sys.argv[2]
    print(f"Loading existing model from {model_file}...")
    model = load_model(model_file)

    if (sys.argv[3] == '--test'):
        test_file = sys.argv[4]
        print(f"Testing model with image at {test_file}...")

        img_path = "v_data/"+test_file
        img = image.load_img(img_path, target_size=(img_width, img_height))
        img = image.img_to_array(img)
        img = np.expand_dims(img, axis=0)

        pred = model.predict(img)
        if (pred[0][0] == 0):
            print(f'Image @ "{img_path}" is of a car.')
        else:
            print(f'Image @ "{img_path}" is of a plane.')


else:
    # DEFINE MODEL
    print("Defining new model...")
    model = Sequential()

    # Convolution layers
    conv_filters = 32
    conv_layers = 1
    if (len(sys.argv) > 2 and sys.argv[1] == '--conv_layers'):
        conv_layers = int(sys.argv[2])

    for l in range(0, conv_layers):
        model.add(Conv2D(conv_filters, (3, 3), input_shape=input_shape))
        model.add(Activation('relu'))
        model.add(MaxPooling2D(pool_size=(2, 2)))

        # double filters each layer
        conv_filters *= 2


    model.add(Flatten())
    model.add(Dense(64))
    model.add(Activation('relu'))
    model.add(Dropout(0.5))
    model.add(Dense(1))

    model.add(Activation('sigmoid'))


    model.compile(
        loss='binary_crossentropy',
        optimizer='rmsprop',
        metrics=['accuracy']
    )

    train_datagen = ImageDataGenerator(
        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(
        train_data_dir,
        target_size=(img_width, img_height),
        batch_size=batch_size,
        class_mode='binary'
    )

    validation_generator = test_datagen.flow_from_directory(
        validation_data_dir,
        target_size=(img_width, img_height),
        batch_size=batch_size,
        class_mode='binary'
    )

    model.fit_generator(
        train_generator,
        steps_per_epoch=nb_train_samples // batch_size,
        epochs=epochs,
        validation_data=validation_generator,
        validation_steps=nb_validation_samples // batch_size
    )

    model_file = f"model_saved-{conv_layers}_conv.h5"
    model.save(model_file)
    print(f"Model with {conv_layers} convolution layers saved to '{model_file}'.")
    # DEFINE MODEL [END]