PyTorchセグメンテーションモデルの変換とOpenCVでの実行

目標

このチュートリアルでは、以下の方法を学ぶ:

• PyTorchのセグメンテーションモデルを変換する
• 変換したPyTorchモデルをOpenCVで実行する
• PyTorchモデルとOpenCV DNNモデルの評価結果を得る

上記の項目を、FCN ResNet-50アーキテクチャを例として詳しく見ていく。

はじめに

PyTorch の分類モデルとセグメンテーションモデルを OpenCV API へ移行するパイプラインの要点は共通である。最初のステップは、PyTorch の組み込み関数 torch.onnx.export を使ってモデルを ONNX 形式に変換することである。次に、得られた .onnx モデルを cv.dnn.readNetFromONNX に渡すと、DNN 操作を行う準備が整った cv.dnn.Net オブジェクトが返される。

実践

この章では、以下の点を扱う。

1. セグメンテーションモデルの変換パイプラインを作成し、推論を提供する
2. セグメンテーションモデルを評価・テストする

評価またはテストのモデルパイプラインを実行するだけであれば、「モデル変換パイプライン」の部分は省略できる。

モデル変換パイプライン

本節のコードは dnn_model_runner モジュール内にあり、次の行で実行できる:

python -m dnn_model_runner.dnn_conversion.pytorch.segmentation.py_to_py_fcnresnet50

以下のコードには、下に列挙する各ステップの説明が含まれている:

1. PyTorchモデルをインスタンス化する
2. PyTorchモデルを.onnxへ変換する
3. OpenCV APIで変換したネットワークを読み込む
4. 入力データを準備する
5. 推論を実行する
6. 予測結果から色付けされたマスクを取得する
7. 結果を可視化する

# initialize PyTorch FCN ResNet-50 model
original_model = models.segmentation.fcn_resnet50(pretrained=True)
 
# get the path to the converted into ONNX PyTorch model
full_model_path = get_pytorch_onnx_model(original_model)
 
# read converted .onnx model with OpenCV API
opencv_net = cv2.dnn.readNetFromONNX(full_model_path)
print("OpenCV model was successfully read. Layer IDs: \n", opencv_net.getLayerNames())
 
# get preprocessed image
img, input_img = get_processed_imgs("test_data/sem_segm/2007_000033.jpg")
 
# obtain OpenCV DNN predictions
opencv_prediction = get_opencv_dnn_prediction(opencv_net, input_img)
 
# obtain original PyTorch ResNet50 predictions
pytorch_prediction = get_pytorch_dnn_prediction(original_model, input_img)
 
pascal_voc_classes, pascal_voc_colors = read_colors_info("test_data/sem_segm/pascal-classes.txt")
 
# obtain colored segmentation masks
opencv_colored_mask = get_colored_mask(img.shape, opencv_prediction, pascal_voc_colors)
pytorch_colored_mask = get_colored_mask(img.shape, pytorch_prediction, pascal_voc_colors)
 
# obtain palette of PASCAL VOC colors
color_legend = get_legend(pascal_voc_classes, pascal_voc_colors)
 
cv2.imshow('PyTorch Colored Mask', pytorch_colored_mask)
cv2.imshow('OpenCV DNN Colored Mask', opencv_colored_mask)
cv2.imshow('Color Legend', color_legend)
 
cv2.waitKey(0)

モデル推論を行うために、PASCAL VOC検証用データセットから以下の画像を使用する:

目標とするセグメンテーション結果は次のとおりである:

PASCAL VOCの色のデコードおよび予測されたマスクへのマッピングのために、PASCAL VOCクラスの完全な一覧と対応する色を含むpascal-classes.txtファイルも必要となる。

事前学習済みのPyTorch FCN ResNet-50を例に、各コードステップを詳しく見ていこう:

• PyTorch FCN ResNet-50モデルをインスタンス化する:

# initialize PyTorch FCN ResNet-50 model
original_model = models.segmentation.fcn_resnet50(pretrained=True)

• PyTorchモデルをONNX形式へ変換する:

# define the directory for further converted model save
onnx_model_path = "models"
# define the name of further converted model
onnx_model_name = "fcnresnet50.onnx"
 
# create directory for further converted model
os.makedirs(onnx_model_path, exist_ok=True)
 
