Kinect_streaming_segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# coding: utf-8\n",
    "import numpy as np\n",
    "import os\n",
    "import cv2\n",
    "import sys\n",
    "from pylibfreenect2 import Freenect2, SyncMultiFrameListener\n",
    "from pylibfreenect2 import FrameType, Registration, Frame\n",
    "from pylibfreenect2 import createConsoleLogger, setGlobalLogger\n",
    "from pylibfreenect2 import LoggerLevel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Packet pipeline: OpenGLPacketPipeline\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    from pylibfreenect2 import OpenGLPacketPipeline\n",
    "    pipeline = OpenGLPacketPipeline()\n",
    "except:\n",
    "    try:\n",
    "        from pylibfreenect2 import OpenCLPacketPipeline\n",
    "        pipeline = OpenCLPacketPipeline()\n",
    "    except:\n",
    "        from pylibfreenect2 import CpuPacketPipeline\n",
    "        pipeline = CpuPacketPipeline()\n",
    "print(\"Packet pipeline:\", type(pipeline).__name__)\n",
    "\n",
    "# Create and set logger\n",
    "logger = createConsoleLogger(LoggerLevel.Debug)\n",
    "setGlobalLogger(logger)\n",
    "\n",
    "fn = Freenect2()\n",
    "num_devices = fn.enumerateDevices()\n",
    "if num_devices == 0:\n",
    "    print(\"No device connected!\")\n",
    "    sys.exit(1)\n",
    "\n",
    "serial = fn.getDeviceSerialNumber(0)\n",
    "device = fn.openDevice(serial, pipeline=pipeline)\n",
    "\n",
    "listener = SyncMultiFrameListener(FrameType.Color | FrameType.Ir | FrameType.Depth)\n",
    "\n",
    "# Register listeners\n",
    "device.setColorFrameListener(listener)\n",
    "device.setIrAndDepthFrameListener(listener)\n",
    "\n",
    "device.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: must be called after device.start()\n",
    "registration = Registration(device.getIrCameraParams(),\n",
    "                            device.getColorCameraParams())\n",
    "\n",
    "# Optinal parameters for registration\n",
    "# set True if you need\n",
    "need_bigdepth = False\n",
    "need_color_depth_map = False\n",
    "\n",
    "#bigdepth = Frame(1920, 1082, 3) if need_bigdepth else None\n",
    "color_depth_map = np.zeros((424, 512),  np.int32).ravel() \\\n",
    "    if need_color_depth_map else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "# For Model\n",
    "import numpy as np\n",
    "import os\n",
    "import six.moves.urllib as urllib\n",
    "import sys\n",
    "import tarfile\n",
    "import tensorflow as tf\n",
    "import zipfile\n",
    "import cv2\n",
    "\n",
    "import random\n",
    "import math\n",
    "import skimage.io\n",
    "import matplotlib\n",
    "\n",
    "from collections import defaultdict\n",
    "from io import StringIO\n",
    "from matplotlib import pyplot as plt\n",
    "from PIL import Image\n",
    "\n",
    "# This is needed to display the images.\n",
    "%matplotlib inline\n",
    "sys.path.append(\"..\")\n",
    "\n",
    "# Root directory of the project\n",
    "ROOT_DIR = os.getcwd()\n",
    "\n",
    "# Import Mask RCNN# Import \n",
    "sys.path.append(ROOT_DIR)  # 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 SUN config\n",
    "sys.path.append(os.path.join(ROOT_DIR, \"samples/sun/\"))  # To find local version\n",
    "import sun\n",
    "\n",
    "# Directory to save logs and trained model\n",
    "MODEL_DIR = os.path.join(ROOT_DIR, \"logs\")\n",
    "\n",
    "# Local path to trained weights file\n",
    "SUN_MODEL_PATH = os.path.join(ROOT_DIR, \"mask_rcnn_sun.h5\")\n",
    "\n",
    "# Directory of images to run detection on\n",
    "IMAGE_DIR = os.path.join(ROOT_DIR, \"images\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Configurations:\n",
      "BACKBONE                       resnet101\n",
      "BACKBONE_STRIDES               [4, 8, 16, 32, 64]\n",
      "BATCH_SIZE                     1\n",
      "BBOX_STD_DEV                   [0.1 0.1 0.2 0.2]\n",
      "COMPUTE_BACKBONE_SHAPE         None\n",
      "DETECTION_MAX_INSTANCES        100\n",
      "DETECTION_MIN_CONFIDENCE       0.9\n",
      "DETECTION_NMS_THRESHOLD        0.3\n",
      "FPN_CLASSIF_FC_LAYERS_SIZE     1024\n",
      "GPU_COUNT                      1\n",
      "GRADIENT_CLIP_NORM             5.0\n",
      "IMAGES_PER_GPU                 1\n",
      "IMAGE_MAX_DIM                  1024\n",
      "IMAGE_META_SIZE                26\n",
      "IMAGE_MIN_DIM                  800\n",
      "IMAGE_MIN_SCALE                0\n",
      "IMAGE_RESIZE_MODE              square\n",
      "IMAGE_SHAPE                    [1024 1024    3]\n",
      "LEARNING_MOMENTUM              0.9\n",
      "LEARNING_RATE                  0.001\n",
      "LOSS_WEIGHTS                   {'rpn_class_loss': 1.0, 'rpn_bbox_loss': 1.0, 'mrcnn_class_loss': 1.0, 'mrcnn_bbox_loss': 1.0, 'mrcnn_mask_loss': 1.0}\n",
      "MASK_POOL_SIZE                 14\n",
      "MASK_SHAPE                     [28, 28]\n",
      "MAX_GT_INSTANCES               100\n",
      "MEAN_PIXEL                     [123.7 116.8 103.9]\n",
      "MINI_MASK_SHAPE                (56, 56)\n",
      "NAME                           sun\n",
      "NUM_CLASSES                    14\n",
      "POOL_SIZE                      7\n",
      "POST_NMS_ROIS_INFERENCE        1000\n",
      "POST_NMS_ROIS_TRAINING         2000\n",
      "ROI_POSITIVE_RATIO             0.33\n",
      "RPN_ANCHOR_RATIOS              [0.5, 1, 2]\n",
      "RPN_ANCHOR_SCALES              (32, 64, 128, 256, 512)\n",
      "RPN_ANCHOR_STRIDE              1\n",
      "RPN_BBOX_STD_DEV               [0.1 0.1 0.2 0.2]\n",
      "RPN_NMS_THRESHOLD              0.7\n",
      "RPN_TRAIN_ANCHORS_PER_IMAGE    256\n",
      "STEPS_PER_EPOCH                100\n",
      "TOP_DOWN_PYRAMID_SIZE          256\n",
      "TRAIN_BN                       False\n",
      "TRAIN_ROIS_PER_IMAGE           200\n",
      "USE_MINI_MASK                  True\n",
      "USE_RPN_ROIS                   True\n",
      "VALIDATION_STEPS               50\n",
      "WEIGHT_DECAY                   0.0001\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "class InferenceConfig(sun.SunConfig):\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()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Image segmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create model object in inference mode.\n",
    "model = modellib.MaskRCNN(mode=\"inference\", model_dir=MODEL_DIR, config=config)\n",
    "\n",
    "# Load weights trained on MS-SUN\n",
    "model.load_weights(SUN_MODEL_PATH, by_name=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class_names = ['BG', 'bed', 'books', 'ceiling', 'chair', 'floor',\n",
    "               'furniture', 'objects', 'picture', 'sofa', 'table',\n",
    "               'tv', 'wall', 'window']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def  random_colors(N):\n",
    "    np.random.seed(1)\n",
    "    colors = [tuple(255 * np.random.rand(3)) for _ in range(N)]\n",
    "    return colors\n",
    "\n",
    "def apply_mask(image, mask, color, alpha=0.5):\n",
    "    \"\"\"apply mask to image\"\"\"\n",
    "    for n, c in enumerate(color):\n",
    "        image[:, :, n] = np.where(\n",
    "            mask == 1,\n",
    "            image[:, :, n] * (1 - alpha) + alpha * c,\n",
    "            image[:, :, n]\n",
    "        )\n",
    "    return image\n",
    "\n",
    "def display_instances(image, boxes, masks, ids, names, scores):\n",
    "    \"\"\"\n",
    "        take the image and results and apply the mask, box, and Label\n",
    "    \"\"\"\n",
    "    n_instances = boxes.shape[0]\n",
    "    colors = random_colors(n_instances)\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, color in enumerate(colors):\n",
    "        if not np.any(boxes[i]):\n",
    "            continue\n",
    "\n",
    "        y1, x1, y2, x2 = boxes[i]\n",
    "        label = names[ids[i]]\n",
    "        score = scores[i] if scores is not None else None\n",
    "        caption = '{} {:.2f}'.format(label, score) if score else label\n",
    "        mask = masks[:, :, i]\n",
    "\n",
    "        image = apply_mask(image, mask, color)\n",
    "        image = cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)\n",
    "        image = cv2.putText(image, caption, (x1, y1), cv2.FONT_HERSHEY_COMPLEX, 0.7, color, 2)\n",
    "\n",
    "    return image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    from pylibfreenect2 import OpenGLPacketPipeline\n",
    "    pipeline = OpenGLPacketPipeline()\n",
    "except:\n",
    "    try:\n",
    "        from pylibfreenect2 import OpenCLPacketPipeline\n",
    "        pipeline = OpenCLPacketPipeline()\n",
    "    except:\n",
    "        from pylibfreenect2 import CpuPacketPipeline\n",
    "        pipeline = CpuPacketPipeline()\n",
    "print(\"Packet pipeline:\", type(pipeline).__name__)\n",
    "\n",
    "# Create and set logger\n",
    "logger = createConsoleLogger(LoggerLevel.Debug)\n",
    "setGlobalLogger(logger)\n",
    "\n",
    "fn = Freenect2()\n",
    "num_devices = fn.enumerateDevices()\n",
    "if num_devices == 0:\n",
    "    print(\"No device connected!\")\n",
    "    sys.exit(1)\n",
    "\n",
    "serial = fn.getDeviceSerialNumber(0)\n",
    "device = fn.openDevice(serial, pipeline=pipeline)\n",
    "\n",
    "listener = SyncMultiFrameListener(\n",
    "    FrameType.Color | FrameType.Ir | FrameType.Depth)\n",
    "\n",
    "# Register listeners\n",
    "device.setColorFrameListener(listener)\n",
    "device.setIrAndDepthFrameListener(listener)\n",
    "\n",
    "device.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: must be called after device.start()\n",
    "registration = Registration(device.getIrCameraParams(),\n",
    "                            device.getColorCameraParams())\n",
    "\n",
    "undistorted = Frame(512, 424, 4)\n",
    "registered = Frame(512, 424, 4)\n",
    "\n",
    "# Optinal parameters for registration\n",
    "# set True if you need\n",
    "need_bigdepth = False\n",
    "need_color_depth_map = False\n",
    "\n",
    "bigdepth = Frame(1920, 1082, 4) if need_bigdepth else None\n",
    "color_depth_map = np.zeros((424, 512),  np.int32).ravel() \\\n",
    "    if need_color_depth_map else None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Working with a video stream from Kinect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "while True:\n",
    "    frames = listener.waitForNewFrame()\n",
    "    frame = frames[\"color\"]\n",
    "    frame = frame.asarray(np.uint8)\n",
    "    frame = cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR)\n",
    "    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
    "    \n",
    "    results = model.detect([frame], verbose=1)\n",
    "    # Visualize results\n",
    "    r = results[0]\n",
    "    frame = display_instances(frame, r['rois'], r['masks'], r['class_ids'], class_names, r['scores'])\n",
    "    \n",
    "    cv2.imshow('image segmentation', cv2.resize(frame, (800,600)))\n",
    "    #listener.release(frame)\n",
    "\n",
    "    key = cv2.waitKey(delay=1)\n",
    "    if key == ord('q'):\n",
    "        break\n",
    "\n",
    "device.stop()\n",
    "device.close()\n",
    "sys.exit(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Working with writing all color frames on disc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "dirname = 'videos'\n",
    "os.mkdir(dirname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = 1\n",
    "while True:\n",
    "    frames = listener.waitForNewFrame()\n",
    "    #os.remove(\"/Users/ekaterina/Desktop/diploma/object_detection/frame.png\")\n",
    "    color = frames[\"color\"].asarray(np.uint8)\n",
    "    color = cv2.cvtColor(color, cv2.COLOR_BGRA2BGR)\n",
    "    \n",
    "    image_np_expanded = np.expand_dims(color, axis=0)\n",
    "    #cv2.imshow('object detection', color)\n",
    "    frame_name = \"frame\" + str(a) + \".png\"\n",
    "    cv2.imwrite(os.path.join(dirname, frame_name), cv2.resize(color, (1280,720)))\n",
    "    cv2.imshow('object detection', color)\n",
    "    a += 1\n",
    "    \n",
    "    key = cv2.waitKey(0) & 0xFF\n",
    "    if key == ord('q'):\n",
    "        break\n",
    "device.stop()\n",
    "device.close()\n",
    "\n",
    "sys.exit(0)"
   ]
  }
 ],
 "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
}