Pythonによる背景除去(Image Matting)モデル・ライブラリの紹介

はじめに

これまでに

• 前々回はopenCVによる背景除去
• 前回はDeepLabV3 + IndexNet Mattingによる背景の除去

を行いましたが、今回はその他のモデルやライブラリによる背景の除去を紹介します。 紹介の内容は以下になります。 他にも候補はありましたが、githubからcloneしただけではコードが不足しているものやバグが直っていないものが多くあり2つのみとなります。

• MODNet
• pymatting

MODNet

MODNetはPortrait Matting、つまり人物と背景の分離を行うDNNのモデルです。 MODNetはtrimapを使用せず、エンコーダー/デコーダー構造モデルではなくシングルモデルであるため軽量で速いことを 特徴として主張しています。よって精度と速度を両立していることから動画向きであるといえるかもしれません。

MODNetデモコード

例によってgoogle colaboratoryを使用して、コードを作成し実行しました。 ソースコードは以下になります。 流れとしては、MODNetをCloneして画像ファイルをアップロードし、demo.pyを実行するだけです。

    import os

%cd /content
if not os.path.exists('MODNet'):
    !git clone https://github.com/ZHKKKe/MODNet
%cd MODNet/

pretrained_ckpt = 'pretrained/modnet_photographic_portrait_matting.ckpt'
if not os.path.exists(pretrained_ckpt):
    !gdown --id 1mcr7ALciuAsHCpLnrtG_eop5-EYhbCmz \
            -O pretrained/modnet_photographic_portrait_matting.ckpt

import shutil
from google.colab import files

# clean and rebuild the image folders
input_folder = 'demo/image_matting/colab/input'
if os.path.exists(input_folder):
    shutil.rmtree(input_folder)
os.makedirs(input_folder)

output_folder = 'demo/image_matting/colab/output'
if os.path.exists(output_folder):
    shutil.rmtree(output_folder)
os.makedirs(output_folder)

# upload images (PNG or JPG)
image_names = list(files.upload().keys())
for image_name in image_names:
    shutil.move(image_name, os.path.join(input_folder, image_name))

!python -m demo.image_matting.colab.inference \
--input-path demo/image_matting/colab/input \
--output-path demo/image_matting/colab/output \
--ckpt-path ./pretrained/modnet_photographic_portrait_matting.ckpt
    

出力結果の画像に確認用の処理を行った結果は以下になります。またDeepLabV3との比較を行いました。 使用した画像はopenCVの付属データになります。

MODNetの出力結果

DeepLabV3の出力結果

試した枚数は多くなく小さい画像を使用したのでデモのモデルでは良さを実感することはできませんでしたが、 MODNetはself-supervised sub-objectives consistency (SOC) strategyという現実世界に適応させる戦略をとり、 また現実世界のデータセットを使用して学習を行っているので、 上記の画像よりは傾向としては日常的な画像の方が比較してよい性能が出るかもしれません。 文献によれば従来のものよりはこのような現実的な画像のデータセットでは分離性能は良いようです。 ただ処理の速度は速いので動画に向いているといえると思います。

pixabay から別の肖像画で確認しました。

また人物だけではなく以下のような鳥の画像でも前景を抽出することができましたが、なんでも抽出できるわけではありませんでした。 これくらいの精度であればtrimapを作成してきれいに抽出することが可能です。

pymattingライブラリ

pymattingは元画像とtrimapがあれば、それをを用いて手軽に前景抽出を行うことのできるライブラリです。 得られたalpha画像を用いて別の背景とブレンドすることも可能です。 まずpipでpymattingをインストールします。

    !pip install pymatting

trimapを生成します。画像は先ほどのMODNetから得られたmask画像を処理してtrimapを作成しました。 以下ではあいまいな部分は128にしていますが、確実の前景である部分を100%、確実に背景である部分を0%として連続的であっても構いません。

    %cd /content
img = cv2.imread("base.png", cv2.IMREAD_COLOR)
mask = cv2.imread("mask.png", cv2.IMREAD_GRAYSCALE)
mask[mask > 127] = 255
mask[mask <= 127] = 0
kernel = np.ones((3,3),np.uint8)
sure_bg = cv2.dilate(mask, kernel, iterations=5)
sure_fg = cv2.erode(mask, kernel, iterations=5)
trimap = np.full(mask.shape, 128)
trimap[sure_fg == 255] = 255
trimap[sure_bg == 0] = 0
cv2.imwrite("trimap.png", trimap)
    

cutout関数に元画像、trimap、出力画像パスを引数をとるだけです。 これでcutout.pngに背景が除去された画像が出力されます。

    from pymatting import cutout

cutout(
    "base.png",
    "trimap.png",
    "cutout.png")
    

きれいに背景が除去された画像が得られたことを確認できました。

まとめ

シンプルなコードで高い精度で前景画像を抽出することができました。 ただtrimapなしでは背景が残ってしまうこともあるようでした。 深層学習において大変な部分が学習データ集めとtrainの部分かと思いますが、 オブジェクトにこだわらなければ事前学習モデルが使用可能です。

Reference