LRASPP MobileNet v3 と OpenVINO によるセマンティクスのセグメント化

この Jupyter ノートブックはオンラインで起動でき、ブラウザーのウィンドウで対話型環境を開きます。ローカルにインストールすることもできます。次のオプションのいずれかを選択します。

Binder Google Colab GitHub

torchvision.models サブパッケージには、画像分類、ピクセル単位のセマンティック・セグメント化、オブジェクト検出、インスタンスのセグメント化、人物キーポイント検出、ビデオ分類、オプティカル・フローなど、さまざまなタスクに対処するモデルの定義が含まれています。このノートブックでは、そのいずれかの使用方法を説明します。LRASPP モデルは、MobileNetV3 の検索に関する論文に基づいています。この論文によると、MobileNetV3、LR-ASPP、または Lite Reduced Atrous Spatial Pyramid Pooling の検索には、軽量で効率的なセグメント化デコーダーのアーキテクチャーが採用されています。 このモデルは MS COCO データセットで事前トレーニングされています。80 のクラスすべてでトレーニングする代わりに、セグメント化モデルは PASCAL VOC データセットの 20 クラスでトレーニングされました: 背景、飛行機、自転車、鳥、ボート、ボトル、バス、車、猫、椅子、牛、ダイニングテーブル、犬、馬、バイク、人、鉢植え、羊、ソファ、電車、テレビモニター

モデルの詳細については、torchvision のドキュメントを参照してください。

目次

必要条件

%pip install -q --extra-index-url https://download.pytorch.org/whl/cpu torch torchvision
%pip install -q matplotlib
%pip install -q "openvino>=2023.2.0"
Note: you may need to restart the kernel to use updated packages.
Note: you may need to restart the kernel to use updated packages.
Note: you may need to restart the kernel to use updated packages.
from pathlib import Path

import openvino as ov
import torch

テスト画像を取得

まず、開いたデータセットからテスト画像を取得します。

import urllib.request

from torchvision.io import read_image
import torchvision.transforms as transforms


img_path = 'cats_image.jpeg'
urllib.request.urlretrieve(
    url='https://huggingface.co/datasets/huggingface/cats-image/resolve/main/cats_image.jpeg',
    filename=img_path
)
image = read_image(img_path)
display(transforms.ToPILImage()(image))
../_images/125-lraspp-segmentation-with-output_5_0.png

モデルをダウンロードして準備

推論中にネットワークによって使用される画像の幅と高さを定義します。入力変換関数に従って、モデルは高さ 480、幅 640 の画像で事前トレーニングされます。

IMAGE_WIDTH = 640
IMAGE_HEIGHT = 480

Torchvision は、利用可能なモデルをリストし、取得するメカニズムを提供します。

import torchvision.models as models

# List available models
all_models = models.list_models()
# List of models by type
segmentation_models = models.list_models(module=models.segmentation)

print(segmentation_models)
['deeplabv3_mobilenet_v3_large', 'deeplabv3_resnet101', 'deeplabv3_resnet50', 'fcn_resnet101', 'fcn_resnet50', 'lraspp_mobilenet_v3_large']

lraspp_mobilenet_v3_large を使用します。models.get_model("lraspp_mobilenet_v3_large", weights='DEFAULT') を使用して名前でモデルを取得するか、対応する関数を直接呼び出すことができます。ここでは、torchvision.models.segmentation.lraspp_mobilenet_v3_large を使用します。重み列挙型 LRASPP_MobileNet_V3_Large_Weights.COCO_WITH_VOC_LABELS_V1 を使用して、事前トレーニングされたモデルの重みをモデル初期化関数に直接渡すことができます。これはデフォルトの重みです。モデルで使用可能なすべての重みを取得するため、weights_enum = models.get_model_weights("lraspp_mobilenet_v3_large") を呼び出すことができますが、このモデルには 1 つしかありません。

weights = models.segmentation.LRASPP_MobileNet_V3_Large_Weights.COCO_WITH_VOC_LABELS_V1
model = models.segmentation.lraspp_mobilenet_v3_large(weights=weights)

前処理を定義して入力データを準備

