[Python] 機械学習で写真から影を削除する [Ghost-free Shadow Removal]

2021年11月29日月曜日

Artificial Intelligence

本記事では、GANをベースとした機械学習手法を用いて写真から影を消す方法を紹介します。

アイキャッチ

概要

本記事ではTowards Ghost-free Shadow Removal via Dual Hierarchical Aggregation Network and Shadow Matting GANで提唱されている手法を用いて、任意の影を含む画像から学習済みモデルを用いて、影の削除を行っていきます。

なお、学習データを用意すれば、自身でカスタマイズしたモデルを生成することも可能です。
学習データは以下のような形式となっています。

学習データセット

上図の左の画像から順に、影を含む画像(shadow image)、影の場所を示したマスク画像(shadow mask)、影を除いた画像(shadowfree image)の3画像となっています。

Ghost-free Shadow Removalの導入手順

セットアップ1: conda環境構築

それでは早速、開発環境にGhost-free Shadow Removalをセットアップしていきます。
動作確認は下記の環境で行っています。

OS: Ubuntu 18.04.3 LTS
GPU: GeForce GTX 1080

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

それでは、Ghost-free Shadow Removalをインストールしていきます。他の機械学習環境に影響を与えないためにMinicondaの仮想環境上に構築していきます。Minicondaのインストール手順は公式ドキュメントをご参照ください。

conda環境構築手順を下記に記載します。

# codeをgit clone
$ git clone https://github.com/ceciliavision/perceptual-reflection-removal.git
$ cd perceptual-reflection-removal

# conda環境をymlファイルを使って作成
$ conda env create -f env.yml

# 作成したconda環境をactivate 
$ conda activate perceptual-reflection-removal

$ conda install -c anaconda pillow scikit-image

# tensorflowのインストール確認
$ python
>>> from tensorflow.python.client import device_lib
>>> device_lib.list_local_devices()

# codeをgit clone
$ git clone https://github.com/vinthony/ghost-free-shadow-removal.git
$ cd ghost-free-shadow-removal/

以上でconda環境構築は完了です。

セットアップ2: 学習済みモデルとテストデータ準備

続いて学習済みモデルとテストデータを準備していきます。
テストデータは下記の4画像を使用していきます。

テスト用画像

学習済みモデルは以下のように準備していきます。

$ pwd
/perceptual-reflection-removal/ghost-free-shadow-removal

$ mkdir Models
$ cd Models
$ wget https://www.vlfeat.org/matconvnet/models/imagenet-vgg-verydeep-19.mat
# 下記リンクからSRD+をダウンロード
# https://uofmacau-my.sharepoint.com/personal/yb87432_umac_mo/_layouts/15/onedrive.aspx?id=%2Fpersonal%2Fyb87432%5Fumac%5Fmo%2FDocuments%2Fghost%2Dfree%2Dshadow%2Dremoval%2Fsrdplus%2Dpretrained%2Ezip&parent=%2Fpersonal%2Fyb87432%5Fumac%5Fmo%2FDocuments%2Fghost%2Dfree%2Dshadow%2Dremoval
$ unzip srdplus-pretrained.zip

$ cd ..
$ mkdir outputs inputs
# inputsにテスト用画像を格納してください

ここまでの作業で以下のようなディレクトリ構造になっていれば準備完了です。

ghost-free-shadow-removal
│
├── Models
│   ├── imagenet-vgg-verydeep-19.mat
│   ├── srdplus-pretrained
│   │   ├── checkpoint
│   │   ├── model_best.ckpt.data-00000-of-00001
│   │   ├── model_best.ckpt.index
│   │   └── model_best.ckpt.meta
│   └── srdplus-pretrained.zip
├── inputs
├── outputs

影の除去

最後にテスト用の4写真から影を除去してみます。

# カレントディレクトリ確認
$ pwd
中略/envs/perceptual-reflection-removal/ghost-free-shadow-removal

# 影の除去実行
$ python demo.py --model ./Models/srdplus-pretrained/ --vgg_19_path ./Models/imagenet-vgg-verydeep-19.mat --input_dir ./inputs --result_dir ./outputs

出力結果は以下の通りとなります。

影除去結果

上段が入力画像、下段が影除去結果となります。

左2つは非常に綺麗に除去できています。一方で右2つは、少し跡が残っているか、一部残ってしまっていますね。 この点追加で学習が必要ですね。

影の合成モデルのトレーニング

ここからは、自作モデルのトレーニング方法を紹介していきます。

ISTDデータセットを使って学習を行う方法を記載していきます。

# カレントディレクトリ確認
$ pwd
中略/envs/perceptual-reflection-removal/ghost-free-shadow-removal

