画像の差分からセグメンテーションのアノテーションを作りたい！

Introduction

背景

現在業務でセグメンテーションモデルの構築・精度検証に携わっているのですが、その一環で「モデルの予測前と予測後の画像を比較して、予測値を抽出する」必要性が発生しました。ちょっと何言ってんだこいつ、という感じだと思うので詳しく説明します。

まずセグメンテーションモデルについてですが、こちらについてはネット上に色々と資料が転がっているのでそちらを参照してもらえればと思います。端的にいうとある画像上に存在する物体をピクセルレベルで検出する、という用途を持っています。

Dive into Deep Learningより抜粋

さてこのようなセグメンテーションモデルを構築したい場合、学習データを揃えることが必要になってくるのですが、セグメンテーションにおける学習データというは基本的に

• 画像
• アノテーション

の二つで構成されています。ここでいうアノテーションとは、学習対象となるセグメンテーションについての情報が記述されたファイルであり、上記の例でいうと画像上のどのピクセルが「dog」「cat」「background」に該当するのか書かれているファイルとなります。
アノテーションファイルの書式はいくつか存在するのですが、一般的によく使われているのはCOCOフォーマットというものであり、こちらは下記のような構成をしています。

COCOフォーマット

```json { "images": [ { "id" : int, "width" : int, "height" : int, "file_name" : str, } ], "categories": [ { "id" : int, "name": str, "supercategory": str, } ], "annotations": [ { "id": int, "image_id": int, "category_id": int, "segmentation": RLE (or [Polygon]), "area": float, "bbox": [x,y,width,height], "iscrowd": 0 or 1, } ] } ```

上記segmentationの中にマスクの情報（写真上どのピクセル・箇所が何に該当するのか）を含んでいるのですが、矩形の形を座標で表す [x_0, y_0, x_1, y_1, x_2, y_2, ...] というPolygonフォーマットか、これらを圧縮したRLEフォーマットで記述されます。

問題提起

ここまでセグメンテーションの一般的な話をしてきましたが、ここから本題に入ろうと思います。本来セグメンテーション用の学習データ（あるいは学習済みモデルの出力値）というのは、前述のとおり画像と対応するアノテーションファイルで構成されるのですが、今回業務の中で「アノテーションファイルがない」という事態に陥りました。
手元にあるのは「元画像 (orginal)」と「セグメンテーションされた画像 (masked)」のみ。百聞は一見に如かず。下記を見ていただけると一目瞭然かなと。

さてこのように2枚の画像から「セグメンテーションの箇所（赤く塗られたところ）」のみ抽出しアノテーションファイルを作成するにはどうすればよいのか？

こちらが本稿で回答したいと問題となっています。

ちなみにここで紹介する手法は差分からのマスク抽出以外にも、バイナリマスクからのCOCOファイルへの変換にも使えたりします。

手法

今回画像の差分からアノテーションファイルを抽出するにあたっては二つの関数を作成しました。

1. 画像の差分を取り検知された物体をマスクに変換する関数
2. バイナリマスクをCOCOアノテーションファイルに変換する関数

それぞれ実装を見ていきたいと思います。

画像差分からマスクの生成

画像差分の抽出

画像を取り扱うにあたって、今回はOpenCVを中心に利用します。
まずは画像の差分の検出ですが、こちらは双方の画像間にて各ピクセルをそれぞれのチャンネル（RGB）で差分を取ることで可能となります。

import cv2
orginal_image_path = "hoge/original.jpeg"
detected_image_path = "hoge/detected.jpeg"

original_image = cv2.imread(original_image_path)
detected_image = cv2.imread(detected_image_path)

diff = cv2.absdiff(original_image, detected_image)  # 差分の抽出
gray_diff = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)  # 後続のためにグレースケールに変換

しかし差分を取るだけではマスクの色の濃ゆさの違いが現れたり、ノイズが入ってしまう可能性があります。というのも、今回は簡単のためオリジナルの画像とセグメンテーションありの画像が同サイズであることを仮定していますが、画像保存時の圧縮率等が異なる場合、微弱な差分が発生しうるためです。
ノイズは適当な閾値を決めた上でハイパスフィルタリングを実施します。閾値の設定は平均値などを活用することでよりシステマチックに決めることも可能です。

