Demo-Video.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "This demo uses mask-RCNN to accomplish the video processing. Input an original video (suppose this video include a person or some people), this program will recognize the main character and save him/her as RGB image while other background information will be changed as grayscale. \n",
    "\n",
    "|**Demo 1**|**Demo 2**|\n",
    "| :--: | :--: |\n",
    "|![](demo/fujing.gif)|![](demo/nikki.gif)|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hu/anaconda3/lib/python3.6/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\n",
      "  from ._conv import register_converters as _register_converters\n",
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "import os\n",
    "import sys\n",
    "from samples import coco\n",
    "from mrcnn import utils\n",
    "from mrcnn import model as modellib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare Model File and Configuration Information\n",
    "\n",
    "Now load the pre-trained model data (Mask-RCNN trained by COCO dataset)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the pre-trained model data\n",
    "ROOT_DIR = os.getcwd()\n",
    "MODEL_DIR = os.path.join(ROOT_DIR, \"logs\")\n",
    "COCO_MODEL_PATH = os.path.join(ROOT_DIR, \"mask_rcnn_coco.h5\")\n",
    "if not os.path.exists(COCO_MODEL_PATH):\n",
    "    utils.download_trained_weights(COCO_MODEL_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The original configuration information is saved in config.py file. It can be changed if necessary. \n",
    "\n",
    "It's better to use the default value, but you can also change the GPU information to suit the personal GPU well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change the config infermation\n",
    "class InferenceConfig(coco.CocoConfig):\n",
    "    GPU_COUNT = 1\n",
    "    \n",
    "    # Number of images to train with on each GPU. A 12GB GPU can typically\n",
    "    # handle 2 images of 1024x1024px.\n",
    "    # Adjust based on your GPU memory and image sizes. Use the highest\n",
    "    # number that your GPU can handle for best performance.\n",
    "    IMAGES_PER_GPU = 1\n",
    "    \n",
    "config = InferenceConfig()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# COCO dataset object names\n",
    "model = modellib.MaskRCNN(\n",
    "    mode=\"inference\", model_dir=MODEL_DIR, config=config\n",
    ")\n",
    "model.load_weights(COCO_MODEL_PATH, by_name=True)\n",
    "class_names = [\n",
    "    'BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',\n",
    "    'bus', 'train', 'truck', 'boat', 'traffic light',\n",
    "    'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',\n",
    "    'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',\n",
    "    'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',\n",
    "    'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',\n",
    "    'kite', 'baseball bat', 'baseball glove', 'skateboard',\n",
    "    'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',\n",
    "    'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',\n",
    "    'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',\n",
    "    'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',\n",
    "    'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',\n",
    "    'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',\n",
    "    'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',\n",
    "    'teddy bear', 'hair drier', 'toothbrush'\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define The Image Process Function\n",
    "\n",
    "Now define two image process functions to process each frame of the input video. \n",
    "\n",
    "apply_mask is used to change the background information to grayscale.\n",
    "\n",
    "display_instances is used to show the object detection result in original image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This function is used to change the colorful background information to grayscale.\n",
    "# image[:,:,0] is the Blue channel,image[:,:,1] is the Green channel, image[:,:,2] is the Red channel\n",
    "# mask == 0 means that this pixel is not belong to the object.\n",
    "# np.where function means that if the pixel belong to background, change it to gray_image.\n",
    "# Since the gray_image is 2D, for each pixel in background, we should set 3 channels to the same value to keep the grayscale.\n",
    "\n",
    "def apply_mask(image, mask):\n",
    "    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n",
    "    image[:, :, 0] = np.where(\n",
    "        mask == 0,\n",
    "        gray_image[:, :],\n",
    "        image[:, :, 0]\n",
    "    )\n",
    "    image[:, :, 1] = np.where(\n",
    "        mask == 0,\n",
    "        gray_image[:, :],\n",
    "        image[:, :, 1]\n",
    "    )\n",
    "    image[:, :, 2] = np.where(\n",
    "        mask == 0,\n",
    "        gray_image[:, :],\n",
    "        image[:, :, 2]\n",
    "    )\n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This function is used to show the object detection result in original image.\n",
    "def display_instances(image, boxes, masks, ids, names, scores):\n",
    "    # max_area will save the largest object for all the detection results\n",
    "    max_area = 0\n",
    "    \n",
    "    # n_instances saves the amount of all objects\n",
    "    n_instances = boxes.shape[0]\n",
    "\n",
    "    if not n_instances:\n",
    "        print('NO INSTANCES TO DISPLAY')\n",
    "    else:\n",
    "        assert boxes.shape[0] == masks.shape[-1] == ids.shape[0]\n",
    "\n",
    "    for i in range(n_instances):\n",
    "        if not np.any(boxes[i]):\n",
    "            continue\n",
    "\n",
    "        # compute the square of each object\n",
    "        y1, x1, y2, x2 = boxes[i]\n",
    "        square = (y2 - y1) * (x2 - x1)\n",
    "\n",
    "        # use label to select person object from all the 80 classes in COCO dataset\n",
    "        label = names[ids[i]]\n",
    "        if label == 'person':\n",
    "            # save the largest object in the image as main character\n",
    "            # other people will be regarded as background\n",
    "            if square > max_area:\n",
    "                max_area = square\n",
    "                mask = masks[:, :, i]\n",
    "            else:\n",
    "                continue\n",
    "        else:\n",
    "            continue\n",
    "\n",
    "        # apply mask for the image\n",
    "        image = apply_mask(image, mask)\n",
    "        \n",
    "    return image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Process Video\n",
    "\n",
    "Now use the functions above to accomplish the video processing and save the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hu/anaconda3/lib/python3.6/site-packages/skimage/transform/_warps.py:110: UserWarning: Anti-aliasing will be enabled by default in skimage 0.15 to avoid aliasing artifacts when down-sampling images.\n",
      "  warn(\"Anti-aliasing will be enabled by default in skimage 0.15 to \"\n"
     ]
    }
   ],
   "source": [
    "input_video = 'demo/1.mp4'\n",
    "capture = cv2.VideoCapture(input_video)\n",
    "\n",
    "# these 2 lines can be removed if you dont have a 1080p camera.\n",
    "# capture.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)\n",
    "# capture.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)\n",
    "\n",
    "# Recording Video\n",
    "fps = 25.0\n",
    "width = int(capture.get(3))\n",
    "height = int(capture.get(4))\n",
    "fcc = cv2.VideoWriter_fourcc('D', 'I', 'V', 'X')\n",
    "out = cv2.VideoWriter(\"new_video.avi\", fcc, fps, (width, height))\n",
    "\n",
    "while True:\n",
    "    ret, frame = capture.read()\n",
    "    results = model.detect([frame], verbose=0)\n",
    "    r = results[0]\n",
    "    frame = display_instances(\n",
    "        frame, r['rois'], r['masks'], r['class_ids'], class_names, r['scores']\n",
    "    )\n",
    "    cv2.imshow('video', frame)\n",
    "\n",
    "    # Recording Video\n",
    "    out.write(frame)\n",
    "\n",
    "    if cv2.waitKey(1) & 0xFF == ord('q'):\n",
    "        break\n",
    "\n",
    "capture.release()\n",
    "cv2.destroyAllWindows()"
   ]
  }
 ],
 "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}