ふとsemantic segmentationモデルを学習してみたくなったので、自作データセットのアノテーションからモデル学習までを既存ツールの組み合わせでやってみました。備忘録がてら、手順をメモしておきたいと思います。

アノテーション

アノテーションには COCO Annotatorを使いました。GUIが直感的で使いやすく気に入っています。

細かい使い方に関しては解説記事などがわかりやすいと思います。Dockerが使える環境なら非常に簡単にセットアップすることができます。

今回は自作クラスの領域をpolygonツールでアノテーションしました。

アノテーションした後はデータセットのページで「Export COCO」のボタンを押し、対象クラスを選んでエクスポートします。エクスポートされたデータは、上から二段目のタブの「Exports」ボタンを押したページからダウンロードすることができます。

Pascal VOCフォーマットへの変換

COCO AnnotatorはMS COCOフォーマットでデータを出力するため、これをPascal VOCフォーマットに変換するために以下のツールを使います。

github.com

Gitでクローンして必要な環境を整えます。

git clone https://github.com/alicranck/coco2voc.git
pip install pycocotools

次にexample.pyを以下のように修正します。

# L37-39の変数をそれぞれ手元のデータのパスに合わせて修正
annotations_file = "<COCO AnnotatorからダウンロードしたJSONファイル>"
labels_target_folder = "<出力先ディレクトリ名>"
data_folder = "<データセットの画像が含まれるディレクトリ名>"

# その下に以下を追記
import sys
sys.path.append(data_folder)

さらに、coco2voc関数が呼び出されている箇所の引数のうちnの部分をn=Noneのように修正します。

その後、python example.pyを実行すると出力先ディレクトリに指定した場所にclass_labels/が生成されます。このディレクトリ以下の画像を学習に利用することになります。

DeepLabv3の学習

今回は以下のPyTorch実装を利用します。

github.com

環境整備

まずGitからクローンを行い、環境を整備します。

git clone https://github.com/VainF/DeepLabV3Plus-Pytorch.git

cd DeepLabV3Plus-Pytorch
pip install -r requirements.txt

自作データセット読み込みクラスの作成

続いて自作データセットの読み込みクラスを作成します。

今回はとにかくモデルを学習することを優先し、やり方はスマートではないですが、既存データセットのDatasetクラスを自作クラスのものに置き換えてしまう方針を取ります。

そのために、datasets/mydata.pyの名前でファイルを作成し、以下のように記述します。手元のデータに合わせて修正すべき箇所は次の２点です。

クラスコンストラクタ冒頭のorig_json, mask_dir, image_dirの3変数。
- それぞれCOCOフォーマットでのJSONファイル、coco2vocの出力先ディレクトリ以下にあるclass_labelsのパス、画像が保存されているディレクトリパスに修正します。
最後のdecode_target関数
- cmapは2次元配列で、各行が各クラスを表示する際の色となっています。クラス数に合わせて変更してください

import json
import os

import numpy as np
from PIL import Image
import torch.utils.data as data

class MyDataSegmentation(data.Dataset):
    def __init__(self, root,
                 year='2012',
                 image_set='train',
                 download=False,
                 transform=None):
        # 適宜必要に応じてimage_setがtrainの時とvalの時で分けることもできます
        orig_json = "path/to/json/file/of/dataset"
        mask_dir = "path/to/output/dir/class_labels"
        image_dir = "path/to/image/directory"

        self.images = []
        self.masks = []

        with open(orig_json) as h:
            coco_data = json.load(h)
        for img_data in coco_data["images"]:
            self.images.append(
                os.path.join(image_dir, os.path.basename(img_data["path"]))
            )
            self.masks.append(
                os.path.join(mask_dir, f"{img_data['id']}.png")
            )
        
        self.transform = transform


    def __getitem__(self, index):
        img = Image.open(self.images[index]).convert('RGB')
        target = Image.open(self.masks[index])
        if self.transform is not None:
            img, target = self.transform(img, target)

        return img, target
    
    def __len__(self):
        return len(self.images)
    
    @classmethod
    def decode_target(cls, mask):
        cmap = np.array([
            [0, 0, 0],
            [255, 0, 0]
        ])
        return cmap[mask]

続いて、datasets/__init__.pyを以下のように修正します。

#from .voc import VOCSegmentation
from .mydata import MyDataSegmentation as VOCSegmentation
from .cityscapes import Cityscapes

学習

--model引数により学習するモデルアーキテクチャを変更できるようです。詳細はレポジトリのREADMEをご参照ください。

python main.py --model deeplabv3plus_mobilenet  --gpu_id 0 --year 2012_aug --crop_val --lr 0.01 --crop_size 513 --batch_size 16 --output_stride 16

評価

resultsディレクトリ以下に画像で結果が出力されます。

python main.py --model deeplabv3plus_mobilenet --gpu_id 0 --year 2012_aug --crop_val --lr 0.01 --crop_size 513 --batch_size 16 --output_stride 16 --ckpt checkpoints/ --test_only --save_val_results

とあの方舟

COCO Annotatorでアノテーションした自作データセットでSemantic Segmentationモデル(DeepLabv3)を学習する

アノテーション

Pascal VOCフォーマットへの変換

DeepLabv3の学習

環境整備

自作データセット読み込みクラスの作成

学習

評価