猫画像から猫部分のみを抽出する(matting/semantig segmentation)

18729 ワード
Python ComputerVision Keras Python テキストリンク

概要

Pixabayから収集した猫画像から、できるだけ少ない手間で猫部分のみを抽出させてみました。

前景抽出

画像から前景と背景を分離する処理をマッティング(matting)と呼ぶそうです。Photoshopではメジャーな機能のようで、gimpであれば「前景抽出」という名称の機能がそれに該当します。
画像を「前景」「背景」「そのどちらか」の3種類に分割して、それぞれの領域における画像特徴を見て境界を決定してゆくのが一般的な挙動のようです。

この処理を自動化することはなかなか困難なようで、ざっと調べた範囲では、ポートレート写真を対象とした、人物のみを抽出するDeep Automatic Portrait Mattingという手法しか見つけられませんでした。

これとは別に、マッティングに用いる情報(trimap)を自動生成するという手法の論文はありました。

しかし適切なtrimapさえ生成できれば自動マッティングができるとも限りません。実際、gimp等で処理を行う際には、細かいところで前景・背景の判定が期待通りに分類されず微調整を必要とする場面はよくあります。

セマンティックセグメンテーション

別のアプローチとして、セマンティックセグメンテーションを用いることを考えてみました。セマンティックセグメンテーションは、一枚の入力画像に対して、各種オブジェクトの判別をピクセル単位で行うものです。

以前書いた記事では、ChainerCVはセマンティックセグメンテーションにも対応していると書きましたが、事前学習モデルとして用意されているものはCamVid datasetとよばれる訓練データに基づいたもので、空・道路・自動車・歩行者などといった11クラスの分類に対応しています。今回扱いたい猫の情報はありません。

他のセマンティックセグメンテーション実装やデータセットについて探してみたところ、DilatedNetのKeras実装が、PASCAL VOC2011 challangeのデータセットをベースにしたデータセットを用いた訓練済みモデルを含めて配布していたので、今回はそれを用いることにしました。

セットアップ

コードをcloneしてrequirements.txtの依存関係を満たすpythonモジュールを一通りインストールし、事前学習モデルをダウンロードすれば利用可能です。
古いKerasとTensorFlow backend(0.12.1)を用いるので、専用のvirtualenv環境を用意するのがよいでしょう。

$ pip install -r requirements.txt
$ curl -L https://github.com/nicolov/segmentation_keras/releases/download/model/nicolov_segmentation_model.tar.gz \
 | tar xvf -
$ python predict.py --weights_path \ 
 conversion/converted/dilation8_pascal_voc.npy \
 images/cat.jpg

dilation8_pascal_voc.npyが事前訓練モデルファイルになります。同梱のimages/cat.jpgを処理すると、images/cat_seg.pngというファイルが生成されます。

未知の画像に適用

実際に、このモデルを使って未知の画像を処理してみます。残念ながら、この実装では入力画像の大きさに制限があり、幅と高さがおおよそ500未満でないと処理できないようです。あらかじめその制限を満たすサイズに画像をリサイズした上で、処理を行ってみました。

処理対象画像:

セグメンテーション出力

猫画像と判別された領域は#400000(R:0x40, G:0x00, B:0x00)の色で識別されます。一見うまくいっているようです。あとはこの画像の各ピクセルを参照しつつ、オリジナル画像から猫画像以外に判定された部分を塗りつぶせば望みの結果が得られそうです。

しかしながら境界部分をよくよくみると、きっちりと数値通り綺麗に分かれているわけではありません。

絶対値比較の場合

このようにエッジ周辺の値は若干ぶれているので、単純にセグメンテーションピクセルの絶対値で比較して前景画像かどうかを判断させると、この部分の画像が欠落してしまいます。

色空間同士のノルムを見る

いろいろと考えたり試したりした結果、numpyのlinalg.normを使ってセグメンテーションピクセルの色と#400000との距離を計算して、一定数以内であれば猫画像ピクセルである、と判別する方法をとりました。64の半分である32にあたる範囲を有効とみなした結果が以下になります。

割と良い感じになりました。この結果を出力するためのコードを次に示します。

# -*- coding: utf-8 -*-

import argparse
import os, sys
from PIL import Image
import numpy as np

def get_args():
    p = argparse.ArgumentParser()
    p.add_argument("--contents-dir", '-c', default=None)
    p.add_argument("--segments-dir", '-s', default=None)
    p.add_argument("--output-dir", '-o', default=None)
    p.add_argument("--cat-label-vals", '-v', default="64,0,0")
    args = p.parse_args()
    if args.contents_dir is None or args.segments_dir is None or args.output_dir is None:
        p.print_help()
        sys.exit(1)

    return args

def cat_col_array(val_str):
    vals = val_str.split(',')
    vals = [int(i) for i in vals]
    return vals

def cmpary(a1, a2):
    v1 = np.asarray(a1)
    v2 = np.asarray(a2)
    norm = np.linalg.norm(v1-v2)
    if norm <= 32:
        return True
    return False

def make_images(cont_dir, seg_dir, out_dir, cat_vals):
    files = []
    for fname in os.listdir(cont_dir): # make target file list
        seg_fname = os.path.join(seg_dir, fname)
        if os.path.exists(seg_fname):
            files.append(fname)
    for fname in files:
        print("processing: %s" % fname)
        c_fname = os.path.join(cont_dir, fname)
        cont = np.asarray(Image.open(c_fname)).copy()
        s_fname = os.path.join(seg_dir, fname)
        seg = np.asarray(Image.open(s_fname))
        width, height, _ = seg.shape
        for y in range(height):
            for x in range(width):
                vals = seg[x, y]
                if not cmpary(vals, cat_vals):
                    cont[x, y] = [255, 255, 255]
        out = Image.fromarray(cont)
        o_fname = os.path.join(out_dir, fname)
        out.save(o_fname)

def main():
    args = get_args()
    cat_vals = cat_col_array(args.cat_label_vals)
    make_images(args.contents_dir, args.segments_dir, args.output_dir,
                cat_vals)

if __name__ == '__main__':
    main()

課題

今回用いたモデルをPixabayから収集、選別した画像に一通り適用すると、必ずしも100%期待した結果を示すとは限りませんでした。猫以外の残り19種類どれかの特徴があると判断されるケースがあります。

胴体の一部が犬判定される

以下の画像を処理したところ、一部の胴体部分が犬として判別されました。

バイオレットっぽい色合い(#400080)が犬と判別された領域です。確かにこんな色合いの犬は割といるような気がします。この結果を先のスクリプトに与えると次のようになります。

今回用意した画像に犬はほとんど含まれていないので、犬と判断された領域も残すようにしても良いかもしれません。

黒猫の扱い

黒い猫はそもそも綺麗に写真へ収めるのが難しいので、判別に関しても問題が起きやすいようです。

影部分が猫の一部になっていたり、真っ黒にしか見えない胴体部分が消えたりしています。

データセット

今回セグメンテーション/マッティング対象としたデータへのURLをリスト化してgithubにおいています。

1100ちょっとの画像になります。

今後

この程度のデータでもpix2pixのedges2catsはある程度再現できるのではないかと思っているので、実際に学習をさせてみるつもりです。

他には、今回使ったDilatedNetやChainerCVで利用できるSegNetの構造について、論文を読んで理解を深めたいところです。また、DilatedNet Keras実装で用いている訓練用データをChainerCVのSegNetで学習させてみるといったこともやってみたいところです。