こんにちは。

以前の記事で、YOLACTというinstance segmentation手法を紹介しました。そのとき、自作のデータセット(custom dataset)を訓練データとして用いていました。今回は、その訓練データの生成についての記事です。

f:id:ume-boshi:20200911045104p:plain:w400 — こういうやつ

Instance segmentationの訓練データを作る際、下記の様な手順で行うことが主流だと思います。

web上からスクレイピングして対象物の画像を収集
LabelMeというツールを用いて、手動で輪郭の代表点をプロット & 各画像ごとにJSON出力
labelme2coco.pyを用いてjsonファイルを統合

annotationする画像枚数が少なければこの手法でも問題ないですが、枚数が多くなったり対象物の形状が複雑なほど手作業は面倒くさくなります。マゾな人は手作業して苦労したいのかもしれませんが、私はsemi-マゾなので苦労しないために実装に苦労しました。

概要

まあ、そんなこんなで、たった1つの短時間の映像から訓練データを半自動的に生成するようなtoolを作成しようと考えたわけです。

youtu.be

github.com

入出力としては、対象物を周囲から撮影した1分程度の動画を入力として、輪郭抽出とデータ拡張された画像とJSONファイルを出力とします。

我々はちょっとしたお絵描きをするだけで、画像1枚当たり10秒もかからないで処理できます。データ拡張も雑に実装してあるため、1枚の画像からn倍のデータを生成できます。

なので、1分の動画を入力しデータ拡張を10倍とするならば、10分程度の時間で{60 + 60*10} = 660枚の訓練データを生成できます。我ながら革命やな。

f:id:ume-boshi:20210206203646p:plain:w400

使用方法

想定している大まかな使用手順は、次の5段階です。

githubから本ツールをcloneしてくる
スマホやwebカメラで対象物を周囲から撮影し、PCに保存
backgroundのフォルダに、背景にしたい画像を保存
ツールを起動する：python annotateFromVideo.py targetVideo.mp4 targetObjectName
前景と背景を選択していく（jpgとJSONが自動的に生成される）
ツール終了後、githubからlabelme2cocoをcloneし、それを用いてJSONファイルを統合（trainval.jsonを出力）

背景画像について

MS COCOなどのデータセットから適当に1000枚程度コピーしてきて、保存することを想定しています。

ファイル名は下記の通りにしてください。

1 (1).jpg
1 (2).jpg
1 (3).jpg
︙

f:id:ume-boshi:20210207043537p:plain:w400 — こんな感じに

Windowsでは、画像ファイルを全選択して「名前を変更」をし、1つ目の画像を「1.jpg」とすると自動的に他の画像も変更できます。

LinuxやMacの人は適当にコマンドを打って解決してください（なげやり）。

対象物抽出について

輪郭抽出のためにGraphCutという前景と背景の分離手法を用いており、ユーザはマウスカーソルで前景と背景を色塗りしていくだけです。PowerPointの画像背景削除ツールに似た使用感。

ツール実行中に使用できるコマンドは5つあり、下記の機能を持ちます。

[ b ]必ず背景としたい箇所をマークできるモードに変更
[ w ]必ず前景としたい箇所をマークできるモードに変更
[ g ]背景と前景の抽出を試行する
[esc]背景と前景の抽出を確定し、次の画像へ進む
[ q ]途中でツールを終了できる

注意点

前景と背景の色塗り（白黒）を間違えた場合、そこは正しい色ですべて塗りつぶす必要があります
輪郭抽出をミスした場合、その画像に関する「./outputImages/○.json」ファイルを削除してプログラムを実行しなおすと再挑戦できます
「outputImages」と「contourResults」のフォルダを事前に作成していなければ正常に実行できません
labelme2cocoの実行時にはすべてのjsonファイルを参照するようになっており、変なjsonファイルを保存しないようにしてください。

実行環境

Python 3.8
- numpy 
- opencv-python

実行前後のディレクトリ構成は下記の通りです。

【実行前】

project
 /annotateFromVideo.py
 /targetVideo.mp4
 /outputImages    //訓練用データ
    /label.txt
 /contourResults    //輪郭確認用データ
 /backgound    //背景画像
    /1 (1).jpg
    /1 (2).jpg
    /1 (3).jpg
      ︙

【実行後】

project
 /annotateFromVideo.py
 /targetVideo.mp4
 /outputImages    //訓練用データ
    /label.txt
    /0.jpg
    /0.json
    /0-0.jpg
    /0-0.json
      ︙

 /contourResults    //輪郭確認用データ
    /0.jpg
    /0-0.jpg
      ︙

 /backgound    //背景画像
    /1 (1).jpg
    /1 (2).jpg
    /1 (3).jpg
      ︙

labelme2cocoでJSONを1つに統合する際、categoryのIDが0として出力されます。これはYOLACTの学習時には問題が生じるため、下記の様にコードを書き換えてもいいかもしれません。
もしくは、labeme2coco実行後のtrainval.jsonについて、"category"のIDを置換すると良いでしょう。

変更後labelme2coco ↓

     ︙

    def annotation(self, points, label, num):
        ︙
        annotation["category_id"] = label[0] + 1  # self.getcatid(label)
        annotation["id"] = self.annID
        return annotation]

    ︙
    def category(self, label):
        category = {}
        category["supercategory"] = label[0]
        category["id"] = len(self.categories) + 1
        category["name"] = label[0]
        return category
       ︙