Pythonによる画像背景除去方法(DeepLearning)

はじめに

前回ではopencCVを用いて二値化を行い、背景と前景を分離して背景の除去を行いました。 しかし、グレースケールにおける背景と前景のヒストグラムの重なりにより画像によってはうまく抽出できないケースがありました。

今回はDeepLearningのセマンティックセグメンテーション(Semantic Segmentation) を用いて背景除去を行います。 前景と背景を正確に推定するプロセスはイメージマッティング(Image Matting)とも呼ばれています。 DNNにより効果的に特徴抽出を行うことができるため、前回の手法よりは良い結果が期待できます。 Pytorchライブラリにはセマンティックセグメンテーションを使用するのに有効な学習済みのモデルが搭載されており、 気軽に使用することが可能です。使用する学習モデルはDeepLabV3 with ResNet-101です。 またDeepLabV3モデルで得られた結果をさらに向上させるためにIndexNet Mattingを使用します。 IndexNet MattingはDeepLabV3のデコーダーにおける境界の詳細描画の弱点を改善するものとなっている。

DeepLabV3 ResNet-101

DeepLabV3モデルは入力画像を特徴マップにエンコードする部分と、特徴マップをセグメンテーションに分割した画像にデコードする部分からなります。 エンコード部分にはAtrous畳み込みという各解像度・視野ごとにより効果的に特徴を抽出しすることが可能な手法で高速化にも貢献しています。 Pytorchの学習モデルはPASCAL VOC2012のデータセットで学習したモデルであるため、 そのため、識別可能なオブジェクトは以下の通りとなります。

Person: person
Animal: bird, cat, cow, dog, horse, sheep
Vehicle: aeroplane, bicycle, boat, bus, car, motorbike, train
Indoor: bottle, chair, dining table, potted plant, sofa, tv/monitor

これ以外のオブジェクトを識別したい場合は転移学習によるfinetuneが必要になります。

DeepLabV3による背景除去

以下のソースコードを作成して実行し、背景除去を行いました。 実行にはGoogle Colaboratoryを使用しました。

    %cd /content
import numpy as np
import cv2
import matplotlib.pyplot as plt
import torch
from torchvision import models
import torchvision.transforms as T

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = models.segmentation.deeplabv3_resnet101(pretrained=1).to(device).eval()
def segmentation(model, img, dev):
    h, w = img.shape[:2]
    trf = T.Compose([T.ToPILImage(),
                     T.Resize(800),
                     T.ToTensor(),
                     T.Normalize(mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225])])
    inpt = trf(img).unsqueeze(0).to(dev)
    out = model.to(dev)(inpt)['out']
    mask = torch.argmax(out.squeeze(), dim=0).detach().byte().cpu().numpy()
    mask = cv2.resize(mask, (w,h))
    return mask
    

画像をアップロードして、処理結果を出力します。

    from google.colab import files
upload_file = files.upload()
file_name = next(iter(upload_file))

img = cv2.imread(file_name)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

mask = segmentation(model, img, device)
mask = np.where(mask == 0, 0, 255).astype(np.uint8)

img = cv2.cvtColor(img, cv2.COLOR_RGB2BGRA)
img_rmv = img.copy()
mask2 = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGRA)
img_rmv[:, :, 3] = np.where(np.all(mask2 == [0, 0, 0, 255], axis=-1), 0, 255)

img_cat = cv2.hconcat([img, mask2, img_rmv])
#cv2.imwrite('./result.png', img_cat)
cv2.imwrite('./result.png', img_rmv)
    

結果は以下の画像になります。 前回使用した自作の小鳥画像はbirdとは判定されませんでした。

きちんとした本物の鳥画像で試すためpixabay より画像をダウンロードして使用しました。 結果は以下になります。この画像ではうまくいきました。

比較用に前回のopenCVとの比較を行います。

比較するとopenCVでは除去しきれなかった複雑な背景をうまく除去できていました。 前回の手法よりパフォーマンスがよいことがわかりました。 ただよく見ると少し背景が残ってしまっています。 これをIndexNet Mattingを使用して、さらにきれいに背景を除去しました。

IndexNet Matting

デコーダーにおいて境界の描画に重要なのはインデックスであり、そのインデックスに境界の位置を投影して特徴マップを構築することで 境界の描画の改善が行われます。IndexNetでは局所特徴マップの関数としてインデックスをモデル化しインデックス関数を学習します。 学習されたインデックス関数によって境界と領域の特徴を同時に抽出します。

このモデルを使用するには、cv2.erodeとcv2.dilateを用いて、前景となる部分・背景となる部分・境界部分に分割したtrimap画像を作成します。

    kernel = np.ones((3,3),np.uint8)
sure_bg = cv2.dilate(mask, kernel, iterations=2)
sure_fg = cv2.erode(mask, kernel, iterations=2)
trimap = np.full(mask.shape, 128)
trimap[sure_fg == 255] = 255
trimap[sure_bg == 0] = 0
cv2.imwrite('./bird_trimap.png', trimap)
    

後は、IndexNet Mattingをgithubからクローンしてきてdemo.pyを実行すれば、mattesディレクトリにmask画像が出力されます

    import os
if not os.path.exists('indexnet_matting'):
    !git clone https://github.com/poppinace/indexnet_matting
%cd indexnet_matting/
!python scripts/demo.py
%cd /content
    

最後にmaskを用いて背景部分に透過処理を行えばよいです。

    file1 = "./indexnet_matting/examples/images/bird.jpg"
file2 = "./indexnet_matting/examples/mattes/bird.jpg"


img = cv2.imread(file1)
img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
mask = cv2.imread(file2)
mask = cv2.cvtColor(mask, cv2.COLOR_BGR2BGRA)

img = img.astype(np.float32)
mask = mask.astype(np.float32)/255.
mask[:,:,3] = mask[:,:,0]

out = img * mask

img_cat = cv2.hconcat([img, mask*255., out])
cv2.imwrite('./result_indexnet_matting.png', img_cat)
    

まとめ

DeepLabV3で背景と前景の識別を行い、IndexNet Mattingで境界のガイドをしてやれば 容易に精度の高い背景除去が可能になります。 今回はPascal VOCの学習モデルを使用しているので識別可能なオブジェクトの種類は限られますので 別のオブジェクトを使いたいときには転移学習でfinetuneを行う必要があります。 転移学習や別のモデルも別の機会に紹介できたらと思います。

Reference