# Google Driveからファイルをダウンロードするツールをインストール
$ pip3 install gdown

# データセットを格納するディレクトリを作成
$ mkdir isid_dataset
$ cd isid_dataset

# ISTDデータセットをダウンロード
$ gdown --id 1I0qw-65KBA6np8vIZzO6oeiOvcDBttAY -O ./

# 解凍
$ sudo apt install unrar
$ unrar x ISTD_Dataset.rar

# shadow free データセットをダウンロード
$ gdown --id 1PPAX0W4eyfn1cUrb2aBefnbrmhB1htoJ -O ./
$ unzip shadow_USR.zip
$ mv ./shadow_USR/shadow_free ./

ISTDデータセットは以下のような構成となっています。

ISTD_Dataset
├── test
│   ├── test_A # 影を含む画像
│   ├── test_B # 影のみを切り出したmask画像
│   └── test_C # 影を除いた画像
└── train
    ├── train_A # 影を含む画像
    ├── train_B # 影のみを切り出したmask画像
    └── train_C # 影を除いた画像

以下は左からtrain_A、train_B、train_Cの順に並んでいます。

ISTDデータセットサンプル

USR shadow freeは以下のような画像が格納されています。

USR shadow freeサンプル画像

同様のデータセットを自作する場合は、mask画像や、影を除いた画像が必要になりますね。
下記の記事がデータセットの作成に役立つかもしれません。ご確認ください。

続いて、ISTDデータセットを基に、合成の影画像を生成します。
--taskは任意のタスク名称を
--data_dirはtrainまでのディレクトリを
--use_gpuは使用するgpuの数を、CPU飲みであれば0を
--is_trainingはtrainの場合は1を、テストの場合は0を設定してください。

# カレントディレクトリ確認
$ pwd
中略/envs/perceptual-reflection-removal/ghost-free-shadow-removal

$ python train_ss.py \
--task "istd_synth" \
--data_dir ./isid_dataset/ISTD_Dataset/train/ \
--use_gpu 2 \
--is_training 1

TypeError: expected str, bytes or os.PathLike object, not dict

上記エラーが発生した場合は、train_ss.pyを以下のように修正ください。

loss_percep = compute_percep_loss(shadowed_image, target, vgg_19_path=vgg_19_path)

# 上記を以下の通り修正
vgg_mat_path = "./Models/imagenet-vgg-verydeep-19.mat"
loss_percep = compute_percep_loss(shadowed_image, target, vgg_19_path=vgg_mat_path)

AttributeError: 'NoneType' object has no attribute 'model_checkpoint_path'

上記エラーが発生した場合において、is_trainingが正常にTrueになっていない場合は、
train_ss.pyを以下のように修正ください。

parser.add_argument("--is_training", default=1, help="training or testing")
# 上記を以下の通り修正
parser.add_argument("--is_training", default=True, help="training or testing")

is_training=ARGS.is_training==1
# 上記を以下の通り修正
is_training=ARGS.is_training

その後コマンドを下記のように変更します。

# カレントディレクトリ確認
$ pwd
中略/envs/perceptual-reflection-removal/ghost-free-shadow-removal

$ python train_ss.py \
--task "istd_synth" \
--data_dir ./isid_dataset/ISTD_Dataset/train/ \
--use_gpu 2 \
--is_training True

影の除去モデルのトレーニング

続いて影の除去モデルのトレーニング方法をご紹介します。
なお、データセットは前述までと同様です。

以下のコマンドでトレーニングが開始します。

なお、前述の影の合成をスキップした場合は --use_daを0.0に設定してください。

$ python3 train_sr.py \
--task "istd_train" \
--data_dir ./isid_dataset/ISTD_Dataset/train/ \
--use_gpu 1 \
--is_training 1 \
--use_da 0.5

影の合成モデルと同様のエラーが発生した場合は前述のエラー解消方法をご確認ください。

まとめ

本記事では、Ghost-free Shadow Removalで画像から影を除去する方法を紹介しました。

上記のように動かすことは容易かもしれませんが、より精度を求める場合、質の高い学習データセットが求められます。
しかしながら学習データセットの作成はコストがかかる作業になり得ます。
極力自動生成を第一候補に考えると良いかもしれません。

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


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

参考文献

1. 論文 - Towards Ghost-free Shadow Removal via Dual Hierarchical Aggregation Network and Shadow Matting GAN

2. GitHub - vinthony/ghost-free-shadow-removal

3. GitHub - ceciliavision/perceptual-reflection-removal

AIで副業ならココから!

まずは無料会員登録

プロフィール

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

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


Twitter

カテゴリ

このブログを検索

ブログ アーカイブ

TeDokology