[MODNet] 機械学習で写真の背景を削除・合成する [Python]

2022年3月1日火曜日

Artificial Intelligence

本記事では、MODNetを使用して、人物写真の背景を削除し、別の背景を合成する方法を紹介します。

アイキャッチ

MODNet

概要

MODNetは、2020年11月に論文発表された単一の入力画像を使用してリアルタイムでポートレートマットを作成するlight-weight Matting Objective Decomposition NETwork(MODNet)です。

ポートレートマットとは下図の右の白黒画像です。前景と背景の分離に用いるセグメンテーションフォーマットのことであり、このマスク画像を用いることで前景部分のみを切り出すことができます。

portrait_matte_sample

MODNetは、Atrous Spatial Pyramid Pooling(e-ASPP)と自己監視サブオブジェクト整合性(self-supervised sub-objectives consistency:SOC)戦略により、従来手法と比較しモデルの効率と堅牢性を向上させています。このことにより1080TiGPUで毎秒67フレームで実行可能という高速性を獲得しています。

MODNET_Arch
出典: MODNet: Real-Time Trimap-Free Portrait Matting via Objective Decomposition

MODNetの技術詳細はこちらの論文をご確認ください。
本記事では、MODNetを用いて写真からポートレートマットを生成し、ポートレートマットを使用して任意の背景を合成する方法をご紹介します。

また、以下の記事では有償にはなりますが、論文内容の解説、詳細なデモコードを記載しております。よろしければご参照ください。

デモ(Colaboratory)

それでは、実際に動かしながらMODNetを用いた背景合成を行っていきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

また、下記から直接Google Colaboratoryで開くこともできます。
Open In Colab

なお、このデモはPythonで実装しています。
Pythonの実装に不安がある方、Pythonを使った機械学習について詳しく勉強したい方は、以下の書籍やオンライン講座などがおすすめです。

環境セットアップ

それではセットアップしていきます。 Colaboratoryを開いたら下記を設定しGPUを使用するようにしてください。

「ランタイムのタイプを変更」→「ハードウェアアクセラレータ」をGPUに変更

まず、使用するライブラリをインポートしておきます。

import shutil
from google.colab import files
import os
import numpy as np
from PIL import Image

次にGitHubからソースコードを取得します。

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

以上で環境セットアップは完了です。

学習済みモデルのセットアップ

次にポートレートマットの予測に用いるモデルをセットアップしていきます。
論文発表元より学習済みモデルが提供されているため、こちらをGoogle Colaboratoryにダウンロードします。

%cd /content/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

/content/MODNet/pretrainedmodnet_photographic_portrait_matting.ckptがダウンロードされます。

テスト画像のセットアップ

ここでは、背景を除去したい画像をPCからGoogle Colaboratoryにアップロードします。

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)

background_folder = 'demo/image_matting/colab/background'
if os.path.exists(background_folder):
  shutil.rmtree(background_folder)
os.makedirs(background_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))

本記事では、ぱくたそ様の以下の画像を使用させて頂きます。

テスト画像

予測(ポートレートマット生成)

それでは、アップロードした画像をモデルに入力してポートレートマットを生成していきます。

%%time

!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

モデルのロード等を含みますが、Tesla T4で約6秒程で完了します。

処理時間

ポートレートマットを使用して、前景画像を生成して表示します。
下記関数でオリジナル画像、前景画像、ポートレートマット画像を連結したイメージを生成します。

def combined_display(image, matte):
  # calculate display resolution
  w, h = image.width, image.height
  rw, rh = 800, int(h * 800 / (3 * w))
  
  # obtain predicted foreground
  image = np.asarray(image)
  if len(image.shape) == 2:
    image = image[:, :, None]
  if image.shape[2] == 1:
    image = np.repeat(image, 3, axis=2)
  elif image.shape[2] == 4:
    image = image[:, :, 0:3]
  matte = np.repeat(np.asarray(matte)[:, :, None], 3, axis=2) / 255
  foreground = image * matte + np.full(image.shape, 255) * (1 - matte)
  
  # combine image, foreground, and alpha into one line
  combined = np.concatenate((image, foreground, matte * 255), axis=1)
  combined = Image.fromarray(np.uint8(combined)).resize((rw, rh))
  return combined

下記で、上記関数で生成したイメージを表示します。

# visualize all images
image_names = os.listdir(input_folder)
for image_name in image_names:
  matte_name = image_name.split('.')[0] + '.png'
  image = Image.open(os.path.join(input_folder, image_name))
  matte = Image.open(os.path.join(output_folder, matte_name))
  display(combined_display(image, matte))
  print(image_name, '\n')

出力結果は以下の通りです。
非常に正確に人物の領域を検出していることがわかります。

ポートレートマット生成結果

背景合成

ポートレートマットを使って、背景を合成していきます。
まず、合成したい背景をアップロードします。

# upload image (PNG or JPG)
backimage_name = list(files.upload().keys())
shutil.move(backimage_name[0], os.path.join(background_folder, backimage_name[0]))

本記事では、pixabay様のこちらの画像を使用させて頂きます。

背景画像

ポートレートマットを用いた背景合成用関数を定義します。

def background_composition(image_path, matte_path, bg_image_path):
  foreground = Image.open(image_path)
  matte = Image.open(matte_path)
  background = Image.open(bg_image_path).resize(foreground.size)
  
  composition_img = Image.composite(foreground, background, matte)

  return composition_img

最後に背景合成を行います。

image_names = os.listdir(input_folder)
for image_name in image_names:
  matte_name = image_name.split('.')[0] + '.png'
  image_path = os.path.join(input_folder, image_name)
  matte_path = os.path.join(output_folder, matte_name)
  bg_image_path = os.path.join(background_folder, backimage_name[0])
  compos_img = background_composition(image_path, matte_path, bg_image_path)
  display(compos_img)
  print(image_name, '\n')

  # 合成画像を保存
  compos_img.save(os.path.join(output_folder, "composition_" + image_name))

出力結果は以下の通りです。

合成画像

非常に自然に背景を合成できています。

まとめ

本記事では、MODNetを用いて写真の背景を削除・合成する方法をご紹介しました。
iPhoneやGoogle Pixelなどのポートレート機能も前景の人物を際立たせるために前景領域を検出しています。
日常生活で何気なく使っている機能にも機械学習は潜んでいますね。
これを機に機械学習に興味を持つ方が一人でもいらっしゃいましたら幸いです。

また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
もう少し学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。


また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。

参考文献

1.  論文 - MODNet: Real-Time Trimap-Free Portrait Matting via Objective Decomposition

2. GitHub - ZHKKKe/MODNet

AIで副業ならココから!

まずは無料会員登録

プロフィール

メーカーで研究開発を行う現役エンジニア
組み込み機器開発や機会学習モデル開発に従事しています

本ブログでは最新AI技術を中心にソースコード付きでご紹介します


Twitter

カテゴリ

このブログを検索

ブログ アーカイブ

TeDokology