natas_hw
/
openvino
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
openvino/docs/notebooks/222-vision-image-colorizati...

423 lines
19 KiB
ReStructuredText
Raw Blame History

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
%pip install "openvino-dev>=2023.1.0"
.. parsed-literal::
Requirement already satisfied: openvino-dev>=2023.1.0 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (2023.1.0)
Requirement already satisfied: addict>=2.4.0 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2023.1.0) (2.4.0)
Requirement already satisfied: defusedxml>=0.7.1 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2023.1.0) (0.7.1)
Requirement already satisfied: jstyleson>=0.0.2 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2023.1.0) (0.0.2)
Requirement already satisfied: networkx<=3.1 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2023.1.0) (2.8.2)
Requirement already satisfied: numpy>=1.16.6 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2023.1.0) (1.24.3)
Requirement already satisfied: opencv-python in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2023.1.0) (4.8.1.78)
Requirement already satisfied: openvino-telemetry>=2022.1.0 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2023.1.0) (2023.2.1)
Requirement already satisfied: pillow>=8.1.2 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2023.1.0) (10.0.1)
Requirement already satisfied: pyyaml>=5.4.1 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2023.1.0) (6.0.1)
Requirement already satisfied: requests>=2.25.1 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2023.1.0) (2.31.0)
Requirement already satisfied: texttable>=1.6.3 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2023.1.0) (1.7.0)
Requirement already satisfied: tqdm>=4.54.1 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2023.1.0) (4.66.1)
Requirement already satisfied: openvino==2023.1.0 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2023.1.0) (2023.1.0)
Requirement already satisfied: scipy<1.11,>=1.8 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from openvino-dev>=2023.1.0) (1.10.1)
Requirement already satisfied: charset-normalizer<4,>=2 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from requests>=2.25.1->openvino-dev>=2023.1.0) (3.3.1)
Requirement already satisfied: idna<4,>=2.5 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from requests>=2.25.1->openvino-dev>=2023.1.0) (3.4)
Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from requests>=2.25.1->openvino-dev>=2023.1.0) (2.0.7)
Requirement already satisfied: certifi>=2017.4.17 in /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from requests>=2.25.1->openvino-dev>=2023.1.0) (2023.7.22)
DEPRECATION: pytorch-lightning 1.6.5 has a non-standard dependency specifier torch>=1.8.*. pip 24.0 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
import sys
from pathlib import Path
import cv2
import matplotlib.pyplot as plt
import numpy as np
import openvino as ov
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.
.. 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-534/.workspace/scm/ov-notebook/.venv/bin/python -- /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.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
ONNX check passed successfully.
========== Converting colorization-v2 to IR (FP16)
Conversion command: /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/.venv/bin/python -- /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.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 ] 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/2023.0/openvino_2_0_transition_guide.html
[ SUCCESS ] Generated IR version 11 model.
[ SUCCESS ] XML file: /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/notebooks/222-vision-image-colorization/models/public/colorization-v2/FP16/colorization-v2.xml
[ SUCCESS ] BIN file: /opt/home/k8sworker/ci-ai/cibuilds/ov-notebook/OVNotebookOps-534/.workspace/scm/ov-notebook/notebooks/222-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:: 222-vision-image-colorization-with-output_files/222-vision-image-colorization-with-output_21_0.png
.. code:: ipython3
plot_output(test_img_1, color_img_1)
.. image:: 222-vision-image-colorization-with-output_files/222-vision-image-colorization-with-output_22_0.png
Reference in New Issue View Git Blame Copy Permalink