torchvision.transforms で前処理を行ったり、モデル の重みから前処理変換したりできます。

import numpy as np


preprocess = models.segmentation.LRASPP_MobileNet_V3_Large_Weights.COCO_WITH_VOC_LABELS_V1.transforms()
preprocess.resize_size = (IMAGE_HEIGHT, IMAGE_WIDTH)  # change to an image size

input_data = preprocess(image)
input_data = np.expand_dims(input_data, axis=0)

PyTorch モデルで推論を実行

model.eval()
with torch.no_grad():
    result_torch = model(torch.as_tensor(input_data).float())['out']

元のモデルをOpenVINO IR形式に変換

元のモデルを FP16 精度の OpenVINO IR に変換するには、モデル変換 API を使用します。モデルは現在のディレクトリー内に保存されます。モデルの変換方法の詳細については、このページを参照してください。

ov_model_xml_path = Path('models/ov_lraspp_model.xml')


if not ov_model_xml_path.exists():
    ov_model_xml_path.parent.mkdir(parents=True, exist_ok=True)
    dummy_input = torch.randn(1, 3, IMAGE_HEIGHT, IMAGE_WIDTH)
    ov_model = ov.convert_model(model, example_input=dummy_input)
    ov.save_model(ov_model, ov_model_xml_path)
else:
    print(f"IR model {ov_model_xml_path} already exists.")

OpenVINO モデルで推論を実行

OpenVINO を使用して推論を実行するデバイスをドロップダウン・リストから選択します。

import ipywidgets as widgets

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

device
Dropdown(description='Device:', index=1, options=('CPU', 'AUTO'), value='AUTO')

推論の実行

compiled_model = core.compile_model(ov_model_xml_path, device_name=device.value)
res_ir = compiled_model(input_data)[0]

結果の表示

OpenVINO IR および PyTorch モデルの予測を比較して、セグメント化の結果が期待どおりであることを確認します。

pytorch チュートリアルを使用してセグメント化マスクを視覚化できます。以下は、PyTorch モデルの cat マスクを使用して画像を視覚化する簡単な例です。

import torch
import matplotlib.pyplot as plt

import torchvision.transforms.functional as F


plt.rcParams["savefig.bbox"] = 'tight'


def show(imgs):
    if not isinstance(imgs, list):
        imgs = [imgs]
    fix, axs = plt.subplots(ncols=len(imgs), squeeze=False)
    for i, img in enumerate(imgs):
        img = img.detach()
        img = F.to_pil_image(img)
        axs[0, i].imshow(np.asarray(img))
        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])

cat マスクを用意して表示します。

sem_classes = [
    '__background__', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus',
    'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike',
    'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor'
]
sem_class_to_idx = {cls: idx for (idx, cls) in enumerate(sem_classes)}

normalized_mask = torch.nn.functional.softmax(result_torch, dim=1)

cat_mask = normalized_mask[0, sem_class_to_idx['cat']]

show(cat_mask)
../_images/125-lraspp-segmentation-with-output_28_0.png

draw_segmentation_masks() 関数を使用すると、これらのマスクを元の画像の上にプロットできます。この関数はマスクがブールマスクであることを想定していますが、上記のマスクには [0, 1] の確率が含まれています。ブールマスクを取得するには、次の手順を実行します。

class_dim = 1
boolean_cat_mask = (normalized_mask.argmax(class_dim) == sem_class_to_idx['cat'])

そして、元の画像の上にブールマスクをプロットできるようになりました。

from torchvision.utils import draw_segmentation_masks

show(draw_segmentation_masks(image, masks=boolean_cat_mask, alpha=0.7, colors='yellow'))
../_images/125-lraspp-segmentation-with-output_32_0.png

OpenVINO IR モデルの結果を表示

normalized_mask = torch.nn.functional.softmax(torch.from_numpy(res_ir), dim=1)
boolean_cat_mask = (normalized_mask.argmax(class_dim) == sem_class_to_idx['cat'])
show(draw_segmentation_masks(image, masks=boolean_cat_mask, alpha=0.7, colors='yellow'))
../_images/125-lraspp-segmentation-with-output_34_0.png