Image Colorization with OpenVINO

This tutorial is also available as a Jupyter notebook that can be cloned directly from GitHub. See the installation guide for instructions to run this tutorial locally on Windows, Linux or macOS. To run without installing anything, click the launch binder button.

Binder Github

This notebook demonstrates how to colorize images with OpenVINO using the Colorization model colorization-v2 or colorization-siggraph from Open Model Zoo based on the paper Colorful Image Colorization models from Open Model Zoo.

Let there be color

Let there be color

Given a grayscale image as input, the model generates colorized version of the image as the output.

About Colorization-v2

  • The colorization-v2 model is one of the colorization group of models designed to perform image colorization.

  • Model trained on the ImageNet dataset.

  • Model consumes L-channel of LAB-image as input and produces predict A- and B-channels of LAB-image as output.

About Colorization-siggraph

  • The colorization-siggraph model is one of the colorization group of models designed to real-time user-guided image colorization.

  • Model trained on the ImageNet dataset with synthetically generated user interaction.

  • Model consumes L-channel of LAB-image as input and produces predict A- and B-channels of LAB-image as output.

See the colorization repository for more details.

Imports

import os
import sys
from pathlib import Path

import cv2
import matplotlib.pyplot as plt
import numpy as np
from openvino.runtime import Core

sys.path.append("../utils")
import notebook_utils as utils

Configurations

  • PRECISION - {FP16, FP32}, default: FP16

  • MODEL_DIR - directory where the model is to be stored, default: public.

  • MODEL_NAME - name of the model used for inference, default: colorization-v2

  • DATA_DIR - directory where test images are stored, default: data

  • DEVICE - {CPU, GPU, GNA,VPU} device to used for inference, default: CPU

PRECISION = "FP16"
MODEL_DIR = "models"
MODEL_NAME = "colorization-v2"
# MODEL_NAME="colorization-siggraph"
MODEL_PATH = f"{MODEL_DIR}/public/{MODEL_NAME}/{PRECISION}/{MODEL_NAME}.xml"
DATA_DIR = "data"
DEVICE = "CPU"

Download the model

omz_downloader downloads model files from online sources and, if necessary, patches them to make them more usable with Model Convertor.

In our case omz_downloader downloads the checkpoint and pytorch model of colorization-v2 or colorization-siggraph from Open Model Zoo and saves it under MODEL_DIR as specified in above configuration

download_command = (
    f"omz_downloader "
    f"--name {MODEL_NAME} "
    f"--output_dir {MODEL_DIR} "
    f"--cache_dir {MODEL_DIR}"
)
! $download_command

################|| Downloading colorization-v2 ||################

========== Downloading models/public/colorization-v2/ckpt/colorization-v2-eccv16.pth


========== Downloading models/public/colorization-v2/model/__init__.py


========== Downloading models/public/colorization-v2/model/base_color.py


========== Downloading models/public/colorization-v2/model/eccv16.py


========== Replacing text in models/public/colorization-v2/model/__init__.py
========== Replacing text in models/public/colorization-v2/model/__init__.py
========== Replacing text in models/public/colorization-v2/model/eccv16.py

Convert the model to OpenVINO IR

omz_converter converts the models that are not in the OpenVINO™ IR format into that format using Model Optimizer.

Our downloaded pytorch model is not in OpenVINO IR format which is required for inference with OpenVINO runtime, omz_converter is used to convert our downloaded pytorch model into ONNX and OpenVINO IR format respectively

if not os.path.exists(MODEL_PATH):
    convert_command = (
        f"omz_converter "
        f"--name {MODEL_NAME} "
        f"--download_dir {MODEL_DIR} "
        f"--precisions {PRECISION}"
    )
    ! $convert_command

========== Converting colorization-v2 to ONNX
Conversion to ONNX command: /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-275/.workspace/scm/ov-notebook/.venv/bin/python -- /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-275/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages/openvino/model_zoo/internal_scripts/pytorch_to_onnx.py --model-path=models/public/colorization-v2 --model-name=ECCVGenerator --weights=models/public/colorization-v2/ckpt/colorization-v2-eccv16.pth --import-module=model --input-shape=1,1,256,256 --output-file=models/public/colorization-v2/colorization-v2-eccv16.onnx --input-names=data_l --output-names=color_ab

/opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-275/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages/torch/nn/functional.py:3454: UserWarning: Default upsampling behavior when mode=bilinear is changed to align_corners=False since 0.4.0. Please specify align_corners=True if the old behavior is desired. See the documentation of nn.Upsample for details.
  warnings.warn(
ONNX check passed successfully.

========== Converting colorization-v2 to IR (FP16)
Conversion command: /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-275/.workspace/scm/ov-notebook/.venv/bin/python -- /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-275/.workspace/scm/ov-notebook/.venv/bin/mo --framework=onnx --data_type=FP16 --output_dir=models/public/colorization-v2/FP16 --model_name=colorization-v2 --input=data_l --output=color_ab --input_model=models/public/colorization-v2/colorization-v2-eccv16.onnx '--layout=data_l(NCHW)' '--input_shape=[1, 1, 256, 256]'

Model Optimizer arguments:
Common parameters:
    - Path to the Input Model:  /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-275/.workspace/scm/ov-notebook/notebooks/222-vision-image-colorization/models/public/colorization-v2/colorization-v2-eccv16.onnx
    - Path for generated IR:    /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-275/.workspace/scm/ov-notebook/notebooks/222-vision-image-colorization/models/public/colorization-v2/FP16
    - IR output name:   colorization-v2
    - Log level:    ERROR
    - Batch:    Not specified, inherited from the model
    - Input layers:     data_l
    - Output layers:    color_ab
    - Input shapes:     [1, 1, 256, 256]
    - Source layout:    Not specified
    - Target layout:    Not specified
    - Layout:   data_l(NCHW)
    - Mean values:  Not specified
    - Scale values:     Not specified
    - Scale factor:     Not specified
    - Precision of IR:  FP16
    - Enable fusing:    True
    - User transformations:     Not specified
    - Reverse input channels:   False
    - Enable IR generation for fixed input shape:   False
    - Use the transformations config file:  None
Advanced parameters:
    - Force the usage of legacy Frontend of Model Optimizer for model conversion into IR:   False
    - Force the usage of new Frontend of Model Optimizer for model conversion into IR:  False
OpenVINO runtime found in:  /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-275/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages/openvino
OpenVINO runtime version:   2022.2.0-7713-af16ea1d79a-releases/2022/2
Model Optimizer version:    2022.2.0-7713-af16ea1d79a-releases/2022/2
[ SUCCESS ] Generated IR version 11 model.
[ SUCCESS ] XML file: /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-275/.workspace/scm/ov-notebook/notebooks/222-vision-image-colorization/models/public/colorization-v2/FP16/colorization-v2.xml
[ SUCCESS ] BIN file: /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-275/.workspace/scm/ov-notebook/notebooks/222-vision-image-colorization/models/public/colorization-v2/FP16/colorization-v2.bin
[ SUCCESS ] Total execution time: 0.81 seconds.
[ SUCCESS ] Memory consumed: 313 MB.
[ INFO ] The model was converted to IR v11, the latest model format that corresponds to the source DL framework input/output format. While IR v11 is backwards compatible with OpenVINO Inference Engine API v1.0, please use API v2.0 (as of 2022.1) to take advantage of the latest improvements in IR v11.
Find more information about API v2.0 and IR v11 at https://docs.openvino.ai

Loading the Model

Load the model in OpenVINO Runtime with ie.read_model and compile it for the specified device with ie.compile_model

ie = Core()
model = ie.read_model(model=MODEL_PATH)
compiled_model = ie.compile_model(model=model, device_name=DEVICE)
input_layer = compiled_model.input(0)
output_layer = compiled_model.output(0)
N, C, H, W = list(input_layer.shape)

Utility Functions

def read_image(impath: str) -> np.ndarray:
    """
    Returns an image as ndarra, given path to an image reads the
    (BGR) image using opencv's imread() API.

        Parameter:
            impath (string): Path of the image to be read and returned.

        Returns:
            image (ndarray): Numpy array representing the read image.
    """

    raw_image = cv2.imread(impath)
    if raw_image.shape[2] > 1:
        image = cv2.cvtColor(
            cv2.cvtColor(raw_image, cv2.COLOR_BGR2GRAY), cv2.COLOR_GRAY2RGB
        )
    else:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)

    return image


