【AIアプリ開発 Day1】画像内の人間をセグメンテーションして、モザイクパターンで埋める

アプリ開発記録です。前回はどんなアプリを作るか決めました。

【AIアプリ開発 Day0】どんなアプリを作るか考える

今回は、核となる処理が実現できるか見ていこうと思います。

おしながき

・COCOデータセット学習済みモデルを用いて、画像内の Person をセグメンテーション

・Person の箇所だけ除去

・除去した箇所を、モザイクパターン画像で埋める

Before

After

それではやっていこうと思います。

COCOデータセットでの学習済みモデルを用いて、セグメンテーション

Person クラスが分類ラベルに含まれている代表的なデータセットでいうと、COCOデータセットがあります。（Person は代表的なデータセットにはだいたい含まれているような気もしますが・・）

ですので、COCOで学習された重みを用いれば、わざわざ再学習せずとも人間をセグメンテーションすることが簡単にできるはずです。

torchvision に組み込みのモデルが COCO データセットで Pascal VOC の 21クラス（中に Person 含む）を学習させたものでしたので、今回はそれを使わせてもらうことにしましょう。

まずは画像を読み込みます。

import torch
import torchvision
from torchvision import models
from torchvision import transforms

import numpy as np
from PIL import Image

# 画像の読み込み
sample1 = Image.open("data/sample1.jpg")

import torch

import torchvision

from torchvision import models

from torchvision import transforms

import numpy as np

from PIL import Image

# 画像の読み込み

sample1 = Image.open("data/sample1.jpg")

前処理を行います。この画像が 4016 x 6016 の高解像度画像だったので、リサイズして小さくします。（高解像度画像の対応は、ミニマムにいったん一通りできてから考えます。）

transform = transforms.Compose([
    transforms.Resize((800, 1200)),
    transforms.ToTensor()
    ])

inputs = transform(sample1)
inputs = inputs.unsqueeze(0)

transform = transforms.Compose([

transforms.Resize((800, 1200)),

transforms.ToTensor()

])

inputs = transform(sample1)

inputs = inputs.unsqueeze(0)

準備ができたのでモデルを用意して推論。

#推論
model = models.segmentation.deeplabv3_resnet101(pretrained=True, num_classes=21)
model.eval()

pred = model(inputs)

# 後処理
mask = torch.argmax(pred["out"][0], 0)

print(mask.size())

"""
torch.Size([800, 1200])
"""

print(mask.unique())

"""
tensor([ 0, 15])
"""こ

#推論

model = models.segmentation.deeplabv3_resnet101(pretrained=True, num_classes=21)

model.eval()

pred = model(inputs)

# 後処理

mask = torch.argmax(pred["out"][0], 0)

print(mask.size())

"""

torch.Size([800, 1200])

"""

print(mask.unique())

"""

tensor([ 0, 15])

"""こ

マスク画像を獲得することができました。今回はラベル 15 が Person です。

Person の箇所だけ除去

では、マスク画像の情報をもとに、元画像から Person の箇所だけ除去します。numpy で処理します。

# 元画像を再度読み込みリサイズ
img = Image.open("data/sample1.jpg")
img = np.array(img.resize((1200, 800)))

mask = np.stack([mask, mask, mask], 2)
print(mask.shape)

"""
(800, 1200, 3)
"""

# 元画像を再度読み込みリサイズ

img = Image.open("data/sample1.jpg")

img = np.array(img.resize((1200, 800)))

mask = np.stack([mask, mask, mask], 2)

print(mask.shape)

"""

(800, 1200, 3)

"""

元画像の中で、mask 画像のラベルが 15 出ない位置を、輝度 0 で埋めます。

import matplotlib.pyplot as plt
%matplotlib inline

result = np.where(mask != 15, img , 0)
plt.imshow(result)

import matplotlib.pyplot as plt

%matplotlib inline

result = np.where(mask != 15, img , 0)

plt.imshow(result)

除去した箇所を、モザイクパターンで埋める

では、最後に先程できた画像で輝度が 0 の箇所を、任意のモザイク画像で埋めていきます。今回は木目のようなパターンでやってみます。

mosaic = Image.open("mosaic.png")
mosaic = np.array(mosaic.resize((1200, 800)))

print(mosaic.shape)

"""
(800, 1200, 3)
"""

output = np.where(result == 0, mosaic, result)
plt.imshow(out)