Using Neural Network.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using the trained neural network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Importing the Neural Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convolutional Neural Network is imported and the probability model is constructed for making predictions:\n",
    "conv_model = tf.keras.models.load_model('saved_model/CNNDiceModel2')\n",
    "conv_prob_mod = tf.keras.Sequential([conv_model, tf.keras.layers.Softmax()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the AI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, the same processing of image frames as the data-collecting stage is performed on the video feed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = cv2.SimpleBlobDetector_Params()\n",
    "\n",
    "params.filterByInertia\n",
    "params.minInertiaRatio = 0.5\n",
    "\n",
    "params.filterByArea\n",
    "params.minArea = 500\n",
    "\n",
    "BLUR = 19\n",
    "\n",
    "detector = cv2.SimpleBlobDetector_create(params)\n",
    "\n",
    "def get_blobs(frame):\n",
    "    frame_blurred = cv2.medianBlur(frame, BLUR)\n",
    "    blobs = detector.detect(frame_blurred)\n",
    "    return blobs\n",
    "\n",
    "def cropped(frame, blobs):\n",
    "    if len(blobs) == 1:\n",
    "        height = frame.shape[0]\n",
    "        width = frame.shape[1]\n",
    "        pos = blobs[0].pt\n",
    "        r = 20\n",
    "        #X and Y are adjusted for fish-eye effect of my webcam\n",
    "        x = int(pos[0] - r/2)\n",
    "        x = int(x + 0.035 * (x - width / 2.0))\n",
    "        y = int(pos[1] - r/2)\n",
    "        y = int(y + 0.035 * (y - height / 2.5))\n",
    "        \n",
    "        cropped_frame = frame[y:y+r, x:x+r]\n",
    "\n",
    "        return cropped_frame\n",
    "    \n",
    "def extractYellow(frame):\n",
    "    df = frame.copy()\n",
    "    redFrame = frame[:,:,2]\n",
    "    greenFrame = frame[:,:,1]\n",
    "    blueFrame = frame[:,:,0]\n",
    "    \n",
    "    result = np.zeros((20,20))\n",
    "    \n",
    "    for x in range(20):\n",
    "        for y in range(20):\n",
    "            red = redFrame[x,y]\n",
    "            blue = blueFrame[x,y]\n",
    "            green = greenFrame[x,y]\n",
    "            if red < 20 or green < 20 or blue < 90:\n",
    "                df[x,y,:] = 0\n",
    "            if blue > 1.1*red or blue > 3.0*green:\n",
    "                df[x,y,:] = 0\n",
    "            if blue > 1.3 * green and blue > 0.5 * red:\n",
    "                df[x,y,:] = 0\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Methods for applying the probability model to the frame\n",
    "def PredictSide(img):\n",
    "    img = np.array(img) / (255.0)\n",
    "    img = np.expand_dims(img,0)\n",
    "    conv_img = img.reshape(img.shape[0], img.shape[1], img.shape[2], 1)\n",
    "    pred = conv_prob_mod.predict(conv_img)\n",
    "    PredictedSide = np.argmax(pred[0]) + 1\n",
    "    Probability = np.max(pred[0]) * 100\n",
    "    return PredictedSide, Probability\n",
    "    \n",
    "def PredictionOverlay(frame, blobs, side, probability):\n",
    "    newFrame = frame.copy()\n",
    "    pos = blobs[0].pt\n",
    "    r = 30\n",
    "    x1 = int(pos[0] - r)\n",
    "    y1 = int(pos[1] - r)\n",
    "    x2 = int(pos[0] + r)\n",
    "    y2 = int(pos[1] + r)\n",
    "\n",
    "    newFrame = cv2.rectangle(newFrame, (x1,y1), (x2,y2), (0,0,255), 2)\n",
    "\n",
    "    newFrame = cv2.putText(newFrame, \"{} ({:3.0f}%)\".format(side, probability), (x1,y2 + 20), cv2.FONT_HERSHEY_SIMPLEX, \n",
    "                          0.5, (255, 0, 0), 2)\n",
    "    return newFrame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Video-Stream Loop:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "#initialize a video feed (argument is the OS-defined device number of desired camera)\n",
    "cap = cv2.VideoCapture(3)\n",
    "\n",
    "\n",
    "while(True):\n",
    "\n",
    "    ret, frame = cap.read()\n",
    "    \n",
    "    grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)\n",
    "    GBlur = cv2.GaussianBlur(grayscale, (15,15), 0)    \n",
    "    ret1, globalthreshold = cv2.threshold(GBlur,120,255,cv2.THRESH_BINARY)\n",
    "    M1Blur = cv2.medianBlur(globalthreshold, BLUR)\n",
    "    MBlur = cv2.medianBlur(M1Blur, BLUR)\n",
    "    \n",
    "    Gblobs = get_blobs(MBlur)\n",
    "\n",
    "    if len(Gblobs) == 1:\n",
    "        diceframe = cropped(frame, Gblobs)\n",
    "        cv2.imshow(\"Die\", diceframe)\n",
    "        \n",
    "        blue = diceframe[:,:,0]\n",
    "        green = diceframe[:,:,1]\n",
    "  \n",
    "        subtract = cv2.subtract(green, np.uint8(blue/8.0))\n",
    "        resizeSubtract = subtract.copy()\n",
    "        resizeSubtract = cv2.resize(resizeSubtract, (100,100))\n",
    "        cv2.imshow(\"subtract\", resizeSubtract)\n",
    "        \n",
    "        YellowFrame = extractYellow(diceframe)\n",
    "        YellowFrame = cv2.resize(YellowFrame, (100,100))\n",
    "        cv2.imshow(\"Yellow\", YellowFrame)\n",
    "        \n",
    "        #Send dice-cropped frame to Neural network probability model:\n",
    "        side, odds = PredictSide(subtract)\n",
    "        overlayedFrame = PredictionOverlay(frame, Gblobs, side, odds)\n",
    "        cv2.imshow(\"overlay\", overlayedFrame)\n",
    "    else:\n",
    "        cv2.imshow(\"overlay\", frame) #So that video feed is not cut off when there is no die in frame\n",
    "        \n",
    "    \n",
    "    res = cv2.waitKey(1)\n",
    "        \n",
    "    if res & 0xFF == ord('q'):\n",
    "        break\n",
    "        \n",
    "cap.release()\n",
    "cv2.destroyAllWindows()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}