traffic.py

import cv2
import numpy as np
import os
import sys
import tensorflow as tf

from sklearn.model_selection import train_test_split

EPOCHS = 10
IMG_WIDTH = 30
IMG_HEIGHT = 30
NUM_CATEGORIES = 43
TEST_SIZE = 0.4


def main():

    # Check command-line arguments
    if len(sys.argv) not in [2, 3]:
        sys.exit("Usage: python traffic.py data_directory [model.h5]")

    # Get image arrays and labels for all image files
    images, labels = load_data(sys.argv[1])

    # Split data into training and testing sets
    labels = tf.keras.utils.to_categorical(labels)
    x_train, x_test, y_train, y_test = train_test_split(
        np.array(images), np.array(labels), test_size=TEST_SIZE
    )

    # Get a compiled neural network
    model = get_model()

    # Fit model on training data
    model.fit(x_train, y_train, epochs=EPOCHS)

    # Evaluate neural network performance
    model.evaluate(x_test,  y_test, verbose=2)

    # Save model to file
    if len(sys.argv) == 3:
        filename = sys.argv[2]
        model.save(filename)
        print(f"Model saved to {filename}.")


def load_data(data_dir):
    """
    Load image data from directory `data_dir`.

    Assume `data_dir` has one directory named after each category, numbered
    0 through NUM_CATEGORIES - 1. Inside each category directory will be some
    number of image files.

    Return tuple `(images, labels)`. `images` should be a list of all
    of the images in the data directory, where each image is formatted as a
    numpy ndarray with dimensions IMG_WIDTH x IMG_HEIGHT x 3. `labels` should
    be a list of integer labels, representing the categories for each of the
    corresponding `images`.
    """

    print(f'Loading images from dataset in directory "{data_dir}"')

    images = []
    labels = []

    # Iterate through sign folders in directory:
    for foldername in os.listdir(data_dir):
        # Error Checking Data Folder
        try:
            int(foldername)
        except ValueError:
            print("Warning! Non-integer folder name in data directory! Skipping...")
            continue
    # Iterate through images in each folder
        for filename in os.listdir(os.path.join(data_dir, foldername)):
            # Open each image and resize to be IMG_WIDTH X IMG HEIGHT
            img = cv2.imread(os.path.join(data_dir, foldername, filename))
            img = cv2.resize(img, (IMG_WIDTH, IMG_HEIGHT))

            # Normalise image pixel intensities:
            img=img/255

            # Append Resized Image and its label to lists
            images.append(img)
            labels.append(int(foldername))

    # Check number of Images Matches Number of Labels:
    if len(images) != len(labels):
        sys.exit('Error when loading data, number of images did not match number of labels!')
    else:
        print(f'{len(images)}, {len(labels)} labelled images loaded successfully from dataset!')

    return (images, labels)


def get_model():
    """
    Returns a compiled convolutional neural network model. Assume that the
    `input_shape` of the first layer is `(IMG_WIDTH, IMG_HEIGHT, 3)`.
    The output layer should have `NUM_CATEGORIES` units, one for each category.
    """

    # Create the Neural Network Model using keras:
    model = tf.keras.models.Sequential([

    # Add 2 sequential 64 filter, 3x3 Convolutional Layers Followed by 2x2 Pooling
    tf.keras.layers.Conv2D(64, (3, 3), activation="relu", input_shape=(IMG_WIDTH, IMG_HEIGHT, 3)),
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
    tf.keras.layers.Conv2D(64, (3, 3), activation="relu"),
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),

    # Flatten layers
    tf.keras.layers.Flatten(),

    # Add A Dense Hidden layer with 512 units and 50% dropout
    tf.keras.layers.Dense(512, activation="relu"),
    tf.keras.layers.Dropout(0.5),

    # Add Dense Output layer with 43 output units
    tf.keras.layers.Dense(NUM_CATEGORIES, activation="softmax")
    ])

    # Set additional model settings and compile:
    model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

    # Return model for training and testing
    return model


if __name__ == "__main__":
    main()