{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "This demo uses mask-RCNN to accomplish the image processing. Input an original image (suppose this image 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",
    "Here are some demo images:\n",
    "\n",
    "|**Original Image 1**|**New Image 1**|**Original Image 2**|**New Image 2**|\n",
    "| :--: | :--: | :--: | :--: |\n",
    "|![](demo/14.jpg)|![](demo/14N.jpg)|![](demo/1.jpg)|![](demo/1N.jpg)|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/yingshaoxo/.local/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 Image\n",
    "\n",
    "First, load the original image from disk and show it in window by using OpenCV."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input the original image name\n",
    "original_image = 'demo/1.jpg'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Read original image successfully! The original image shape is:\n",
      "(854, 1280, 3)\n",
      "Press ESC to exit or press s to save and exit.\n"
     ]
    }
   ],
   "source": [
    "# Use OpenCV to read and show the original image\n",
    "image = cv2.imread(original_image)\n",
    "cv2.imshow('original_image', image)\n",
    "\n",
    "print(\"Read original image successfully! The original image shape is:\")\n",
    "print(image.shape)\n",
    "print(\"Press ESC to exit or press s to save and exit.\")\n",
    "\n",
    "# Wait for keys to exit or save\n",
    "k = cv2.waitKey(0)\n",
    "if k == 27:                 \n",
    "    cv2.destroyAllWindows()\n",
    "elif k == ord('s'):        \n",
    "    cv2.imwrite('original_image.jpg', image)\n",
    "    cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, use cvtColor to accomplish image transformation from RGB image to gray image.\n",
    "\n",
    "**Attention: the original image is an colorful image which has 3 channels(RGB), while the gray image is a grayscale, which only has 2D format.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Change gray image successfully! The gray image shape is:\n",
      "(854, 1280)\n",
      "Press ESC to exit or press s to save and exit.\n"
     ]
    }
   ],
   "source": [
    "# Use cvtColor to accomplish image transformation from RGB image to gray image\n",
    "gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n",
    "cv2.imshow('gray_image', gray_image)\n",
    "\n",
    "print(\"Change gray image successfully! The gray image shape is:\")\n",
    "print(gray_image.shape)\n",
    "print(\"Press ESC to exit or press s to save and exit.\")\n",
    "\n",
    "# Wait for keys to exit or save\n",
    "k = cv2.waitKey(0)\n",
    "if k == 27:                 \n",
    "    cv2.destroyAllWindows()\n",
    "elif k == ord('s'):        \n",
    "    cv2.imwrite('gray_image.jpg', image)\n",
    "    cv2.destroyAllWindows()"
   ]
  },
  {
   "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": 5,
   "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": 6,
   "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": 7,
   "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. \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": 8,
   "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",
    "    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": 9,
   "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",
    "    # by mistake you put apply_mask inside for loop or you can write continue in if also\n",
    "    image = apply_mask(image, mask)\n",
    "        \n",
    "    return image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Process Image\n",
    "\n",
    "Now use the functions above to accomplish the image processing and save the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/dist-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": [
    "results = model.detect([image], verbose=0)\n",
    "r = results[0]\n",
    "frame = display_instances(\n",
    "    image, r['rois'], r['masks'], r['class_ids'], class_names, r['scores']\n",
    ")\n",
    "cv2.imshow('save_image', frame)\n",
    "\n",
    "# Wait for keys to exit or save\n",
    "k = cv2.waitKey(0)\n",
    "if k == 27:                 \n",
    "    cv2.destroyAllWindows()\n",
    "elif k == ord('s'):        \n",
    "    cv2.imwrite('save_image.jpg', image)\n",
    "    cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now try to get rid of background, then add a new background "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}