461 lines
19 KiB
ReStructuredText
461 lines
19 KiB
ReStructuredText
Image Colorization with OpenVINO
|
|
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
|
|
|
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.
|
|
|
|
.. figure:: https://user-images.githubusercontent.com/18904157/180923280-9caefaf1-742b-4d2f-8943-5d4a6126e2fc.png
|
|
:alt: 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 <https://github.com/richzhang/colorization>`__
|
|
repository for more details.
|
|
|
|
Table of contents:
|
|
^^^^^^^^^^^^^^^^^^
|
|
|
|
- `Imports <#imports>`__
|
|
- `Configurations <#configurations>`__
|
|
|
|
- `Select inference device <#select-inference-device>`__
|
|
|
|
- `Download the model <#download-the-model>`__
|
|
- `Convert the model to OpenVINO
|
|
IR <#convert-the-model-to-openvino-ir>`__
|
|
- `Loading the Model <#loading-the-model>`__
|
|
- `Utility Functions <#utility-functions>`__
|
|
- `Load the Image <#load-the-image>`__
|
|
- `Display Colorized Image <#display-colorized-image>`__
|
|
|
|
.. code:: ipython3
|
|
|
|
import platform
|
|
|
|
%pip install "openvino-dev>=2024.0.0" opencv-python tqdm
|
|
|
|
if platform.system() != "Windows":
|
|
%pip install -q "matplotlib>=3.4"
|
|
else:
|
|
%pip install -q "matplotlib>=3.4,<3.7"
|
|
|
|
|
|
.. parsed-literal::
|
|
|
|
Requirement already satisfied: openvino-dev>=2024.0.0 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (2024.1.0)
|
|
Requirement already satisfied: opencv-python in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (4.10.0.82)
|
|
Requirement already satisfied: tqdm in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (4.66.4)
|
|
Requirement already satisfied: defusedxml>=0.7.1 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2024.0.0) (0.7.1)
|
|
Requirement already satisfied: networkx<=3.1.0 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2024.0.0) (3.1)
|
|
Requirement already satisfied: numpy<2.0.0,>=1.16.6 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2024.0.0) (1.23.5)
|
|
Requirement already satisfied: openvino-telemetry>=2023.2.1 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2024.0.0) (2024.1.0)
|
|
Requirement already satisfied: packaging in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2024.0.0) (24.0)
|
|
Requirement already satisfied: pyyaml>=5.4.1 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2024.0.0) (6.0.1)
|
|
Requirement already satisfied: requests>=2.25.1 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2024.0.0) (2.32.3)
|
|
Collecting openvino==2024.1.0 (from openvino-dev>=2024.0.0)
|
|
Using cached openvino-2024.1.0-15008-cp38-cp38-manylinux2014_x86_64.whl.metadata (8.8 kB)
|
|
Requirement already satisfied: charset-normalizer<4,>=2 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from requests>=2.25.1->openvino-dev>=2024.0.0) (3.3.2)
|
|
Requirement already satisfied: idna<4,>=2.5 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from requests>=2.25.1->openvino-dev>=2024.0.0) (3.7)
|
|
Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from requests>=2.25.1->openvino-dev>=2024.0.0) (2.2.1)
|
|
Requirement already satisfied: certifi>=2017.4.17 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from requests>=2.25.1->openvino-dev>=2024.0.0) (2024.6.2)
|
|
Using cached openvino-2024.1.0-15008-cp38-cp38-manylinux2014_x86_64.whl (38.7 MB)
|
|
DEPRECATION: pytorch-lightning 1.6.5 has a non-standard dependency specifier torch>=1.8.*. pip 24.1 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063
|
|
Installing collected packages: openvino
|
|
Attempting uninstall: openvino
|
|
Found existing installation: openvino 2024.3.0.dev20240605
|
|
Uninstalling openvino-2024.3.0.dev20240605:
|
|
Successfully uninstalled openvino-2024.3.0.dev20240605
|
|
ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
|
|
openvino-tokenizers 2024.3.0.0.dev20240605 requires openvino~=2024.3.0.0.dev, but you have openvino 2024.1.0 which is incompatible.
|
|
Successfully installed openvino-2024.1.0
|
|
Note: you may need to restart the kernel to use updated packages.
|
|
DEPRECATION: pytorch-lightning 1.6.5 has a non-standard dependency specifier torch>=1.8.*. pip 24.1 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063
|
|
Note: you may need to restart the kernel to use updated packages.
|
|
|
|
|
|
Imports
|
|
-------
|
|
|
|
|
|
|
|
.. code:: ipython3
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import openvino as ov
|
|
|
|
# Fetch `notebook_utils` module
|
|
import requests
|
|
|
|
r = requests.get(
|
|
url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py",
|
|
)
|
|
|
|
open("notebook_utils.py", "w").write(r.text)
|
|
|
|
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.
|
|
|
|
.. code:: ipython3
|
|
|
|
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"
|
|
|
|
Select inference device
|
|
~~~~~~~~~~~~~~~~~~~~~~~
|
|
|
|
|
|
|
|
select device from dropdown list for running inference using OpenVINO
|
|
|
|
.. code:: ipython3
|
|
|
|
import ipywidgets as widgets
|
|
|
|
core = ov.Core()
|
|
|
|
device = widgets.Dropdown(
|
|
options=core.available_devices + ["AUTO"],
|
|
value="AUTO",
|
|
description="Device:",
|
|
disabled=False,
|
|
)
|
|
|
|
device
|
|
|
|
|
|
|
|
|
|
.. parsed-literal::
|
|
|
|
Dropdown(description='Device:', index=1, options=('CPU', 'AUTO'), value='AUTO')
|
|
|
|
|
|
|
|
Download the model
|
|
------------------
|
|
|
|
|
|
|
|
``omz_downloader`` downloads model files from online sources and, if
|
|
necessary, patches them to make them more usable with Model Converter.
|
|
|
|
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.
|
|
|
|
.. code:: ipython3
|
|
|
|
download_command = f"omz_downloader " f"--name {MODEL_NAME} " f"--output_dir {MODEL_DIR} " f"--cache_dir {MODEL_DIR}"
|
|
! $download_command
|
|
|
|
|
|
.. parsed-literal::
|
|
|
|
################|| 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 conversion API.
|
|
|
|
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
|
|
|
|
.. code:: ipython3
|
|
|
|
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
|
|
|
|
|
|
.. parsed-literal::
|
|
|
|
========== Converting colorization-v2 to ONNX
|
|
Conversion to ONNX command: /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/bin/python -- /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages/omz_tools/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
|
|
|
|
ONNX check passed successfully.
|
|
|
|
========== Converting colorization-v2 to IR (FP16)
|
|
Conversion command: /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/bin/python -- /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/.venv/bin/mo --framework=onnx --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]' --compress_to_fp16=True
|
|
|
|
[ INFO ] Generated IR will be compressed to FP16. If you get lower accuracy, please consider disabling compression explicitly by adding argument --compress_to_fp16=False.
|
|
Find more information about compression to FP16 at https://docs.openvino.ai/2023.0/openvino_docs_MO_DG_FP16_Compression.html
|
|
[ INFO ] MO command line tool is considered as the legacy conversion API as of OpenVINO 2023.2 release. Please use OpenVINO Model Converter (OVC). OVC represents a lightweight alternative of MO and provides simplified model conversion API.
|
|
Find more information about transition from MO to OVC at https://docs.openvino.ai/2023.2/openvino_docs_OV_Converter_UG_prepare_model_convert_model_MO_OVC_transition.html
|
|
[ SUCCESS ] Generated IR version 11 model.
|
|
[ SUCCESS ] XML file: /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/notebooks/vision-image-colorization/models/public/colorization-v2/FP16/colorization-v2.xml
|
|
[ SUCCESS ] BIN file: /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-697/.workspace/scm/ov-notebook/notebooks/vision-image-colorization/models/public/colorization-v2/FP16/colorization-v2.bin
|
|
|
|
|
|
|
|
Loading the Model
|
|
-----------------
|
|
|
|
|
|
|
|
Load the model in OpenVINO Runtime with ``ie.read_model`` and compile it
|
|
for the specified device with ``ie.compile_model``.
|
|
|
|
.. code:: ipython3
|
|
|
|
core = ov.Core()
|
|
model = core.read_model(model=MODEL_PATH)
|
|
compiled_model = core.compile_model(model=model, device_name=device.value)
|
|
input_layer = compiled_model.input(0)
|
|
output_layer = compiled_model.output(0)
|
|
N, C, H, W = list(input_layer.shape)
|
|
|
|
Utility Functions
|
|
-----------------
|
|
|
|
|
|
|
|
.. code:: ipython3
|
|
|
|
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
|
|
--------------
|
|
|
|
|
|
|
|
.. code:: ipython3
|
|
|
|
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")
|
|
|
|
.. code:: ipython3
|
|
|
|
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
|
|
|
|
.. code:: ipython3
|
|
|
|
color_img_0 = colorize(test_img_0)
|
|
color_img_1 = colorize(test_img_1)
|
|
|
|
Display Colorized Image
|
|
-----------------------
|
|
|
|
|
|
|
|
.. code:: ipython3
|
|
|
|
plot_output(test_img_0, color_img_0)
|
|
|
|
|
|
|
|
.. image:: vision-image-colorization-with-output_files/vision-image-colorization-with-output_21_0.png
|
|
|
|
|
|
.. code:: ipython3
|
|
|
|
plot_output(test_img_1, color_img_1)
|
|
|
|
|
|
|
|
.. image:: vision-image-colorization-with-output_files/vision-image-colorization-with-output_22_0.png
|
|
|