Video_Masker.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import argparse\n",
    "import imutils\n",
    "import time\n",
    "import cv2\n",
    "import os\n",
    "import sys\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "# Root directory of the project\n",
    "  # To find local version of the library\n",
    "from mrcnn import utils\n",
    "import mrcnn.model as modellib\n",
    "from mrcnn import visualize\n",
    "# Import COCO config\n",
    "import coco\n",
    "\n",
    "%matplotlib inline \n",
    "\n",
    "# Directory to save logs and trained model\n",
    "MODEL_DIR = \"logs\"\n",
    "\n",
    "# Local path to trained weights file\n",
    "COCO_MODEL_PATH = \"pretrained_model.h5\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    " \n",
    "# construct the argument parse and parse the arguments\n",
    "\"\"\"\n",
    "ap = argparse.ArgumentParser()\n",
    "ap.add_argument(\"-i\", \"--input\", required=True,\n",
    "\thelp=\"path to input video file\")\n",
    "ap.add_argument(\"-o\", \"--output\", required=True,\n",
    "\thelp=\"path to output video file\")\n",
    "ap.add_argument(\"-m\", \"--mask-rcnn\", required=True,\n",
    "\thelp=\"base path to mask-rcnn directory\")\n",
    "ap.add_argument(\"-c\", \"--confidence\", type=float, default=0.5,\n",
    "\thelp=\"minimum probability to filter weak detections\")\n",
    "ap.add_argument(\"-t\", \"--threshold\", type=float, default=0.3,\n",
    "\thelp=\"minimum threshold for pixel-wise mask segmentation\")\n",
    "args = vars(ap.parse_args())\n",
    "\"\"\"\n",
    "\n",
    "input_video_path = str(os.getcwd())+'/videos/traffic_pokfulam_road.mov'\n",
    "output_video_path = str(os.getcwd())+'/videos/masked_video_traffic.avi'\n",
    "maskrcnn_path = str(os.getcwd())\n",
    "CONFIDENCE = 0.5\n",
    "THRESHOLD_PIXEL = 0.3\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] loading Mask R-CNN from disk...\n"
     ]
    }
   ],
   "source": [
    "# load the COCO class labels our Mask R-CNN was trained on\n",
    "labelsPath = os.path.sep.join([maskrcnn_path,\"object_detection_classes_coco.txt\"])\n",
    "LABELS = open(labelsPath).read().strip().split(\"\\n\")\n",
    "\n",
    "# initialize a list of colors to represent each possible class label\n",
    "np.random.seed(42)\n",
    "COLORS = np.random.randint(0, 255, size=(len(LABELS), 3),dtype=\"uint8\")\n",
    " \n",
    "#derive the paths to the Mask R-CNN weights and model configuration\n",
    "weightsPath = os.path.sep.join([maskrcnn_path,\"frozen_inference_graph.pb\"])\n",
    "configPath = os.path.sep.join([maskrcnn_path,\"mask_rcnn_inception_v2_coco_2018_01_28.pbtxt\"])\n",
    "\n",
    "#class InferenceConfig(coco.CocoConfig):\n",
    "    # Set batch size to 1 since we'll be running inference on\n",
    "    # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU\n",
    "    #GPU_COUNT = 1\n",
    "    #IMAGES_PER_GPU = 1\n",
    "\n",
    "#config = InferenceConfig()\n",
    "#config.display()\n",
    "\n",
    "print(\"[INFO] loading Mask R-CNN from disk...\")\n",
    "#model = modellib.MaskRCNN(mode=\"inference\", model_dir=MODEL_DIR, config=config)\n",
    "\n",
    "# Load weights trained on MS-COCO\n",
    "#model.load_weights(COCO_MODEL_PATH, by_name=True)\n",
    "\n",
    "# load our Mask R-CNN trained on the COCO dataset (90 classes)\n",
    "# from disk\n",
    "\n",
    "net = cv2.dnn.readNetFromTensorflow(weightsPath, configPath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] 319 total frames in video\n"
     ]
    }
   ],
   "source": [
    "vs = cv2.VideoCapture(input_video_path)\n",
    "writer = None\n",
    " \n",
    "# try to determine the total number of frames in the video file\n",
    "try:\n",
    "\tprop = cv2.cv.CV_CAP_PROP_FRAME_COUNT if imutils.is_cv2() \\\n",
    "\t\telse cv2.CAP_PROP_FRAME_COUNT\n",
    "\ttotal = int(vs.get(prop))\n",
    "\tprint(\"[INFO] {} total frames in video\".format(total))\n",
    " \n",
    "# an error occurred while trying to determine the total\n",
    "# number of frames in the video file\n",
    "except:\n",
    "\tprint(\"[INFO] could not determine # of frames in video\")\n",
    "\ttotal = -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] single frame took 4.6508 seconds\n",
      "[INFO] estimated total time to finish: 1483.5956\n",
      "[INFO] cleaning up...\n"
     ]
    }
   ],
   "source": [
    "while True:\n",
    "\t# read the next frame from the file\n",
    "\t(grabbed, frame) = vs.read()\n",
    "\n",
    "\t# if the frame was not grabbed, then we have reached the end\n",
    "\t# of the stream\n",
    "\tif not grabbed:\n",
    "\t\tbreak\n",
    "\n",
    "\t# construct a blob from the input frame and then perform a\n",
    "\t# forward pass of the Mask R-CNN, giving us (1) the bounding box\n",
    "\t# coordinates of the objects in the image along with (2) the\n",
    "\t# pixel-wise segmentation for each specific object\n",
    "\tblob = cv2.dnn.blobFromImage(frame, swapRB=True, crop=False)\n",
    "\tnet.setInput(blob)\n",
    "\tstart = time.time()\n",
    "\t(boxes, masks) = net.forward([\"detection_out_final\",\n",
    "\t\t\"detection_masks\"])\n",
    "\tend = time.time()\n",
    "\n",
    "\t# loop over the number of detected objects\n",
    "\tfor i in range(0, boxes.shape[2]):\n",
    "\t\t# extract the class ID of the detection along with the\n",
    "\t\t# confidence (i.e., probability) associated with the\n",
    "\t\t# prediction\n",
    "\t\tclassID = int(boxes[0, 0, i, 1])\n",
    "\t\tconfidence = boxes[0, 0, i, 2]\n",
    "\n",
    "\t\t# filter out weak predictions by ensuring the detected\n",
    "\t\t# probability is greater than the minimum probability\n",
    "\t\tif confidence > CONFIDENCE:\n",
    "\t\t\t# scale the bounding box coordinates back relative to the\n",
    "\t\t\t# size of the frame and then compute the width and the\n",
    "\t\t\t# height of the bounding box\n",
    "\t\t\t(H, W) = frame.shape[:2]\n",
    "\t\t\tbox = boxes[0, 0, i, 3:7] * np.array([W, H, W, H])\n",
    "\t\t\t(startX, startY, endX, endY) = box.astype(\"int\")\n",
    "\t\t\tboxW = endX - startX\n",
    "\t\t\tboxH = endY - startY\n",
    "\n",
    "\t\t\t# extract the pixel-wise segmentation for the object,\n",
    "\t\t\t# resize the mask such that it's the same dimensions of\n",
    "\t\t\t# the bounding box, and then finally threshold to create\n",
    "\t\t\t# a *binary* mask\n",
    "\t\t\tmask = masks[i, classID]\n",
    "\t\t\tmask = cv2.resize(mask, (boxW, boxH),\n",
    "\t\t\t\tinterpolation=cv2.INTER_NEAREST)\n",
    "\t\t\tmask = (mask > THRESHOLD_PIXEL)\n",
    "\n",
    "\t\t\t# extract the ROI of the image but *only* extracted the\n",
    "\t\t\t# masked region of the ROI\n",
    "\t\t\troi = frame[startY:endY, startX:endX][mask]\n",
    "\n",
    "\t\t\t# grab the color used to visualize this particular class,\n",
    "\t\t\t# then create a transparent overlay by blending the color\n",
    "\t\t\t# with the ROI\n",
    "\t\t\tcolor = COLORS[classID]\n",
    "\t\t\tblended = ((0.4 * color) + (0.6 * roi)).astype(\"uint8\")\n",
    "\n",
    "\t\t\t# store the blended ROI in the original frame\n",
    "\t\t\tframe[startY:endY, startX:endX][mask] = blended\n",
    "\n",
    "\t\t\t# draw the bounding box of the instance on the frame\n",
    "\t\t\tcolor = [int(c) for c in color]\n",
    "\t\t\tcv2.rectangle(frame, (startX, startY), (endX, endY),\n",
    "\t\t\t\tcolor, 2)\n",
    "\n",
    "\t\t\t# draw the predicted label and associated probability of\n",
    "\t\t\t# the instance segmentation on the frame\n",
    "\t\t\ttext = \"{}: {:.4f}\".format(LABELS[classID], confidence)\n",
    "\t\t\tcv2.putText(frame, text, (startX, startY - 5),\n",
    "\t\t\t\tcv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)\n",
    "\n",
    "\t# check if the video writer is None\n",
    "\tif writer is None:\n",
    "\t\t# initialize our video writer\n",
    "\t\tfourcc = cv2.VideoWriter_fourcc(*\"MJPG\")\n",
    "\t\twriter = cv2.VideoWriter(output_video_path, fourcc, 30,\n",
    "\t\t\t(frame.shape[1], frame.shape[0]), True)\n",
    "\n",
    "\t\t# some information on processing single frame\n",
    "\t\tif total > 0:\n",
    "\t\t\telap = (end - start)\n",
    "\t\t\tprint(\"[INFO] single frame took {:.4f} seconds\".format(elap))\n",
    "\t\t\tprint(\"[INFO] estimated total time to finish: {:.4f}\".format(\n",
    "\t\t\t\telap * total))\n",
    "\n",
    "\t# write the output frame to disk\n",
    "\twriter.write(frame)\n",
    "\n",
    "# release the file pointers\n",
    "print(\"[INFO] cleaning up...\")\n",
    "writer.release()\n",
    "vs.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (py36)",
   "language": "python",
   "name": "py36"
  },
  "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}