# get full path to the converted model
full_model_path = os.path.join(onnx_model_path, onnx_model_name)
 
# generate model input to build the graph
generated_input = Variable(
    torch.randn(1, 3, 500, 500)
)
 
# model export into ONNX format
torch.onnx.export(
    original_model,
    generated_input,
    full_model_path,
    verbose=True,
    input_names=["input"],
    output_names=["output"],
    opset_version=11
)

このステップのコードは分類モデルの変換の場合と変わらない。したがって、上記コードの実行が成功すると、models/fcnresnet50.onnxが得られる。

• 前のステップで得られた ONNX モデルを渡して、cv.dnn.readNetFromONNX で変換済みのネットワークを読み込む:

# read converted .onnx model with OpenCV API
opencv_net = cv2.dnn.readNetFromONNX(full_model_path)

• 入力データを準備する:

# read the image
input_img = cv2.imread(img_path, cv2.IMREAD_COLOR)
input_img = input_img.astype(np.float32)
 
# target image sizes
img_height = input_img.shape[0]
img_width = input_img.shape[1]
 
# define preprocess parameters
mean = np.array([0.485, 0.456, 0.406]) * 255.0
scale = 1 / 255.0
std = [0.229, 0.224, 0.225]
 
# prepare input blob to fit the model input:
# 1. subtract mean
# 2. scale to set pixel values from 0 to 1
input_blob = cv2.dnn.blobFromImage(
    image=input_img,
    scalefactor=scale,
    size=(img_width, img_height),  # img target size
    mean=mean,
    swapRB=True,  # BGR -> RGB
    crop=False # center crop
)
# 3. divide by std
input_blob[0] /= np.asarray(std, dtype=np.float32).reshape(3, 1, 1)

このステップでは画像を読み込み、cv2.dnn.blobFromImage関数でモデル入力を準備する。この関数は4次元のblobを返す。cv2.dnn.blobFromImage では最初に平均値が減算され、その後にピクセル値がスケーリングされる点に注意してほしい。したがって、元の画像前処理の順序を再現するために mean に 255.0 を掛けている:

img /= 255.0
img -= [0.485, 0.456, 0.406]
img /= [0.229, 0.224, 0.225]

• OpenCV cv.dnn_Net による推論:

# set OpenCV DNN input
opencv_net.setInput(preproc_img)
 
# OpenCV DNN inference
out = opencv_net.forward()
print("OpenCV DNN segmentation prediction: \n")
print("* shape: ", out.shape)
 
# get IDs of predicted classes
out_predictions = np.argmax(out[0], axis=0)

上記のコードを実行すると、次の出力が得られる:

OpenCV DNN segmentation prediction:
* shape:  (1, 21, 500, 500)

21個の予測チャンネルのそれぞれ（21はPASCAL VOCクラスの数を表す）には確率が含まれており、ピクセルがそのPASCAL VOCクラスにどれだけ対応しやすいかを示す。

• PyTorch FCN ResNet-50 モデルの推論:

original_net.eval()
preproc_img = torch.FloatTensor(preproc_img)
 
with torch.no_grad():
 # obtaining unnormalized probabilities for each class
    out = original_net(preproc_img)['out']
 
print("\nPyTorch segmentation model prediction: \n")
print("* shape: ", out.shape)
 
# get IDs of predicted classes
out_predictions = out[0].argmax(dim=0)

上記のコードを実行すると、次の出力が得られる:

PyTorch segmentation model prediction:
* shape:  torch.Size([1, 21, 366, 500])

PyTorchの予測にも、各クラス予測に対応する確率が含まれている。

• 予測から色付きマスクを取得する:

# convert mask values into PASCAL VOC colors
processed_mask = np.stack([colors[color_id] for color_id in segm_mask.flatten()])
 
# reshape mask into 3-channel image
processed_mask = processed_mask.reshape(mask_height, mask_width, 3)
processed_mask = cv2.resize(processed_mask, (img_width, img_height), interpolation=cv2.INTER_NEAREST).astype(
    np.uint8)
 
# convert colored mask from BGR to RGB for compatibility with PASCAL VOC colors
processed_mask = cv2.cvtColor(processed_mask, cv2.COLOR_BGR2RGB)

