本記事では、Lamaと呼ばれる機械学習手法を使って写真からの障害物を削除する方法をご紹介します。

Lamaとは

Lamaとは「Resolution-robust Large Mask Inpainting with Fourier Convolutions」の略称で、Samsungの調査機関が発表した機械学習を用いた画像修復の手法です。

イメージをつかんでいただくために下記のGIFをご覧ください。

出典: https://github.com/saic-mdal/lama

ご覧いただくと一目瞭然かと思います。

技術的には、削除範囲を指定し画像から除去
その後、Lamaの画像修復技術によって、削除した範囲を修復しています。

出典: https://arxiv.org/pdf/2109.07161v1.pdf

Lamaの導入手順

セットアップ

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

なお、GPU搭載のPCをお持ちでない場合は、以下からブラウザで動作確認することも可能です。

Google Colaboratory

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

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

# code clone
$ git clone https://github.com/saic-mdal/lama.git
$ cd lama

# Proxy配下の場合は下記
$ export http_proxy="http://"username":"password"@proxy:port"
$ export https_proxy="http://"username":"password"@proxy:port"

# conda環境構築
$ conda env create -f conda_env.yml

# conda環境をactivate
$ conda activate lama
# pytorchのインストール
# https://pytorch.org/ で適するバージョンをインストール
$ conda install pytorch torchvision torchaudio cudatoolkit=10.2 -c pytorch -y
$ pip install pytorch-lightning==1.2.9

以上で、セットアップは完了です。
conda環境をymlに出力してくれているので非常に親切なプロジェクトですね。

画像修復

次に、セットアップした環境を利用して画像修復を実施します。

$ cd lama

# 学習済みモデルのダウンロードに必要なツールをダウンロード
$ pip3 install wldhx.yadisk-direct

# 学習済みモデルのダウンロード
$ curl -L $(yadisk-direct https://disk.yandex.ru/d/ouP6l8VJ0HpMZg) -o big-lama.zip
$ unzip big-lama.zip

# テスト用画像とマスク画像のダウンロード
$ curl -L $(yadisk-direct https://disk.yandex.ru/d/xKQJZeVRk5vLlQ) -o LaMa_test_images.zip
$ unzip LaMa_test_images.zip

以上で学習済みモデルとテスト用画像、マスク画像が用意されました。
下記のディレクトリ構造になっているかご確認ください。

$ /lama$ tree -L 1
.
├── LICENSE
├── LaMa_test_images		←ダウンロードしたテスト用画像とマスク画像
├── LaMa_test_images.zip
├── README.md
├── big-lama			←ダウンロードした学習済みモデル
├── big-lama.zip
├── bin
├── colab
├── conda_env.yml
├── configs
├── docker
├── fetch_data
├── models
├── requirements.txt
└── saicinpainting

最後に、学習済みモデルを使って画像修復を実施します。

$ cd lama

# 出力ディレクトリ作成
$ mkdir output

# 画像修復実施
$ python3 bin/predict.py model.path=$(pwd)/big-lama indir=$(pwd)/LaMa_test_images outdir=$(pwd)/output

outputに予測結果となる画像が出力されます。

ぼんやりと人がいた部分にもやを感じますが、マスク画像に沿って人が削除された後、綺麗に画像が修復されています。

入力画像の作り方

先ほどの画像修復の結果をご覧いただくとわかりますが、入力画像にはmask画像が必要です。
以下に、mask画像の生成方法をご紹介します。

今回はパワーポイントを使って下記の画像のマスク画像を作っていきます。

対象画像をパワーポイントに貼り付け図形の「フリーフォーム:図形」を選択しマスクしたい箇所を囲みます

その後マスク対象の図形を白く塗りつぶします。
最後に、「図の書式設定」→「図の修正」→「明るさ:-100%、コントラスト:100%」にすると以下のような画像に仕上がります。

マスク画像生成時の注意点は以下の2点です。

1. マスク範囲は大きめに取る
   マスクからはみ出た部分は削除されないため、残った障害物を基準に画像修復されます
2. 画像のサイズは合わせる
   画像サイズが合っていないとマスクの範囲がずれます

今回は以下のような入力画像となります

この2つの画像を、「original001.png」、「original001_mask.png」のように、マスク画像のファイル名を「対象ファイル名_mask」に設定した後、先ほどの動作環境に配置して画像修正を実行します。

すると、自作のマスク画像を使って画像から障害物を削除できます。

このように、マスク画像さえ作成してしまえば、任意の画像の任意の対象を画像から削除できます。

トラブルシューティング

ModuleNotFoundError
Traceback (most recent call last): File "bin/predict.py", line 14, in <module> from saicinpainting.evaluation.utils import move_to_device ModuleNotFoundError: No module named 'saicinpainting'
saicinpaintingが読み込めていないことが原因 /lama/bin/predict.pyを修正

import logging
import os
import sys
import traceback

# 下記3行を追加して一つ上のディレクトリを読み込めるように変更
__dir__ = os.path.dirname(os.path.abspath(__file__))
sys.path.append(__dir__)
sys.path.append(os.path.abspath(os.path.join(__dir__, '../')))

from saicinpainting.evaluation.utils import move_to_device

TypeError: load() missing 1 required positional argument: 'Loader'
インストールしたライブラリdaskのエラー
やや強引ではあるが下記修正 lama/lib/python3.6/site-packages/dask/config.py

with open(path) as f:
    # 下記をsafe_loadに変更
    # data = yaml.load(f.read()) or {}
    data = yaml.safe_load(f.read()) or {}

RuntimeError: The NVIDIA driver on your system is too old
RuntimeError: The NVIDIA driver on your system is too old (found version 10010). Please update your GPU driver by downloading and installing a new version from the URL: http://www.nvidia.com/Download/index.aspx Alternatively, go to: https://pytorch.org to install a PyTorch version that has been compiled with your version of the CUDA driver.
pytorch、cudatoolkitのバージョン不一致
cudatoolkitのバージョンを上げるか、pytorchのバージョンを下げる
以下は、pytorchのバージョンを下げる例

# 一度アンインストール
$ conda uninstall pytorch torchaudio cudatoolkit pytorch-mutex

# CUDAバージョン確認
$ nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 418.87.00    Driver Version: 418.87.00    CUDA Version: 10.1     |
|-------------------------------+----------------------+----------------------+

# https://pytorch.org/get-started/previous-versions/ でバージョン確認
$ conda install pytorch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2 cudatoolkit=10.1 -c pytorch

AttributeError: 'builtin_function_or_method' object has no attribute 'rfftn'
名称のconflictが原因
/lama/saicinpainting/training/modules/ffc.pyにrfftnのモジュールを明示的にimportすることで解消

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
# 以下を追加
import torch.fft

まとめ

本記事では、Lamaを使った画像修復を紹介しました。
画像処理系の機械学習は出力結果が目で見てわかるものが多く知的好奇心をくすぐりますね。

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

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

参考文献

1. Resolution-robust Large Mask Inpainting with Fourier Convolutions

2. saic-mdal/lama