train_cnn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rock-Paper-Scissors\n",
    "Training a CNN on rock-paper-scissors hand images to build a simple game application"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import modules -----------------------------------------------------------------------------------\n",
    "import os\n",
    "import cv2\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import random\n",
    "from datetime import datetime\n",
    "from keras import callbacks\n",
    "from keras.preprocessing import image\n",
    "from keras.preprocessing.image import ImageDataGenerator\n",
    "from keras.models import Sequential\n",
    "from keras.layers import InputLayer, Dropout, Flatten, Dense, Conv2D, MaxPooling2D\n",
    "from tensorflow.keras.models import load_model\n",
    "from sklearn.metrics import classification_report, confusion_matrix\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "# define paths -------------------------------------------------------------------------------------\n",
    "train_folder = 'dataset/train'\n",
    "test_folder = 'dataset/test'\n",
    "validation_folder = 'dataset/validation'\n",
    "\n",
    "# Helper methods -----------------------------------------------------------------------------------\n",
    "def load_image(img_path: str) -> np.ndarray:\n",
    "    '''Load an image from a file path and return it as a numpy array.'''\n",
    "    image = cv2.imread(img_path)\n",
    "    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    return image_rgb\n",
    "\n",
    "def display_image(image: np.ndarray) -> None:\n",
    "    '''Display an image in a matplotlib window.'''\n",
    "    plt.imshow(image)\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "\n",
    "def save_rps_model(model: Sequential, name: str) -> None:\n",
    "    '''Save a model to a file with a timestamped filename.'''\n",
    "    timestamp = datetime.now().strftime('%Y%m%d_%H%M')\n",
    "    filename = f'models/{name}_{timestamp}.h5'\n",
    "    model.save(filename)\n",
    "\n",
    "def load_rps_model(model_path: str = None):\n",
    "    '''Load a model from a file path. If no path is provided, the latest model is loaded.'''\n",
    "\n",
    "    if not os.path.isfile(str(model_path)):\n",
    "        # load latest version\n",
    "        model_path = os.path.join('models', sorted([m for m in os.listdir('models') if m.lower().endswith('.h5')])[-1])\n",
    "\n",
    "    model = load_model(model_path)\n",
    "    input_shape = model.layers[0].input_shape\n",
    "\n",
    "    print(f'Loaded model: {model_path} with input shape: {input_shape}')\n",
    "    return model\n",
    "\n",
    "def get_random_picture(set: str = None) -> str:\n",
    "    '''Return a random picture path from the train, test or validation set.'''\n",
    "\n",
    "    if set == 'train':\n",
    "        folder = train_folder\n",
    "    elif set == 'test':\n",
    "        folder = test_folder\n",
    "    elif set == 'validation':\n",
    "        folder = validation_folder\n",
    "    else:\n",
    "        folder = random.choice([train_folder, test_folder, validation_folder])\n",
    "    \n",
    "    subfolder = random.choice(os.listdir(folder))\n",
    "    subfolder_path = os.path.join(folder, subfolder)\n",
    "    test_image = random.choice(os.listdir(subfolder_path))\n",
    "    image_path = os.path.normpath(os.path.join(subfolder_path, test_image))\n",
    "\n",
    "    return image_path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check dataset content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# open test picture\n",
    "image_path = get_random_picture()\n",
    "print('Image:', image_path)\n",
    "img = load_image(image_path)\n",
    "display_image(img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing data for the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_shape = (150, 150, 3)\n",
    "batch_size = 32\n",
    "\n",
    "image_gen = ImageDataGenerator(\n",
    "    rescale=1./255,\n",
    "    rotation_range=90,\n",
    "    width_shift_range=0.2,\n",
    "    height_shift_range=0.2,\n",
    "    shear_range=0.2,\n",
    "    zoom_range=0.2,\n",
    "    horizontal_flip=True,\n",
    "    vertical_flip=False,  # True if you plan to show your hand upside down\n",
    "    fill_mode='nearest'\n",
    ")\n",
    "\n",
    "train_image_gen = image_gen.flow_from_directory(\n",
    "    train_folder,\n",
    "    target_size=image_shape[:2],\n",
    "    batch_size=batch_size,\n",
    "    class_mode='categorical',\n",
    "    shuffle=True  # enable shuffling to make the model more robust\n",
    ")\n",
    "\n",
    "test_image_gen = image_gen.flow_from_directory(\n",
    "    test_folder,\n",
    "    target_size=image_shape[:2],\n",
    "    batch_size=batch_size,\n",
    "    class_mode='categorical',\n",
    "    shuffle=False  # disable shuffling to keep data in same order as labels\n",
    ")\n",
    "\n",
    "validation_generator = image_gen.flow_from_directory(\n",
    "        validation_folder,\n",
    "        target_size=image_shape[:2],\n",
    "        batch_size=batch_size,\n",
    "        class_mode='categorical',\n",
    "        shuffle = False)  # disable shuffling to keep data in same order as labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test generator on random image\n",
    "random_img = image_gen.random_transform(img)\n",
    "display_image(random_img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create sequential modle\n",
    "model = Sequential()\n",
    "\n",
    "model.add(InputLayer(input_shape=image_shape))\n",
    "\n",
    "model.add(Conv2D(filters=32, kernel_size=(3, 3), strides=(1, 1), activation='relu'))\n",
    "model.add(MaxPooling2D(pool_size=(2, 2)))\n",
    "model.add(Conv2D(filters=64, kernel_size=(3, 3), activation='relu'))\n",
    "model.add(MaxPooling2D(pool_size=(2, 2)))\n",
    "model.add(Conv2D(filters=128, kernel_size=(3, 3), activation='relu'))\n",
    "model.add(MaxPooling2D(pool_size=(2, 2)))\n",
    "\n",
    "model.add(Flatten())\n",
    "\n",
    "model.add(Dense(units=256, activation='relu'))\n",
    "\n",
    "# Add 50% dropout to help reduce overfitting\n",
    "model.add(Dropout(0.5))\n",
    "\n",
    "# Output layer\n",
    "model.add(Dense(units=3, activation='softmax'))  # 3 classes\n",
    "\n",
    "model.compile(loss='categorical_crossentropy',\n",
    "              optimizer='adam',\n",
    "              metrics=['accuracy'])\n",
    "\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Moving the model to GPU\n",
    "WIP, I currently cannot see my GPU listed here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.python.client import device_lib\n",
    "print(device_lib.list_local_devices())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# indices\n",
    "class_indices = train_image_gen.class_indices\n",
    "class_labels = {v:k for k, v in class_indices.items()}\n",
    "print(class_indices)\n",
    "print(class_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initializing early stopping callback to monitor the validation loss and prevent overfitting\n",
    "earlystopping = callbacks.EarlyStopping(monitor=\"accuracy\",\n",
    "                                        mode=\"max\",\n",
    "                                        patience=10,\n",
    "                                        restore_best_weights=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the model\n",
    "results = model.fit(\n",
    "    train_image_gen,\n",
    "    epochs=100,\n",
    "    steps_per_epoch=train_image_gen.samples/train_image_gen.batch_size,\n",
    "    validation_data=test_image_gen,\n",
    "    validation_steps=test_image_gen.samples/test_image_gen.batch_size,\n",
    "    verbose=2,\n",
    "    callbacks=[earlystopping]\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize Accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(results.history['accuracy'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_rps_model(model, 'rps_v07_47[epoch]_0.9585[acc]_0.1195[loss]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Confusion Matrix and Classification Report"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = load_rps_model('models/rps_v06_56[epoch]_0.9641[acc]_0.1089[loss].h5')\n",
    "\n",
    "validation_steps = len(validation_generator)\n",
    "validation_generator.reset() # Reset the generator to be sure of avoiding shuffling\n",
    "\n",
    "predictions = model.predict_generator(validation_generator, steps=validation_steps, verbose=1)\n",
    "y_pred_classes = np.argmax(predictions, axis=1)\n",
    "\n",
    "# Print classification report\n",
    "print(\"Classification Report:\")\n",
    "print(classification_report(validation_generator.classes, y_pred_classes))\n",
    "\n",
    "# Print confusion matrix\n",
    "print(\"Confusion Matrix:\")\n",
    "print(confusion_matrix(validation_generator.classes, y_pred_classes))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Image prediction Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = load_rps_model()\n",
    "input_model_size = model.layers[0].input_shape[1:3]\n",
    "\n",
    "class_indices = {'paper': 0, 'rock': 1, 'scissors': 2}\n",
    "class_labels = {v:k for k, v in class_indices.items()}\n",
    "\n",
    "image_path = get_random_picture()\n",
    "print('Image:', image_path)\n",
    "img = image.load_img(image_path, target_size=input_model_size)\n",
    "display_image(img)\n",
    "\n",
    "img = image.img_to_array(img)\n",
    "img = np.expand_dims(img, axis=0)\n",
    "img = img/255\n",
    "\n",
    "prediction_prob = model.predict(img, verbose=0)[0]\n",
    "predicted_class_index = np.argmax(prediction_prob)\n",
    "predicted_class_label = class_labels[predicted_class_index]\n",
    "predicted_class_prob = prediction_prob[predicted_class_index]\n",
    "\n",
    "print('prediction_prob: ', [f'{k}: {prediction_prob[v]}' for k, v in class_indices.items()])\n",
    "print('predicted_class_index:', predicted_class_index)\n",
    "print('predicted_class_label:', predicted_class_label)\n",
    "print('predicted_class_prob:', predicted_class_prob)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Classify Validation Pictures\n",
    "Run all validation pictures through the model and classify them. Display the failed ones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = load_rps_model()\n",
    "input_model_size = model.layers[0].input_shape[1:3]\n",
    "\n",
    "success = 0\n",
    "failure = 0\n",
    "total = 0\n",
    "failed_images = []\n",
    "\n",
    "# load model\n",
    "class_indices = {'paper': 0, 'rock': 1, 'scissors': 2}\n",
    "class_labels = {v:k for k, v in class_indices.items()}\n",
    "\n",
    "test_images = None  # limit number of images to test if needed, None to go through all of them\n",
    "idx = 0\n",
    "\n",
    "for subfolder in ['rock', 'paper', 'scissors']:\n",
    "\n",
    "    current_folder = os.path.join(validation_folder, subfolder)\n",
    "\n",
    "    for img_name in os.listdir(current_folder):\n",
    "\n",
    "        if test_images is not None and idx > test_images - 1:\n",
    "            break\n",
    "        idx += 1\n",
    "\n",
    "        img_file = os.path.normpath(os.path.join(current_folder, img_name))\n",
    "\n",
    "        roi_copy = image.load_img(img_file, target_size=input_model_size)\n",
    "        img_array = image.img_to_array(roi_copy)\n",
    "\n",
    "        # add batch dimension\n",
    "        img_normalized = img_array/255\n",
    "        img_expanded = np.expand_dims(img_normalized, axis=0)\n",
    "\n",
    "        # image prediction probability\n",
    "        prediction_prob = model.predict(img_expanded, verbose=0)\n",
    "\n",
    "        # Get predicted class index\n",
    "        predicted_class_index = np.argmax(prediction_prob)\n",
    "\n",
    "        # Map predicted class index to class label\n",
    "        predicted_class_label = class_labels[predicted_class_index]\n",
    "\n",
    "        # Print predicted class label\n",
    "        if subfolder == predicted_class_label:\n",
    "            success += 1\n",
    "            result = 'success'\n",
    "        else:\n",
    "            failure += 1\n",
    "            result = 'failure'\n",
    "            failed_images.append((img_normalized, predicted_class_label, img_file))\n",
    "\n",
    "        total += 1\n",
    "        to_print = f'Images Classified: {total} | Success: {success} | Failure: {failure} | Accuracy: {success/total*100:.2f}%'\n",
    "        print(f'{to_print}', end='\\r', flush=True)\n",
    "\n",
    "# print failed images with prediction and path\n",
    "if failed_images:\n",
    "    print(f\"\\n{50 * '-'}\\nFAILED IMAGES:\")\n",
    "    for roi_copy, label, model_path in failed_images:\n",
    "        print(model_path)\n",
    "        img_size = 200\n",
    "        plt.figure(figsize=(img_size/100, img_size/100))\n",
    "        plt.imshow(roi_copy, cmap='gray')\n",
    "        plt.title(label, loc='left')\n",
    "        plt.axis('off')\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test on camera\n",
    "Define a ROI from webcam frames and predict the class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = load_rps_model()\n",
    "input_model_size = model.layers[0].input_shape[1:3]\n",
    "\n",
    "class_labels = {0: 'paper', 1: 'rock', 2: 'scissors'}\n",
    "\n",
    "cap = cv2.VideoCapture(0)\n",
    "cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)\n",
    "cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)\n",
    "cv2.namedWindow('Rock Paper Scissors')\n",
    "\n",
    "# define ROI position and size\n",
    "x, y, w, h = 895, 78, 300, 300\n",
    "\n",
    "while True:\n",
    "\n",
    "    try:\n",
    "        ret, frame = cap.read()\n",
    "        frame = cv2.flip(frame, 1)\n",
    "        roi = frame[y:y+h, x:x+w]\n",
    "\n",
    "        key = cv2.waitKey(1) & 0xFF\n",
    "        if key == ord('q'):\n",
    "            break\n",
    "\n",
    "        # convert roi to RGB, resize and normalize\n",
    "        roi_resized = cv2.resize(roi, input_model_size, interpolation=cv2.INTER_AREA)\n",
    "        roi_rgb = cv2.cvtColor(roi_resized, cv2.COLOR_BGR2RGB)\n",
    "        roi_normalized = roi_rgb / 255.0\n",
    "        roi_expanded = np.expand_dims(roi_normalized, axis=0)\n",
    "\n",
    "        # Classify the ROI image\n",
    "        prediction_prob = model.predict(roi_expanded, verbose=0)[0]\n",
    "        predicted_class_index = np.argmax(prediction_prob)\n",
    "        predicted_class_label = class_labels[predicted_class_index]\n",
    "        predicted_class_prob = prediction_prob[predicted_class_index]\n",
    "\n",
    "        # display captured roi top left of the frame for debugging\n",
    "        frame[0:input_model_size[0], 0:input_model_size[1]] = roi_resized\n",
    "\n",
    "        # Display prediction\n",
    "        if predicted_class_prob < 0.9:\n",
    "            label = 'Undefined'\n",
    "        else:\n",
    "            label = f'{predicted_class_label} ({predicted_class_prob:.2f})'\n",
    "\n",
    "        p_paper, p_rock, p_scissors = (f'{p:.2f}' for p in prediction_prob)\n",
    "        print(f'Prediction: Rock {p_rock} - Paper {p_paper} - Scissors {p_scissors}', end='\\r')\n",
    "\n",
    "        # Draw label\n",
    "        position = x, y + h + 20\n",
    "        color = 0, 0, 255\n",
    "        cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)\n",
    "        cv2.putText(frame, label, position, cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA)\n",
    "\n",
    "        # Display frame\n",
    "        cv2.imshow('Rock Paper Scissors', frame)\n",
    "\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        break\n",
    "\n",
    "cap.release()\n",
    "cv2.destroyAllWindows()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "keras",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}