概要

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

前景抽出

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

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

Deep Automatic Portrait Matting

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

Automatic generation of trimap for image matting: International Journal of Machine Intelligence and Sensory Signal Processing: Vol 1, No 3

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

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

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

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

Semantic object classes in video

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

nicolov/segmentation_keras: DilatedNet in Keras for image segmentation

セットアップ

コードを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()