このステップでは、セグメンテーションマスクの確率を、予測されたクラスの適切な色にマッピングする。結果を見てみよう:

モデルの拡張評価には、dnn_model_runner モジュールの py_to_py_segm スクリプトを使用できる。このモジュール部分については次のサブチャプターで説明する。

モデルの評価

dnn/samples で提案されている dnn_model_runner モジュールを使うと、PASCAL VOCデータセット上で完全な評価パイプラインを実行でき、次のPyTorchセグメンテーションモデルのテスト実行が可能になる:

• FCN ResNet-50
• FCN ResNet-101

この一覧は、適切な評価パイプライン設定を追加することでさらに拡張できる。

評価モード

以下の行は、モジュールを評価モードで実行する例である:

python -m dnn_model_runner.dnn_conversion.pytorch.segmentation.py_to_py_segm --model_name <pytorch_segm_model_name>

リストから選択したセグメンテーションモデルがOpenCVの cv.dnn_Net オブジェクトに読み込まれる。PyTorchモデルとOpenCVモデルの評価結果（ピクセル精度、平均IoU、推論時間）がログファイルに書き込まれる。推論時間の値は、得られたモデル情報を概観するためにチャートにも描画される。

必要な評価設定は test_config.py で定義されている:

@dataclass
class TestSegmConfig:
    frame_size: int = 500
    img_root_dir: str = "./VOC2012"
    img_dir: str = os.path.join(img_root_dir, "JPEGImages/")
    img_segm_gt_dir: str = os.path.join(img_root_dir, "SegmentationClass/")
 # reduced val: https://github.com/shelhamer/fcn.berkeleyvision.org/blob/master/data/pascal/seg11valid.txt
    segm_val_file: str = os.path.join(img_root_dir, "ImageSets/Segmentation/seg11valid.txt")
    colour_file_cls: str = os.path.join(img_root_dir, "ImageSets/Segmentation/pascal-classes.txt")

これらの値は、選択したモデルパイプラインに合わせて変更できる。

PyTorch FCN ResNet-50 の評価を開始するには、次の行を実行する:

python -m dnn_model_runner.dnn_conversion.pytorch.segmentation.py_to_py_segm --model_name fcnresnet50

テストモード

以下の行は、モジュールをテストモードで実行する例であり、モデル推論の手順を提供する:

python -m dnn_model_runner.dnn_conversion.pytorch.segmentation.py_to_py_segm --model_name <pytorch_segm_model_name> --test True --default_img_preprocess <True/False> --evaluate False

ここで default_img_preprocess キーは、scale、mean、std などの特定の値でモデルのテスト処理をパラメータ化するか、デフォルト値を使用するかを定義する。

テスト設定は test_config.py の TestSegmModuleConfig クラスで表現される:

@dataclass
class TestSegmModuleConfig:
    segm_test_data_dir: str = "test_data/sem_segm"
    test_module_name: str = "segmentation"
    test_module_path: str = "segmentation.py"
    input_img: str = os.path.join(segm_test_data_dir, "2007_000033.jpg")
    model: str = ""
 
    frame_height: str = str(TestSegmConfig.frame_size)
    frame_width: str = str(TestSegmConfig.frame_size)
    scale: float = 1.0
    mean: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0])
    std: List[float] = field(default_factory=list)
    crop: bool = False
    rgb: bool = True
    classes: str = os.path.join(segm_test_data_dir, "pascal-classes.txt")

デフォルトの画像前処理オプションは default_preprocess_config.py で定義されている:

pytorch_segm_input_blob = {
 "mean": ["123.675", "116.28", "103.53"],
 "scale": str(1 / 255.0),
 "std": ["0.229", "0.224", "0.225"],
 "crop": "False",
 "rgb": "True"
}

モデルテストの基盤は samples/dnn/segmentation.py で表現される。segmentation.py は、--input に変換済みモデルを指定し、cv2.dnn.blobFromImage 用の引数を設定すれば、単独で実行できる。

「Model Conversion Pipeline」で説明したOpenCVのステップを、dnn_model_runner を使って一から再現するには、以下の行を実行する:

python -m dnn_model_runner.dnn_conversion.pytorch.segmentation.py_to_py_segm --model_name fcnresnet50 --test True --default_img_preprocess True --evaluate False