notebooks/222-vision-image-colorization/222-vision-image-colorization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1662ee2c-fdcd-4b68-a65d-3b03a9e6aadf",
   "metadata": {},
   "source": [
    "### Image Colorization with OpenVINO\n",
    "\n",
    "This notebook demonstrates how to colorize images with OpenVINO using the Colorization model [colorization-v2](https://github.com/openvinotoolkit/open_model_zoo/blob/master/models/public/colorization-v2/README.md) or [colorization-siggraph](https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/public/colorization-siggraph) from [Open Model Zoo](https://github.com/openvinotoolkit/open_model_zoo/blob/master/models/public/index.md) based on the paper [Colorful Image Colorization](https://arxiv.org/abs/1603.08511) models from Open Model Zoo.\n",
    "\n",
    "![Let there be color](https://user-images.githubusercontent.com/18904157/180923280-9caefaf1-742b-4d2f-8943-5d4a6126e2fc.png)\n",
    "\n",
    "\n",
    "Given a grayscale image as input, the model generates colorized version of the image as the output.\n",
    "\n",
    "#### About Colorization-v2\n",
    "* The colorization-v2 model is one of the colorization group of models designed to perform image colorization.\n",
    "* Model trained on the ImageNet dataset.\n",
    "* Model consumes L-channel of LAB-image as input and produces predict A- and B-channels of LAB-image as output.\n",
    "\n",
    "#### About Colorization-siggraph\n",
    "* The colorization-siggraph model is one of the colorization group of models designed to real-time user-guided image colorization.\n",
    "* Model trained on the ImageNet dataset with synthetically generated user interaction.\n",
    "* Model consumes L-channel of LAB-image as input and produces predict A- and B-channels of LAB-image as output.\n",
    "\n",
    "See the [colorization](https://github.com/richzhang/colorization) repository for more details.\n",
    "#### Table of contents:\n",
    "- [Imports](#Imports)\n",
    "- [Configurations](#Configurations)\n",
    "    - [Select inference device](#Select-inference-device)\n",
    "- [Download the model](#Download-the-model)\n",
    "- [Convert the model to OpenVINO IR](#Convert-the-model-to-OpenVINO-IR)\n",
    "- [Loading the Model](#Loading-the-Model)\n",
    "- [Utility Functions](#Utility-Functions)\n",
    "- [Load the Image](#Load-the-Image)\n",
    "- [Display Colorized Image](#Display-Colorized-Image)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0aa40247-10ab-49ad-a1e7-aa8908a061b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install \"openvino-dev>=2023.1.0\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ae449c6-8a58-4d2c-8389-fde359ad7d1a",
   "metadata": {},
   "source": [
    "## Imports\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd439f78-a4b6-48b2-946e-a308383cf787",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import openvino as ov\n",
    "\n",
    "sys.path.append(\"../utils\")\n",
    "import notebook_utils as utils"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72d6ca64-aa95-4e50-b5a6-975259d16dcc",
   "metadata": {},
   "source": [
    "## Configurations\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "* `PRECISION` - {FP16, FP32}, default: FP16.\n",
    "* `MODEL_DIR` - directory where the model is to be stored, default: public.\n",
    "* `MODEL_NAME` - name of the model used for inference, default: colorization-v2.\n",
    "* `DATA_DIR` - directory where test images are stored, default: data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02dc1b06-b3ae-40f2-8fca-1368455a369c",
   "metadata": {},
   "outputs": [],
   "source": [
    "PRECISION = \"FP16\"\n",
    "MODEL_DIR = \"models\"\n",
    "MODEL_NAME = \"colorization-v2\"\n",
    "# MODEL_NAME=\"colorization-siggraph\"\n",
    "MODEL_PATH = f\"{MODEL_DIR}/public/{MODEL_NAME}/{PRECISION}/{MODEL_NAME}.xml\"\n",
    "DATA_DIR = \"data\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fab67524-0d85-462f-bfed-74851c53597b",
   "metadata": {},
   "source": [
    "### Select inference device\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "select device from dropdown list for running inference using OpenVINO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf984617-8dbd-440e-b09d-8ac69809b114",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "\n",
    "core = ov.Core()\n",
    "\n",
    "device = widgets.Dropdown(\n",
    "    options=core.available_devices + [\"AUTO\"],\n",
    "    value='AUTO',\n",
    "    description='Device:',\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cf87170-fd22-4fcf-b687-51b404be0db4",
   "metadata": {},
   "source": [
    "## Download the model\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "`omz_downloader` downloads model files from online sources and, if necessary, patches them to make them more usable with Model Converter.\n",
    "\n",
    "In this case, `omz_downloader` downloads the checkpoint and pytorch model of [colorization-v2](https://github.com/openvinotoolkit/open_model_zoo/blob/master/models/public/colorization-v2/README.md) or [colorization-siggraph](https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/public/colorization-siggraph) from [Open Model Zoo](https://github.com/openvinotoolkit/open_model_zoo/blob/master/models/public/index.md) and saves it under `MODEL_DIR`, as specified in the configuration above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "679a06d8-d0ec-440f-ae7a-afc4ebe841e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "download_command = (\n",
    "    f\"omz_downloader \"\n",
    "    f\"--name {MODEL_NAME} \"\n",
    "    f\"--output_dir {MODEL_DIR} \"\n",
    "    f\"--cache_dir {MODEL_DIR}\"\n",
    ")\n",
    "! $download_command"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93715d09-a6b9-4af5-b129-462c0cb9fc27",
   "metadata": {},
   "source": [
    "## Convert the model to OpenVINO IR\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "`omz_converter` converts the models that are not in the OpenVINO™ IR format into that format using model conversion API.\n",
    "\n",
    "The downloaded pytorch model is not in OpenVINO IR format which is required for inference with OpenVINO runtime. `omz_converter` is used to convert the downloaded pytorch model into ONNX and OpenVINO IR format respectively"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78cb56dd-0bd5-4a27-8588-91fb79a9b857",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(MODEL_PATH):\n",
    "    convert_command = (\n",
    "        f\"omz_converter \"\n",
    "        f\"--name {MODEL_NAME} \"\n",
    "        f\"--download_dir {MODEL_DIR} \"\n",
    "        f\"--precisions {PRECISION}\"\n",
    "    )\n",
    "    ! $convert_command"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d8f0b29-dd61-469e-86e8-d7fb956a636c",
   "metadata": {},
   "source": [
    "## Loading the Model\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Load the model in OpenVINO Runtime with `ie.read_model` and compile it for the specified device with `ie.compile_model`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8dc5285-1059-438b-80ef-7f2021b0e4ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "core = ov.Core()\n",
    "model = core.read_model(model=MODEL_PATH)\n",
    "compiled_model = core.compile_model(model=model, device_name=device.value)\n",
    "input_layer = compiled_model.input(0)\n",
    "output_layer = compiled_model.output(0)\n",
    "N, C, H, W = list(input_layer.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7681922b-5be1-4ec0-b48c-39af5e7bb44b",
   "metadata": {},
   "source": [
    "## Utility Functions\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a47519c1-a6dd-468b-845d-96386033e6b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_image(impath: str) -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Returns an image as ndarra, given path to an image reads the\n",
    "    (BGR) image using opencv's imread() API.\n",
    "\n",
    "        Parameter:\n",
    "            impath (string): Path of the image to be read and returned.\n",
    "\n",
    "        Returns:\n",
    "            image (ndarray): Numpy array representing the read image.\n",
    "    \"\"\"\n",
    "\n",
    "    raw_image = cv2.imread(impath)\n",
    "    if raw_image.shape[2] > 1:\n",
    "        image = cv2.cvtColor(\n",
    "            cv2.cvtColor(raw_image, cv2.COLOR_BGR2GRAY), cv2.COLOR_GRAY2RGB\n",
    "        )\n",
    "    else:\n",
    "        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)\n",
    "\n",
    "    return image\n",
    "\n",
    "\n",
    "def plot_image(image: np.ndarray, title: str = \"\") -> None:\n",
    "    \"\"\"\n",
    "    Given a image as ndarray and title as string, display it using\n",
    "    matplotlib.\n",
    "\n",
    "        Parameters:\n",
    "            image (ndarray): Numpy array representing the image to be\n",
    "                             displayed.\n",
    "            title (string): String representing the title of the plot.\n",
    "\n",
    "        Returns:\n",
    "            None\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    plt.imshow(image)\n",
    "    plt.title(title)\n",
    "    plt.axis(\"off\")\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def plot_output(gray_img: np.ndarray, color_img: np.ndarray) -> None:\n",
    "    \"\"\"\n",
    "    Plots the original (bw or grayscale) image and colorized image\n",
    "    on different column axes for comparing side by side.\n",
    "\n",
    "        Parameters:\n",
    "            gray_image (ndarray): Numpy array representing the original image.\n",
    "            color_image (ndarray): Numpy array representing the model output.\n",
    "\n",
    "        Returns:\n",
    "            None\n",
    "    \"\"\"\n",
    "\n",
    "    fig = plt.figure(figsize=(12, 12))\n",
    "\n",
    "    ax1 = fig.add_subplot(1, 2, 1)\n",
    "    plt.title(\"Input\", fontsize=20)\n",
    "    ax1.axis(\"off\")\n",
    "\n",
    "    ax2 = fig.add_subplot(1, 2, 2)\n",
    "    plt.title(\"Colorized\", fontsize=20)\n",
    "    ax2.axis(\"off\")\n",
    "\n",
    "    ax1.imshow(gray_img)\n",
    "    ax2.imshow(color_img)\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68d96b15-0176-45aa-9882-17e493acedff",
   "metadata": {},
   "source": [
    "## Load the Image\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1948813-77e6-434d-967a-5e0940cb40c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_url_0 = \"https://user-images.githubusercontent.com/18904157/180923287-20339d01-b1bf-493f-9a0d-55eff997aff1.jpg\"\n",
    "img_url_1 = \"https://user-images.githubusercontent.com/18904157/180923289-0bb71e09-25e1-46a6-aaf1-e8f666b62d26.jpg\"\n",
    "\n",
    "image_file_0 = utils.download_file(\n",
    "    img_url_0, filename=\"test_0.jpg\", directory=\"data\", show_progress=False, silent=True, timeout=30\n",
    ")\n",
    "assert Path(image_file_0).exists()\n",
    "\n",
    "image_file_1 = utils.download_file(\n",
    "    img_url_1, filename=\"test_1.jpg\", directory=\"data\", show_progress=False, silent=True, timeout=30\n",
    ")\n",
    "assert Path(image_file_1).exists()\n",
    "\n",
    "test_img_0 = read_image(\"data/test_0.jpg\")\n",
    "test_img_1 = read_image(\"data/test_1.jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afccb409-06a0-4e94-8f8a-779628425a3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def colorize(gray_img: np.ndarray) -> np.ndarray:\n",
    "\n",
    "    \"\"\"\n",
    "    Given an image as ndarray for inference convert the image into LAB image, \n",
    "    the model consumes as input L-Channel of LAB image and provides output \n",
    "    A & B - Channels of LAB image. i.e returns a colorized image\n",
    "\n",
    "        Parameters:\n",
    "            gray_img (ndarray): Numpy array representing the original\n",
    "                                image.\n",
    "\n",
    "        Returns:\n",
    "            colorize_image (ndarray): Numpy arrray depicting the\n",
    "                                      colorized version of the original\n",
    "                                      image.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Preprocess\n",
    "    h_in, w_in, _ = gray_img.shape\n",
    "    img_rgb = gray_img.astype(np.float32) / 255\n",
    "    img_lab = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2Lab)\n",
    "    img_l_rs = cv2.resize(img_lab.copy(), (W, H))[:, :, 0]\n",
    "\n",
    "    # Inference\n",
    "    inputs = np.expand_dims(img_l_rs, axis=[0, 1])\n",
    "    res = compiled_model([inputs])[output_layer]\n",
    "    update_res = np.squeeze(res)\n",
    "\n",
    "    # Post-process\n",
    "    out = update_res.transpose((1, 2, 0))\n",
    "    out = cv2.resize(out, (w_in, h_in))\n",
    "    img_lab_out = np.concatenate((img_lab[:, :, 0][:, :, np.newaxis],\n",
    "                                  out), axis=2)\n",
    "    img_bgr_out = np.clip(cv2.cvtColor(img_lab_out, cv2.COLOR_Lab2RGB), 0, 1)\n",
    "    colorized_image = (cv2.resize(img_bgr_out, (w_in, h_in))\n",
    "                       * 255).astype(np.uint8)\n",
    "    return colorized_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9db2da5e-9059-4220-bb58-a1283fa7aae4",
   "metadata": {},
   "outputs": [],
   "source": [
    "color_img_0 = colorize(test_img_0)\n",
    "color_img_1 = colorize(test_img_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68a950c4-7b94-4f84-82c7-3be7fddcc8e0",
   "metadata": {},
   "source": [
    "## Display Colorized Image\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9a7e75f-45a6-4ced-aa0d-29f707ddf55f",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_output(test_img_0, color_img_0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d39f7cc9-b872-4d9a-a144-36480cb88105",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_output(test_img_1, color_img_1)"
   ]
  }
 ],
 "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.11.5"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}