TFLite と OpenVINO を使用したセルフィーのセグメント化#

import platform 

%pip install -q "openvino>=2023.1.0" opencv-python tqdm 

if platform.system() != "Windows":
     %pip install -q "matplotlib>=3.4" 
else:
     %pip install -q "matplotlib>=3.4,<3.7"

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)

from pathlib import Path 
from notebook_utils import download_file 

tflite_model_path = Path("selfie_multiclass_256x256.tflite") 
tflite_model_url = "https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_multiclass_256x256/float32/latest/selfie_multiclass_256x256.tflite" 

download_file(tflite_model_url, tflite_model_path)

import ipywidgets as widgets 

device = widgets.Dropdown( 
    options=core.available_devices + ["AUTO"], 
    value="AUTO", 
    description="Device:", 
    disabled=False, 
) 

device

compiled_model = core.compile_model(ov_model, device.value)

import cv2 
import numpy as np 
from notebook_utils import load_image 

# 入力画像を読み取り RGB に変換 
test_image_url = "https://user-images.githubusercontent.com/29454499/251036317-551a2399-303e-4a4a-a7d6-d7ce973e05c5.png" 
img = load_image(test_image_url) 
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 

# 前処理ヘルパー関数 
def resize_and_pad(image: np.ndarray, height: int = 256, width: int = 256): 
    """ 
    Input preprocessing function, takes input image in np.ndarray format, 
    resizes it to fit specified height and width with preserving aspect ratio 
    and adds padding on bottom or right side to complete target height x width rectangle.

    Parameters: 
        image (np.ndarray): input image in np.ndarray format 
        height (int, *optional*, 256): target height 
        width (int, *optional*, 256): target width 
    Returns: 
        padded_img (np.ndarray): processed image 
        padding_info (Tuple[int, int]): information about padding size, required for postprocessing 
    """ 
    h, w = image.shape[:2] 
    if h < w: 
        img = cv2.resize(image, (width, np.floor(h / (w / width)).astype(int))) 
    else: 
        img = cv2.resize(image, (np.floor(w / (h / height)).astype(int), height)) 

    r_h, r_w = img.shape[:2] 
    right_padding = width - r_w 
    bottom_padding = height - r_h 
    padded_img = cv2.copyMakeBorder(img, 0, bottom_padding, 0, right_padding, cv2.BORDER_CONSTANT) 
    return padded_img, (bottom_padding, right_padding) 

# 前処理ステップを適用 - 入力画像のサイズ変更とパディング 
padded_img, pad_info = resize_and_pad(np.array(img)) 

# 入力データを uint8 [0, 255] から float32 [0, 1] の範囲に変換し、バッチ次元を追加 
normalized_img = np.expand_dims(padded_img.astype(np.float32) / 255, 0)

out = compiled_model(normalized_img)[0]

from typing import Tuple 
from notebook_utils import segmentation_map_to_image, SegmentationMap, Label 

# 視覚化のセグメント化ラベルのヘルパー 
labels = [ 
    Label(index=0, color=(192, 192, 192), name="background"), 
    Label(index=1, color=(128, 0, 0), name="hair"), 
    Label(index=2, color=(255, 229, 204), name="body skin"), 
    Label(index=3, color=(255, 204, 204), name="face skin"), 
    Label(index=4, color=(0, 0, 128), name="clothes"), 
    Label(index=5, color=(128, 0, 128), name="others"), 
] 
SegmentationLabels = SegmentationMap(labels) 

# 後処理出力マスクのヘルパー 
def postprocess_mask(out: np.ndarray, pad_info: Tuple[int, int], orig_img_size: Tuple[int, int]): 
    """ 
    Posptprocessing function for segmentation mask, accepts model output tensor, gets labels for each pixel using argmax, 
    unpads segmentation mask and resizes it to original image size.

    Parameters: 
        out (np.ndarray): model output tensor 
        pad_info (Tuple[int, int]): information about padding size from preprocessing step 
        orig_img_size (Tuple[int, int]): original image height and width for resizing 
    Returns: 
        label_mask_resized (np.ndarray): postprocessed segmentation label mask 
    """ 
    label_mask = np.argmax(out, -1)[0] 
    pad_h, pad_w = pad_info 
    unpad_h = label_mask.shape[0] - pad_h 
    unpad_w = label_mask.shape[1] - pad_w 
    label_mask_unpadded = label_mask[:unpad_h, :unpad_w] 
    orig_h, orig_w = orig_img_size 
    label_mask_resized = cv2.resize(label_mask_unpadded, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST) 
    return label_mask_resized 

# オリジナル画像に関する情報を取得 
image_data = np.array(img) 
orig_img_shape = image_data.shape 

# 置換する背景色を指定 
BG_COLOR = (192, 192, 192) 

# ガウスぼかしをによる背景ぼかしシナリオの画像をぼかす 
blurred_image = cv2.GaussianBlur(image_data, (55, 55), 0) 

# 後処理出力 
postprocessed_mask = postprocess_mask(out, pad_info, orig_img_shape[:2]) 

# 色分けされたセグメント化マップを取得 
output_mask = segmentation_map_to_image(postprocessed_mask, 
SegmentationLabels.get_colormap()) 

# 元の画像の背景を置き換え、 
# 単色の背景色で画像を塗りつぶす 
bg_image = np.full(orig_img_shape, BG_COLOR, dtype=np.uint8) 

# 背景と前景を分離するための条件マスクを定義 
condition = np.stack((postprocessed_mask,) * 3, axis=-1) > 0 
# 背景を単色に置き換え 
output_image = np.where(condition, image_data, bg_image) 
# 背景をぼかした画像のコピーに置き換え 
output_blurred_image = np.where(condition, image_data, blurred_image)

import matplotlib.pyplot as plt 

titles = ["Original image", "Portrait mask", "Removed background", "Blurred background"] 
images = [image_data, output_mask, output_image, output_blurred_image] 
figsize = (16, 16) 
fig, axs = plt.subplots(2, 2, figsize=figsize, sharex="all", sharey="all") 
fig.patch.set_facecolor("white") 
list_axes = list(axs.flat) 
for i, a in enumerate(list_axes): 
    a.set_xticklabels([]) 
    a.set_yticklabels([]) 
    a.get_xaxis().set_visible(False) 
    a.get_yaxis().set_visible(False) 
    a.grid(False) 
    a.imshow(images[i].astype(np.uint8)) 
    a.set_title(titles[i]) 
fig.subplots_adjust(wspace=0.0, hspace=-0.8) 
fig.tight_layout()

WEBCAM_INFERENCE = False 

if WEBCAM_INFERENCE:
     VIDEO_SOURCE = 0 # ウェブカメラ 
else:
     VIDEO_SOURCE = 
"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/CEO%20Pat%20Gelsinger%20on%20Leading%20Intel.mp4"

device

run_background_blurring(source=VIDEO_SOURCE, device=device.value)