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():

    print('in 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('in load_data')
    images = []
    labels = []
    dim = (IMG_WIDTH, IMG_HEIGHT)
    entries = os.listdir(data_dir)

    for entry in entries:
        path1 = os.path.join(data_dir, entry)
        signs = os.listdir(path1)

        for sign in signs:
            path2 = os.path.join(data_dir, entry, sign)
            image = cv2.imread(path2)
            img = cv2.resize(image, dim)
            images.append(img)
            labels.append(int(entry))

    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.
    """
    print('in get_model')
    model = tf.keras.models.Sequential([

    # First convolutional layer. Learn 50 filters using a 3x3 kernel
    tf.keras.layers.Conv2D(
        50, (3, 3), activation="relu", input_shape=(IMG_WIDTH, IMG_HEIGHT, 3)
    ),

    # Max-pooling layer, using 2x2 pool size
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),

    # Second convolutional layer. Learn 30 filters using a 3x3 kernel
    tf.keras.layers.Conv2D(
        30, (3, 3), activation="relu", input_shape=(IMG_WIDTH, IMG_HEIGHT, 3)
    ),

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

    # Add a hidden layer with dropout
    tf.keras.layers.Dense(250, activation="relu"),
    tf.keras.layers.Dropout(0.5),

    # Add an output layer with output units for all 43 digits
    tf.keras.layers.Dense(43, activation="softmax")
    ])

    # Train neural network
    model.compile(
        optimizer="adam",
        loss="categorical_crossentropy",
        metrics=["accuracy"]
    )

    return model


if __name__ == "__main__":
    main()

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():

    print('in 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('in load_data')
    images = []
    labels = []
    dim = (IMG_WIDTH, IMG_HEIGHT)
    entries = os.listdir(data_dir)

    for entry in entries:
        path1 = os.path.join(data_dir, entry)
        signs = os.listdir(path1)

        for sign in signs:
            path2 = os.path.join(data_dir, entry, sign)
            image = cv2.imread(path2)
            img = cv2.resize(image, dim)
            images.append(img)
            labels.append(int(entry))

    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.
    """
    print('in get_model')
    model = tf.keras.models.Sequential([

    # First convolutional layer. Learn 50 filters using a 3x3 kernel
    tf.keras.layers.Conv2D(
        50, (3, 3), activation="relu", input_shape=(IMG_WIDTH, IMG_HEIGHT, 3)
    ),

    # Max-pooling layer, using 2x2 pool size
    tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),

    # Second convolutional layer. Learn 30 filters using a 3x3 kernel
    tf.keras.layers.Conv2D(
        30, (3, 3), activation="relu", input_shape=(IMG_WIDTH, IMG_HEIGHT, 3)
    ),

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

    # Add a hidden layer with dropout
    tf.keras.layers.Dense(250, activation="relu"),
    tf.keras.layers.Dropout(0.5),

    # Add an output layer with output units for all 43 digits
    tf.keras.layers.Dense(43, activation="softmax")
    ])

    # Train neural network
    model.compile(
        optimizer="adam",
        loss="categorical_crossentropy",
        metrics=["accuracy"]
    )

    return model


if __name__ == "__main__":
    main()

# Check command-line arguments # Get image arrays and labels for all image files # Split data into training and testing sets # Get a compiled neural network # Fit model on training data # Evaluate neural network performance # Save model to file 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`. 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. # First convolutional layer. Learn 50 filters using a 3x3 kernel # Max-pooling layer, using 2x2 pool size # Second convolutional layer. Learn 30 filters using a 3x3 kernel # Flatten units # Add a hidden layer with dropout # Add an output layer with output units for all 43 digits # Train neural network