Advertisement
Guest User

traffic

a guest
Jun 12th, 2021
349
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 3.51 KB | None | 0 0
  1. import cv2
  2. import numpy as np
  3. import os
  4. import sys
  5. import tensorflow as tf
  6.  
  7. from sklearn.model_selection import train_test_split
  8.  
  9. EPOCHS = 10
  10. IMG_WIDTH = 30
  11. IMG_HEIGHT = 30
  12. NUM_CATEGORIES = 43
  13. TEST_SIZE = 0.4
  14.  
  15.  
  16. def main():
  17.  
  18.     # Check command-line arguments
  19.     if len(sys.argv) not in [2, 3]:
  20.         sys.exit("Usage: python traffic.py data_directory [model.h5]")
  21.  
  22.     # Get image arrays and labels for all image files
  23.     images, labels = load_data(sys.argv[1])
  24.  
  25.     # Split data into training and testing sets
  26.     labels = tf.keras.utils.to_categorical(labels)
  27.     x_train, x_test, y_train, y_test = train_test_split(
  28.         np.array(images), np.array(labels), test_size=TEST_SIZE
  29.     )
  30.  
  31.     # Get a compiled neural network
  32.     model = get_model()
  33.  
  34.     # Fit model on training data
  35.     model.fit(x_train, y_train, epochs=EPOCHS)
  36.  
  37.     # Evaluate neural network performance
  38.     model.evaluate(x_test,  y_test, verbose=2)
  39.  
  40.     # Save model to file
  41.     if len(sys.argv) == 3:
  42.         filename = sys.argv[2]
  43.         model.save(filename)
  44.         print(f"Model saved to {filename}.")
  45.  
  46.  
  47. def load_data(data_dir):
  48.     """
  49.    Load image data from directory `data_dir`.
  50.  
  51.    Assume `data_dir` has one directory named after each category, numbered
  52.    0 through NUM_CATEGORIES - 1. Inside each category directory will be some
  53.    number of image files.
  54.  
  55.    Return tuple `(images, labels)`. `images` should be a list of all
  56.    of the images in the data directory, where each image is formatted as a
  57.    numpy ndarray with dimensions IMG_WIDTH x IMG_HEIGHT x 3. `labels` should
  58.    be a list of integer labels, representing the categories for each of the
  59.    corresponding `images`.
  60.    """
  61.  
  62.     images = list()
  63.     label = list()
  64.     num = 0
  65.     # Get path of folder
  66.     for root, _, files in os.walk(data_dir):
  67.         for file in files:
  68.             # If file is ppm image
  69.             if file.endswith("ppm"):
  70.                 img = cv2.imread(os.path.join(root, file))
  71.                 img2 = cv2.resize(img, (IMG_WIDTH, IMG_HEIGHT))
  72.  
  73.                 images.append(img2)
  74.                 label.append(int(os.path.basename(root)))
  75.  
  76.     return images, label
  77.  
  78.  
  79. def get_model():
  80.     """
  81.    Returns a compiled convolutional neural network model. Assume that the
  82.    `input_shape` of the first layer is `(IMG_WIDTH, IMG_HEIGHT, 3)`.
  83.    The output layer should have `NUM_CATEGORIES` units, one for each category.
  84.    """
  85.     # Create a convolutional neural network
  86.     model = tf.keras.models.Sequential([
  87.  
  88.         # Convolutional layer. Learn 32 filters using a 3x3 kernel
  89.         tf.keras.layers.Conv2D(
  90.             32, (3, 3), activation="relu", input_shape=(IMG_WIDTH, IMG_HEIGHT, 3)
  91.         ),
  92.  
  93.         # Max-pooling layer, using 2x2 pool size
  94.         tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
  95.  
  96.         # 2nd convolutional layer, same as first
  97.         tf.keras.layers.Conv2D(
  98.             32, (3, 3), activation="relu", input_shape=(IMG_WIDTH, IMG_HEIGHT, 3)
  99.         ),
  100.  
  101.         # Flatten units
  102.         tf.keras.layers.Flatten(),
  103.  
  104.         # Add a hidden layer with dropout
  105.         tf.keras.layers.Dense(128, activation="relu"),
  106.         tf.keras.layers.Dropout(0.5),
  107.  
  108.         # Add an output layer
  109.         tf.keras.layers.Dense(NUM_CATEGORIES, activation="softmax")
  110.     ])
  111.  
  112.     # Train neural network
  113.     model.compile(
  114.         optimizer="adam",
  115.         loss="categorical_crossentropy",
  116.         metrics=["accuracy"]
  117.     )
  118.     return model
  119.  
  120.  
  121. if __name__ == "__main__":
  122.     main()
  123.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement