U^2-Net と OpenVINO™ を使用した画像の背景削除#

import platform 

%pip install -q "openvino>=2023.1.0" 
%pip install -q --extra-index-url https://download.pytorch.org/whl/cpu "torch>=2.1" opencv-python 
%pip install -q "gdown<4.6.4" 

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

import os 
import time 
from collections import namedtuple 
from pathlib import Path 

import cv2 
import matplotlib.pyplot as plt 
import numpy as np 
import openvino as ov 
import torch 
from IPython.display import HTML, FileLink, display

# ローカルモジュールをインポート 
import requests 

if not Path("./notebook_utils.py").exists():     # Fetch `notebook_utils` module 

    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 notebook_utils import load_image, download_file 

if not Path("./model/u2net.py").exists(): 
    download_file( 

url="https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/notebooks/vision-background-removal/model/u2net.py", directory="model" 
    ) 
from model.u2net import U2NET, U2NETP

model_config = namedtuple("ModelConfig", ["name", "url", "model", "model_args"]) 

u2net_lite = model_config( 
    name="u2net_lite", 
    url="https://drive.google.com/uc?id=1W8E4FHIlTVstfRkYmNOjbr0VDXTZm0jD", 
    model=U2NETP, 
    model_args=(), 
) 
u2net = model_config( 
    name="u2net", 
    url="https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ", 
    model=U2NET, 
    model_args=(3, 1), 
) 
u2net_human_seg = model_config( 
    name="u2net_human_seg", 
    url="https://drive.google.com/uc?id=1m_Kgs91b21gayc2XLW0ou8yugAIadWVP", 
    model=U2NET, 
    model_args=(3, 1), 
) 

# u2net_model を上記の 3 つの構成のいずれかに設定 
u2net_model = u2net_lite

# ダウンロードおよび変換されたモデルのファイル名
MODEL_DIR = "model" 
model_path = Path(MODEL_DIR) / u2net_model.name / Path(u2net_model.name).with_suffix(".pth")

if not model_path.exists(): 
    import gdown 

    os.makedirs(name=model_path.parent, exist_ok=True) 
    print("Start downloading model weights file...") 
    with open(model_path, "wb") as model_file: 
        gdown.download(url=u2net_model.url, output=model_file) 
        print(f"Model weights have been downloaded to {model_path}")

# モデルをロード 
net = u2net_model.model(*u2net_model.model_args) 
net.eval() 

# 重みをロード 
print(f"Loading model weights from: '{model_path}'") 
net.load_state_dict(state_dict=torch.load(model_path, map_location="cpu"))

BACKGROUND_FILE = "https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/image/wall.jpg" 
OUTPUT_DIR = "output" 

os.makedirs(name=OUTPUT_DIR, exist_ok=True) 

background_image = cv2.cvtColor(src=load_image(BACKGROUND_FILE), code=cv2.COLOR_BGR2RGB) 
background_image = cv2.resize(src=background_image, dsize=(image.shape[1], image.shape[0])) 

# 結果からの背景画像のすべての前景ピクセルを 0 に設定し、 
# 背景を削除した画像を追加します。 
background_image[resized_result == 1] = 0 
new_image = background_image + bg_removed_result 

# Save the generated image. 
new_image_path = Path(f"{OUTPUT_DIR}/{Path(IMAGE_URI).stem}-{Path(BACKGROUND_FILE).stem}.jpg") 
cv2.imwrite(filename=str(new_image_path), img=cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR)) 

# 元の画像と新しい背景の画像を並べて表示 
fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(18, 7)) 
ax[0].imshow(image) 
ax[1].imshow(new_image) 
for a in ax: 
    a.axis("off") 
plt.show() 

# 画像をダウンロードするリンクを作成 
image_link = FileLink(new_image_path) 
image_link.html_link_str = "<a href='%s' download>%s</a>" 
display( 
    HTML( 
        f"The generated image <code>
        {new_image_path.name}</code> is saved in " f"the directory <code>{new_image_path.parent}</code>.You can also " 
        "download the image by clicking on this link: " 
        f"{image_link._repr_html_()}" 
    ) 
)