本記事では、機械学習手法SpA-Formerを用いて写真に映り込んだ影を除去する方法をご紹介します。

SpA-Former

概要

SpA-Formerは単一の影の映り込んだ画像(shaded image)から影のない画像(shadow-free image)を生成するEnd to Endの技術です。

従来手法は、影の検出、影の除去の2stageの処理をそれぞれの個別のモデルを用いて実現していました。SpA-Formerでは、この2stageを統合し、一つのモデルで実現しています。
Fourier transform residual networkと、two-wheel joint spatial attentionから構成される本手法は、End to Endであり、非常に高速な処理効率を実現しています。

詳細はこちらの論文をご参照ください。

本記事では上記手法を用いて、写真に映り込んだ影を除去していきます。

デモ(Colaboratory)

それでは、実際に動かしながら影を除去していきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

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

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

環境セットアップ

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

初めにGithubからソースコードを取得します。

%cd /content

!git clone https://github.com/zhangbaijin/SpA-Former-shadow-removal.git

次にライブラリをインストールします。

%cd /content/SpA-Former-shadow-removal

!pip install gdown==4.5.1
!pip install einops

最後にライブラリをインポートします。

%cd /content/SpA-Former-shadow-removal

import os
import glob

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

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

次に、論文発表元が公開する学習済みモデルをダウンロードします。

%cd /content/SpA-Former-shadow-removal
!mkdir -p ./checkpoint

if not os.path.exists('checkpoint/gen_model_epoch_160.pth'):
  !gdown https://drive.google.com/uc?id=1gLOu4jkqslu_fWpQUd0N4hcNz4qCD8Xc -O checkpoint/gen_model_epoch_160.pth
if not os.path.exists('checkpoint/dis_model_epoch_160.pth'):
  !gdown https://drive.google.com/uc?id=1AcJSAV4oHiYneYUxCwaV0M_W1QU4MdZX -O checkpoint/dis_model_epoch_160.pth

Google Colaboratoryのストレージに学習済みモデルがダウンロードされます。

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

続いて影を除去する画像を用意します。

%cd /content/SpA-Former-shadow-removal
!rm -rf ./test_imgs
!mkdir -p ./test_imgs
!mkdir -p ./datasets

!wget -c https://cdn.pixabay.com/photo/2016/04/11/03/00/run-1321278_960_720.jpg \
      -O ./test_imgs/test_01.jpg

!wget -c https://cdn.pixabay.com/photo/2013/10/20/22/05/shadow-198682_960_720.jpg \
      -O ./test_imgs/test_02.jpg

!wget -c https://cdn.pixabay.com/photo/2015/09/09/20/33/travel-933171_960_720.jpg \
      -O ./test_imgs/test_03.jpg

!wget -c https://cdn.pixabay.com/photo/2021/08/02/18/21/stairs-6517488_960_720.jpg \
      -O ./test_imgs/test_04.jpg

以下はオプションですが、Google Driveにて公開されているISTDデータセットのショートカットを自身のGoogle Driveに追加し、Google Colaboratoryから参照することでISTDのテストデータを使用することも可能です。

from google.colab import drive
drive.mount('/content/drive')

# https://drive.google.com/file/d/1I0qw-65KBA6np8vIZzO6oeiOvcDBttAY/view
# 上記ファイルのショートカットを自身のGoogle Driveに追加した場合のみ以下を実行

!unrar x /content/drive/MyDrive/ISTD_Dataset.rar /content/SpA-Former-shadow-removal/datasets > /dev/null

Shadow Remove

それでは、セットアップした画像を使って影を除去していきます。

%cd /content/SpA-Former-shadow-removal

import numpy as np
import argparse
from cv2 import cv2
import matplotlib.pyplot as plt
import time

import torch
from torch.autograd import Variable

from utils import gpu_manage, heatmap
from SpA_Former import Generator
from google.colab.patches import cv2_imshow

def predict(args):

    gpu_manage(args)
    ### MODELS LOAD ###
    print('===> Loading models')

    gen = Generator(gpu_ids=args.gpu_ids)

    param = torch.load(args.pretrained)
    gen.load_state_dict(param)

    if args.cuda:
        gen = gen.cuda(0)

    print ('<=== Model loaded')

    print('===> Loading test image')
    img = cv2.imread(args.test_filepath, 1).astype(np.float32)
    img = img / 255
    img = img.transpose(2, 0, 1)
    img = img[None]
    print ('<=== test image loaded')

    with torch.no_grad():
        x = torch.from_numpy(img)
        if args.cuda:
            x = x.cuda()
        
        print('===> Removing the cloud...')
        start_time = time.time()
        att, out = gen(x)
        print('<=== finish! %.3fs cost.' % (time.time()-start_time))

        x_ = x.cpu().numpy()[0]
        x_rgb = x_ * 255
        x_rgb = x_rgb.transpose(1, 2, 0).astype('uint8')
        out_ = out.cpu().numpy()[0]
        out_rgb = np.clip(out_[:3], 0, 1) * 255
        out_rgb = out_rgb.transpose(1, 2, 0).astype('uint8')
        att_ = att.cpu().numpy()[0] * 255
        att_heatmap = heatmap(att_.astype('uint8'))[0]
        att_heatmap = att_heatmap.transpose(1, 2, 0)

        allim = np.hstack((x_rgb, out_rgb, att_heatmap))

        cv2_imshow(allim)

images = glob.glob('./test_imgs/*.jpg')        

args = argparse.ArgumentParser()
args.pretrained = './checkpoint/gen_model_epoch_160.pth'
args.cuda = True
args.gpu_ids = [0]
args.manualSeed = 12

for img in images:
  args.test_filepath = img
  predict(args)

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

ヒートマップを見ると影の検出は良好ですが、影の除去の精度はいまいちです。
この点追加で学習させる必要がありそうです。

続いて、ISTDデータセットのテストデータを試してみます。

import random

datasets = glob.glob('/content/SpA-Former-shadow-removal/datasets/ISTD_Dataset/test/test_A/*.png')

tests = random.sample(datasets, 5)

for img in tests:
  args.test_filepath = img
  predict(args)

先ほどと異なり影の除去もうまくいっています。
この点、学習データに外壁に映り込んだ影や、コンクリートに映り込んだ影が含まれていたなどが考えられます。様々な場所に映る影を除去するためには学習データの拡充が必要そうです。

まとめ

本記事では、SpA-Formerを用いて写真に映り込んだ影を除去する方法をご紹介しました。
実際に動かしてみることで理解が進みますね。

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

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

参考文献

1.  論文 - SpA-Former: Transformer image shadow detection and removal via spatial attention

2. GitHub - zhangbaijin/SpA-Former-shadow-removal