def plot_image(image: np.ndarray, title: str = "") -> None:
    """
    Given a image as ndarray and title as string, display it using
    matplotlib.

        Parameters:
            image (ndarray): Numpy array representing the image to be
                             displayed.
            title (string): String representing the title of the plot.

        Returns:
            None

    """

    plt.imshow(image)
    plt.title(title)
    plt.axis("off")
    plt.show()


def plot_output(gray_img: np.ndarray, color_img: np.ndarray) -> None:
    """
    Plots the original (bw or grayscale) image and colorized image
    on different column axes for comparing side by side.

        Parameters:
            gray_image (ndarray): Numpy array representing the original image.
            color_image (ndarray): Numpy array representing the model output.

        Returns:
            None
    """

    fig = plt.figure(figsize=(12, 12))

    ax1 = fig.add_subplot(1, 2, 1)
    plt.title("Input", fontsize=20)
    ax1.axis("off")

    ax2 = fig.add_subplot(1, 2, 2)
    plt.title("Colorized", fontsize=20)
    ax2.axis("off")

    ax1.imshow(gray_img)
    ax2.imshow(color_img)

    plt.show()

Load the Image

img_url_0 = "https://user-images.githubusercontent.com/18904157/180923287-20339d01-b1bf-493f-9a0d-55eff997aff1.jpg"
img_url_1 = "https://user-images.githubusercontent.com/18904157/180923289-0bb71e09-25e1-46a6-aaf1-e8f666b62d26.jpg"

image_file_0 = utils.download_file(
    img_url_0, filename="test_0.jpg", directory="data", show_progress=False, silent=True, timeout=30
)
assert Path(image_file_0).exists()

image_file_1 = utils.download_file(
    img_url_1, filename="test_1.jpg", directory="data", show_progress=False, silent=True, timeout=30
)
assert Path(image_file_1).exists()

test_img_0 = read_image("data/test_0.jpg")
test_img_1 = read_image("data/test_1.jpg")

def colorize(gray_img: np.ndarray) -> np.ndarray:

    """
    Given an image as ndarray for inference convert the image into LAB image,
    the model consumes as input L-Channel of LAB image and provides output
    A & B - Channels of LAB image. i.e returns a colorized image

        Parameters:
            gray_img (ndarray): Numpy array representing the original
                                image.

        Returns:
            colorize_image (ndarray): Numpy arrray depicting the
                                      colorized version of the original
                                      image.
    """

    # Preprocess
    h_in, w_in, _ = gray_img.shape
    img_rgb = gray_img.astype(np.float32) / 255
    img_lab = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2Lab)
    img_l_rs = cv2.resize(img_lab.copy(), (W, H))[:, :, 0]

    # Inference
    inputs = np.expand_dims(img_l_rs, axis=[0, 1])
    res = compiled_model([inputs])[output_layer]
    update_res = np.squeeze(res)

    # Post-process
    out = update_res.transpose((1, 2, 0))
    out = cv2.resize(out, (w_in, h_in))
    img_lab_out = np.concatenate((img_lab[:, :, 0][:, :, np.newaxis],
                                  out), axis=2)
    img_bgr_out = np.clip(cv2.cvtColor(img_lab_out, cv2.COLOR_Lab2RGB), 0, 1)
    colorized_image = (cv2.resize(img_bgr_out, (w_in, h_in))
                       * 255).astype(np.uint8)
    return colorized_image

color_img_0 = colorize(test_img_0)
color_img_1 = colorize(test_img_1)

Display Colorized Image

plot_output(test_img_0, color_img_0)
../_images/222-vision-image-colorization-with-output_18_0.png 
plot_output(test_img_1, color_img_1)
../_images/222-vision-image-colorization-with-output_19_0.png
Prev Next