notebooks/201-vision-monodepth/201-vision-monodepth.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c2cd5c88",
   "metadata": {
    "id": "moved-collapse",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Monodepth Estimation with OpenVINO\n",
    "\n",
    "This tutorial demonstrates Monocular Depth Estimation with MidasNet in OpenVINO. Model information can be found [here](https://docs.openvino.ai/latest/omz_models_model_midasnet.html).\n",
    "\n",
    "![monodepth](https://user-images.githubusercontent.com/36741649/127173017-a0bbcf75-db24-4d2c-81b9-616e04ab7cd9.gif)\n",
    "\n",
    "### What is Monodepth?\n",
    "Monocular Depth Estimation is the task of estimating scene depth using a single image. It has many potential applications in robotics, 3D reconstruction, medical imaging and autonomous systems. This tutorial uses a neural network model called [MiDaS](https://github.com/intel-isl/MiDaS), which was developed by the [Embodied AI Foundation](https://www.embodiedaifoundation.org/). See the research paper below to learn more.\n",
    "\n",
    "R. Ranftl, K. Lasinger, D. Hafner, K. Schindler and V. Koltun, [\"Towards Robust Monocular Depth Estimation: Mixing Datasets for Zero-shot Cross-dataset Transfer,\"](https://ieeexplore.ieee.org/document/9178977) in IEEE Transactions on Pattern Analysis and Machine Intelligence, doi: 10.1109/TPAMI.2020.3019967."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23fa9f61",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Preparation\n",
    "\n",
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25d08d89",
   "metadata": {
    "id": "ahead-spider"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "import time\n",
    "from pathlib import Path\n",
    "\n",
    "import cv2\n",
    "import matplotlib.cm\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from IPython.display import (\n",
    "    HTML,\n",
    "    FileLink,\n",
    "    Pretty,\n",
    "    ProgressBar,\n",
    "    Video,\n",
    "    clear_output,\n",
    "    display,\n",
    ")\n",
    "from openvino.runtime import Core\n",
    "\n",
    "sys.path.append(\"../utils\")\n",
    "from notebook_utils import load_image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57a29fe4",
   "metadata": {
    "id": "contained-office"
   },
   "source": [
    "### Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94499763",
   "metadata": {
    "id": "amber-lithuania"
   },
   "outputs": [],
   "source": [
    "DEVICE = \"CPU\"\n",
    "MODEL_FILE = \"model/MiDaS_small.xml\"\n",
    "\n",
    "model_xml_path = Path(MODEL_FILE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46a16ede",
   "metadata": {},
   "source": [
    "## Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c9f693b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize_minmax(data):\n",
    "    \"\"\"Normalizes the values in `data` between 0 and 1\"\"\"\n",
    "    return (data - data.min()) / (data.max() - data.min())\n",
    "\n",
    "\n",
    "def convert_result_to_image(result, colormap=\"viridis\"):\n",
    "    \"\"\"\n",
    "    Convert network result of floating point numbers to an RGB image with\n",
    "    integer values from 0-255 by applying a colormap.\n",
    "\n",
    "    `result` is expected to be a single network result in 1,H,W shape\n",
    "    `colormap` is a matplotlib colormap.\n",
    "    See https://matplotlib.org/stable/tutorials/colors/colormaps.html\n",
    "    \"\"\"\n",
    "    cmap = matplotlib.cm.get_cmap(colormap)\n",
    "    result = result.squeeze(0)\n",
    "    result = normalize_minmax(result)\n",
    "    result = cmap(result)[:, :, :3] * 255\n",
    "    result = result.astype(np.uint8)\n",
    "    return result\n",
    "\n",
    "\n",
    "def to_rgb(image_data) -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Convert image_data from BGR to RGB\n",
    "    \"\"\"\n",
    "    return cv2.cvtColor(image_data, cv2.COLOR_BGR2RGB)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7799cbb7",
   "metadata": {
    "id": "sensitive-wagner"
   },
   "source": [
    "## Load the Model\n",
    "\n",
    "Load the model in OpenVINO Runtime with `ie.read_model` and compile it for the specified device with `ie.compile_model`. Get input and output keys and the expected input shape for the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e9e708e",
   "metadata": {
    "id": "complete-brother"
   },
   "outputs": [],
   "source": [
    "ie = Core()\n",
    "ie.set_property({'CACHE_DIR': '../cache'})\n",
    "model = ie.read_model(model_xml_path)\n",
    "compiled_model = ie.compile_model(model=model, device_name=DEVICE)\n",
    "\n",
    "input_key = compiled_model.input(0)\n",
    "output_key = compiled_model.output(0)\n",
    "\n",
    "network_input_shape = list(input_key.shape)\n",
    "network_image_height, network_image_width = network_input_shape[2:]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe5a7171",
   "metadata": {
    "id": "compact-bargain"
   },
   "source": [
    "## Monodepth on Image\n",
    "\n",
    "### Load, resize and reshape input image\n",
    "\n",
    "The input image is read with OpenCV, resized to network input size, and reshaped to (N,C,H,W) (N=number of images,  C=number of channels, H=height, W=width). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1c4ec2d",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "central-psychology",
    "outputId": "d864ee96-3fbd-488d-da1a-88e730f34aad",
    "tags": []
   },
   "outputs": [],
   "source": [
    "IMAGE_FILE = \"../data/image/coco_bike.jpg\"\n",
    "image = load_image(path=IMAGE_FILE)\n",
    "\n",
    "# Resize to input shape for network.\n",
    "resized_image = cv2.resize(src=image, dsize=(network_image_height, network_image_width))\n",
    "\n",
    "# Reshape the image to network input shape NCHW.\n",
    "input_image = np.expand_dims(np.transpose(resized_image, (2, 0, 1)), 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de672364",
   "metadata": {
    "id": "taken-spanking"
   },
   "source": [
    "### Do inference on the image\n",
    "\n",
    "Do inference, convert the result to an image, and resize it to the original image shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c048d5c",
   "metadata": {
    "id": "banner-kruger"
   },
   "outputs": [],
   "source": [
    "result = compiled_model([input_image])[output_key]\n",
    "\n",
    "# Convert the network result of disparity map to an image that shows\n",
    "# distance as colors.\n",
    "result_image = convert_result_to_image(result=result)\n",
    "\n",
    "# Resize back to original image shape. The `cv2.resize` function expects shape\n",
    "# in (width, height), [::-1] reverses the (height, width) shape to match this.\n",
    "result_image = cv2.resize(result_image, image.shape[:2][::-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4989349b",
   "metadata": {},
   "source": [
    "### Display monodepth image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f818fe7b",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 867
    },
    "id": "ranging-executive",
    "outputId": "30373e8e-34e9-4820-e32d-764aa99d4b25"
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 2, figsize=(20, 15))\n",
    "ax[0].imshow(to_rgb(image))\n",
    "ax[1].imshow(result_image);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63bd66b3",
   "metadata": {
    "id": "descending-cache"
   },
   "source": [
    "## Monodepth on Video\n",
    "\n",
    "By default, only the first 100 frames are processed in order to quickly check that everything works. Change `NUM_FRAMES` in the cell below to modify this. Set `NUM_FRAMES` to 0 to process the whole video."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb1cd959",
   "metadata": {},
   "source": [
    "### Video Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "692f8b75",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "terminal-dividend",
    "outputId": "87f5ada0-8caf-49c3-fe54-626e2b1967f3"
   },
   "outputs": [],
   "source": [
    "# Video source: https://www.youtube.com/watch?v=fu1xcQdJRws (Public Domain)\n",
    "VIDEO_FILE = \"../data/video/Coco Walking in Berkeley.mp4\"\n",
    "# Number of seconds of input video to process. Set `NUM_SECONDS` to 0 to process\n",
    "# the full video.\n",
    "NUM_SECONDS = 4\n",
    "# Set `ADVANCE_FRAMES` to 1 to process every frame from the input video\n",
    "# Set `ADVANCE_FRAMES` to 2 to process every second frame. This reduces\n",
    "# the time it takes to process the video.\n",
    "ADVANCE_FRAMES = 2\n",
    "# Set `SCALE_OUTPUT` to reduce the size of the result video\n",
    "# If `SCALE_OUTPUT` is 0.5, the width and height of the result video\n",
    "# will be half the width and height of the input video.\n",
    "SCALE_OUTPUT = 0.5\n",
    "# The format to use for video encoding. The 'vp09` is slow,\n",
    "# but it works on most systems.\n",
    "# Try the `THEO` encoding if you have FFMPEG installed.\n",
    "# FOURCC = cv2.VideoWriter_fourcc(*\"THEO\")\n",
    "FOURCC = cv2.VideoWriter_fourcc(*\"vp09\")\n",
    "\n",
    "# Create Path objects for the input video and the result video.\n",
    "output_directory = Path(\"output\")\n",
    "output_directory.mkdir(exist_ok=True)\n",
    "result_video_path = output_directory / f\"{Path(VIDEO_FILE).stem}_monodepth.mp4\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "babf160c",
   "metadata": {},
   "source": [
    "### Load the Video\n",
    "\n",
    "Load the video from a `VIDEO_FILE`, set in the *Video Settings* cell above. Open the video to read the frame width and height and fps, and compute values for these properties for the monodepth video."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11a80192",
   "metadata": {},
   "outputs": [],
   "source": [
    "cap = cv2.VideoCapture(str(VIDEO_FILE))\n",
    "ret, image = cap.read()\n",
    "if not ret:\n",
    "    raise ValueError(f\"The video at {VIDEO_FILE} cannot be read.\")\n",
    "input_fps = cap.get(cv2.CAP_PROP_FPS)\n",
    "input_video_frame_height, input_video_frame_width = image.shape[:2]\n",
    "\n",
    "target_fps = input_fps / ADVANCE_FRAMES\n",
    "target_frame_height = int(input_video_frame_height * SCALE_OUTPUT)\n",
    "target_frame_width = int(input_video_frame_width * SCALE_OUTPUT)\n",
    "\n",
    "cap.release()\n",
    "print(\n",
    "    f\"The input video has a frame width of {input_video_frame_width}, \"\n",
    "    f\"frame height of {input_video_frame_height} and runs at {input_fps:.2f} fps\"\n",
    ")\n",
    "print(\n",
    "    \"The monodepth video will be scaled with a factor \"\n",
    "    f\"{SCALE_OUTPUT}, have width {target_frame_width}, \"\n",
    "    f\" height {target_frame_height}, and run at {target_fps:.2f} fps\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "425ef846",
   "metadata": {},
   "source": [
    "### Do Inference on a Video and Create Monodepth Video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e9e62ea",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "present-albany",
    "outputId": "600edb69-af12-44dc-ec8e-95005b74179c",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Initialize variables.\n",
    "input_video_frame_nr = 0\n",
    "start_time = time.perf_counter()\n",
    "total_inference_duration = 0\n",
    "\n",
    "# Open the input video\n",
    "cap = cv2.VideoCapture(str(VIDEO_FILE))\n",
    "\n",
    "# Create a result video.\n",
    "out_video = cv2.VideoWriter(\n",
    "    str(result_video_path),\n",
    "    FOURCC,\n",
    "    target_fps,\n",
    "    (target_frame_width * 2, target_frame_height),\n",
    ")\n",
    "\n",
    "num_frames = int(NUM_SECONDS * input_fps)\n",
    "total_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) if num_frames == 0 else num_frames\n",
    "progress_bar = ProgressBar(total=total_frames)\n",
    "progress_bar.display()\n",
    "\n",
    "try:\n",
    "    while cap.isOpened():\n",
    "        ret, image = cap.read()\n",
    "        if not ret:\n",
    "            cap.release()\n",
    "            break\n",
    "\n",
    "        if input_video_frame_nr >= total_frames:\n",
    "            break\n",
    "\n",
    "        # Only process every second frame.\n",
    "        # Prepare a frame for inference.\n",
    "        # Resize to the input shape for network.\n",
    "        resized_image = cv2.resize(src=image, dsize=(network_image_height, network_image_width))\n",
    "        # Reshape the image to network input shape NCHW.\n",
    "        input_image = np.expand_dims(np.transpose(resized_image, (2, 0, 1)), 0)\n",
    "\n",
    "        # Do inference.\n",
    "        inference_start_time = time.perf_counter()\n",
    "        result = compiled_model([input_image])[output_key]\n",
    "        inference_stop_time = time.perf_counter()\n",
    "        inference_duration = inference_stop_time - inference_start_time\n",
    "        total_inference_duration += inference_duration\n",
    "\n",
    "        if input_video_frame_nr % (10 * ADVANCE_FRAMES) == 0:\n",
    "            clear_output(wait=True)\n",
    "            progress_bar.display()\n",
    "            # input_video_frame_nr // ADVANCE_FRAMES gives the number of\n",
    "            # Frames that have been processed by the network.\n",
    "            display(\n",
    "                Pretty(\n",
    "                    f\"Processed frame {input_video_frame_nr // ADVANCE_FRAMES}\"\n",
    "                    f\"/{total_frames // ADVANCE_FRAMES}. \"\n",
    "                    f\"Inference time per frame: {inference_duration:.2f} seconds \"\n",
    "                    f\"({1/inference_duration:.2f} FPS)\"\n",
    "                )\n",
    "            )\n",
    "\n",
    "        # Transform the network result to a RGB image.\n",
    "        result_frame = to_rgb(convert_result_to_image(result))\n",
    "        # Resize the image and the result to a target frame shape.\n",
    "        result_frame = cv2.resize(result_frame, (target_frame_width, target_frame_height))\n",
    "        image = cv2.resize(image, (target_frame_width, target_frame_height))\n",
    "        # Put the image and the result side by side.\n",
    "        stacked_frame = np.hstack((image, result_frame))\n",
    "        # Save a frame to the video.\n",
    "        out_video.write(stacked_frame)\n",
    "\n",
    "        input_video_frame_nr = input_video_frame_nr + ADVANCE_FRAMES\n",
    "        cap.set(1, input_video_frame_nr)\n",
    "\n",
    "        progress_bar.progress = input_video_frame_nr\n",
    "        progress_bar.update()\n",
    "\n",
    "except KeyboardInterrupt:\n",
    "    print(\"Processing interrupted.\")\n",
    "finally:\n",
    "    clear_output()\n",
    "    processed_frames = num_frames // ADVANCE_FRAMES\n",
    "    out_video.release()\n",
    "    cap.release()\n",
    "    end_time = time.perf_counter()\n",
    "    duration = end_time - start_time\n",
    "\n",
    "    print(\n",
    "        f\"Processed {processed_frames} frames in {duration:.2f} seconds. \"\n",
    "        f\"Total FPS (including video processing): {processed_frames/duration:.2f}.\"\n",
    "        f\"Inference FPS: {processed_frames/total_inference_duration:.2f} \"\n",
    "    )\n",
    "    print(f\"Monodepth Video saved to '{str(result_video_path)}'.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe7e5c1d",
   "metadata": {
    "id": "bZ89ZI369KjA"
   },
   "source": [
    "### Display Monodepth Video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9233d3bd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "video = Video(result_video_path, width=800, embed=True)\n",
    "if not result_video_path.exists():\n",
    "    plt.imshow(stacked_frame)\n",
    "    raise ValueError(\"OpenCV was unable to write the video file. Showing one video frame.\")\n",
    "else:\n",
    "    print(f\"Showing monodepth video saved at\\n{result_video_path.resolve()}\")\n",
    "    print(\n",
    "        \"If you cannot see the video in your browser, please click on the \"\n",
    "        \"following link to download the video \"\n",
    "    )\n",
    "    video_link = FileLink(result_video_path)\n",
    "    video_link.html_link_str = \"<a href='%s' download>%s</a>\"\n",
    "    display(HTML(video_link._repr_html_()))\n",
    "    display(video)"
   ]
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}