Panoptic DeepLabをDetectron2で試す

DeepLabはGoogleの研究者が発表した一連のSemantic Segmentation手法です.v1, v2, v3, v3+のバージョンがあり,基本的にはEncoder-Decoderモデルで,Atrous Convolution(Dilated Convolution)を使ったものです.Panoptic DeepLabは,DeepLabのPanoptic Segmentationバージョンです.Semantic Segmentationではクラス単位で領域を塗り分ける手法ですが,Panoptic Segmentationは,物体インスタンス単位で塗り分けるInstance Segmentationと,Semantic Segmentationを統合した手法です.Semantic Segmentationでは画像全体が領域分割されるのに対し,Instance Segmentationでは,検出した物体ごとにそれぞれの物体を区別して領域が分割されます.そして,このPanopti Segmentationでは,画像全体をセグメンテーションしつつ,検出した物体は個々に区別した結果が得られます.

Panoptic Deeplabはオリジナルの実装もありますが,Facebookが公開しているDetectron2でも実装されており,簡単に試すことができます.公開されているコードでは,フォルダ内の画像に対して学習したり評価したりができますが,プログラムに組み込むことが難しかったので,サンプルコードを作ってみました.

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import cv2
from tqdm import tqdm
from pathlib import Path
from argparse import ArgumentParser
import train_net
from detectron2.data import MetadataCatalog
from detectron2.engine import DefaultPredictor
from detectron2.utils.visualizer import Visualizer


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("images", type=Path, help="image directory")
    parser.add_argument("output", type=Path, help="output directory")
    args = parser.parse_args()

    # Detectron2 用の
    args_det = train_net.default_argument_parser().parse_args()
    args_det.config_file = "panoptic_deeplab_R_52_os16_mg124_poly_200k_bs64_crop_640_640_coco_dsconv.yaml"
    args_det.opts = ["MODEL.WEIGHTS", "model_final_5e6da2.pkl"]
    cfg = train_net.setup(args_det)
    model = DefaultPredictor(cfg)

    args.output.mkdir(parents=True, exist_ok=True)

    # ディレクトリ内の画像全てに対して処理
    for img in tqdm(sorted(args.images.glob("*.jpg"))):
        # 画像の読み込み
        rgb_img = cv2.imread(str(img))

        # 推定
        predictions = model(rgb_img)
        panoptic_seg, segments_info = predictions["panoptic_seg"]

        # panoptic_seg は,画素ごとにクラス確率が入った,画像サイズxクラス数の配列
        # segments_info は,各領域ごとに,インスタンスIDが入った画像サイズの配列(?)が物体数分ある

        # 結果の可視化
        visualizer = Visualizer(
            rgb_img[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2
        )
        vis_output = visualizer.draw_panoptic_seg_predictions(
            panoptic_seg.to("cpu"), segments_info
        )

        # 保存
        fname = args.output / img.name
        vis_output.save(fname)

このサンプルコードは,指定したフォルダ内の画像全てに対してPanoptic Segmentationし,結果を可視化した画像を出力するものです.

実行するためには,ネットワークの設定ファイルと,学習済みモデルの重みファイルをダウンロードする必要があります.今回の例では,MSCOCOデータセットで学習したモデルを使います.以下のコマンドで,必要なファイルをダウンロードできます.ダウンロードに加え,ファイルのパスをsedコマンドで修正しています.

wget https://raw.githubusercontent.com/facebookresearch/detectron2/main/projects/Panoptic-DeepLab/train_net.py
wget https://raw.githubusercontent.com/facebookresearch/detectron2/main/projects/Panoptic-DeepLab/configs/Cityscapes-PanopticSegmentation/Base-PanopticDeepLab-OS16.yaml
wget https://raw.githubusercontent.com/facebookresearch/detectron2/main/projects/Panoptic-DeepLab/configs/COCO-PanopticSegmentation/panoptic_deeplab_R_52_os16_mg124_poly_200k_bs64_crop_640_640_coco_dsconv.yaml
wget https://dl.fbaipublicfiles.com/detectron2/PanopticDeepLab/COCO-PanopticSegmentation/panoptic_deeplab_R_52_os16_mg124_poly_200k_bs64_crop_640_640_coco_dsconv/model_final_5e6da2.pkl


sed --in-place -e "s/..\/Cityscapes-PanopticSegmentation\///" panoptic_deeplab_R_52_os16_mg124_poly_200k_bs64_crop_640_640_coco_dsconv.yaml

Panoptic DeepLabの実行方法は以下のとおりです.

python main.py /path/to/images outputs

1つ目の引数が入力画像のディレクトリ,2つ目の引数が出力先ディレクトリです.

コメントする