noise_threshold = 20
ret, binary_diff = cv2.threshold(gray_diff, noise_threshold, 255, 0)  # high pass filtering

画像差分（左）とノイズを除去したもの（右）

より高度な手法としてSSIMなども存在する模様です。

物体の検出

さて上記の流れですでにバイナリマスク（背景が0で物体が1の元画像と同じサイズの白黒画像）を作成できたわけですが、物体検知の場合バイナリマスクだけでは不十分なケースが多く（インスタンスセグメンテーションなど）個々の物体を一つづつアノテーションするのが好ましいです（上記COCOファイルのフォーマットも参照）。そのためバイナリマスク上で物体を切り分ける必要があります。

今回は個の物体がそれぞれ離れており単一のクラスしか存在しないケースを想定しますが、異なる複数の色でセグメンテーションされた多クラスのアノテーションファイル作成も同じ要領で可能かなと思います（画像の差分を取る際に色分けするイメージ）。

バイナリマスク上の物体（集合体）を検出するには輪郭の検知を実施します。白黒画像の輪郭検知はいくつか手法があるのですが、今回はOpenCVに実装されている findContours モジュールを使います（skimageについては後続のメモを参照）。

contours, hierarchy = cv2.findContours(binary_diff, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
contour_threshold = 20  # 輪郭の外周の閾値を設定 
contours = [c for c in contours if len(c) > contour_threshold]  # 外周が極端に小さいオブジェクト（単一ピクセルなど）を除外

ここで重要となるのが、findContours() の引数に cv2.RETR_EXTERNAL をセットすることです。OpenCVの findContours() は階層別に輪郭を抽出してくれる（輪郭の中に存在する輪郭の入れ子構造をツリー形式で返してくれる）のですが引数に RETR_EXTERNAL を指定することで最上位の階層（＝一番外側にある輪郭）のみを返してくれます。
これで仮にセグメンテーション画像に穴などがあっても、きれいに外枠だけ取ってくれます。

skimage ライブラリにも輪郭抽出は実装されていますが、こちらはOpenCVのように階層を考慮しないため、別途穴を埋める必要があります。

skimageを利用した輪郭抽出（左）はopencv（右）と比較して中の穴も捉えてしまう

穴を埋める手法はいくつかあるものの（ skimage.morphology.reconstruction など）、試した結果今回はOpenCVを利用する方が手っ取り早いという結論に至りました。

物体のポリゴン化

物体を検出したら、残るはCOCOアノテーションフォーマットに対応しているPolygonフォーマットに変換するのみです。ここでもやはりOpenCVの標準モジュールを活用します。

polygons = []
for contour in contours:  # 各物体に対してループ
    epsilon = 0.01 * cv2.arcLength(contour, True)  # 輪郭をなめらかにする係数を設定
    polygon = cv2.approxPolyDP(contour, epsilon, True)  # 輪郭からPolygonオブジェクトに変換
    polygons.append(polygon.flatten().tolist())

アノテーションファイルを作成するうえでは上記の通りPolygonオブジェクトのリストのみでも問題ないのですが、COCOアノテーションをPythonで使いやすくする公式のライブラリ pycocotools に実装されてある mask モジュールを利用してRLEフォーマットにエンコードしておくことで、後続の作業が楽になります。

from pycocotools import mask as maskUtils
img_height, img_width = gray_diff.shape[:-1]
rle_masks = maskUtils.frPyObjects(polygons, img_heigth, img_width)

実装関数

import cv2
from pycocotools import mask as maskUtils

def get_mask_from_diff(original_image_path, detected_image_path, noise_threshold=20, contour_threshold=20):
    # Taking image diff
    original_image = cv2.imread(original_image_path)
    detected_image = cv2.imread(detected_image_path)
    diff = cv2.absdiff(original_image, detected_image)
    gray_diff = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)  # グレースケール化

    # 輪郭抽出
    ret, binary_diff = cv2.threshold(gray_diff, noise_threshold, 255, 0)
    cv_contours, hierarchy = cv2.findContours(binary_diff, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 外周が小さい輪郭を除去
    contours = [c for c in cv_contours if len(c) > contour_threshold]
    if not contours:
        # 輪郭がなければマスクもない
        return []

    # ポリゴン化
    polygons = []
    for contour in contours:
        epsilon = 0.01 * cv2.arcLength(contour, True)
        polygon = cv2.approxPolyDP(contour, epsilon, True)
        polygons.append(polygon.flatten().tolist())

    # COCO用RLEマスクに変換
    rle_masks = maskUtils.frPyObjects(polygons, gray_diff.shape[0], gray_diff.shape[1])

    return rle_masks

バイナリマスクからアノテーションファイルの作成

続いての2ステップ目ですが、残るはアノテーションファイルを作成するのみです。バイナリマスクからアノテーションファイルを作成する方法はすでにリソースが幾つかあるので細かい解説は割愛します。一点コードの簡略化ポイントとしては、前述の通りあらかじめ物体のセグメンテーションをRLE型式でエンコードすることで、アノテーションに必要な物体の面積等を簡単に計算できることです。

実装関数

ここでは全てのオリジナル画像が img_dir 配下にあり、同フォルダ内にて差分を取りたい画像が同じファイル名の接尾辞のみ変えた形（オリジナルが hoge.jpeg であれば hoge_detected.jpeg など）であると仮定しています。

import os
import datetime
import cv2
from pycocotools import mask as maskUtils

def convert_diff_to_coco(image_dir, outfile_path=None):
    coco_annotation = {
        "images": [],
        "annotations": [],
        "categories": []
    }

    # カテゴリ情報の付与
    coco_annotation["categories"].append({
        "id": 0,
        "name": "hoge",
    })

    # 各画像に対してイテレーション
    annotation_id = 1
    for image_filename in os.listdir(image_dir):
        
        # ディレクトリ内の画像ファイルだけを抽出
        valid_extensions = ('.jpg', '.jpeg', '.png')
        if not image_filename.lower().endswith(valid_extensions):
            continue
        # マスクあり画像はオリジナル画像と接尾辞のみ異なる想定
        detected_image_filename = image_filename.replace(".jpeg", "_detected.jpeg")

        # 画像読み込み
        original_image_path = os.path.join(original_dir, image_filename)
        detected_image_path = os.path.join(original_dir, detected_image_filename)

        # 画像情報の付与
        original_image = cv2.imread(original_image_path)
        image_id = len(coco_annotation["images"]) + 1
        coco_annotation["images"].append({
            "id": image_id,
            "height": original_image.shape[0],
            "width": original_image.shape[1],
            "file_name": image_filename
        })

        # アノテーション（マスク情報）の抽出
        rle_masks = get_mask_from_diff(original_image_path, detected_image_path)
        if not rle_masks:
            # マスクがなければスキップ
            continue
        areas = maskUtils.area(rle_masks)  # 各物体の面積の計算
        bboxes = maskUtils.toBbox(rle_masks).tolist()  # 各物体のBboxの計算

        # アノテーションの付与
        for rle_mask, area, bbox in zip(rle_masks, areas, bboxes):
            # decode rle object into string
            rle_mask["counts"] = rle_mask["counts"].decode("ascii")
            coco_annotation["annotations"].append({
                "id": annotation_id,
                "image_id": image_id,
                "category_id": 0,  # 今回は単一のクラスしかないことを想定
                "segmentation": rle_mask,
                "area": int(area),
                "bbox": bbox,
                "iscrowd": 0
            })
            annotation_id += 1
    # JSONフォーマットに保存
    if outfile_path is None:
        today = datetime.now().strftime("%Y%m%d")
        outfile_path = f"{today}_annnotation.json"
    with open(outfile_path, "w") as f:
        json.dump(coco_annotation, f)

まとめ

今回は業務上大量の画像の差分からマスクを作成する必要があったため、OpenCVを利用したセグメンテーションのアノテーション方法を実装しました。