OpenCV
DeepLearning
python3
Chainer

画像中の人だけを削除する学習モデル

概要

画像中の「人」だけを削除する。モンスターや動物は削除されないし、削除した跡は背景をうまく補完する。ソースコードは以下にあるので自己責任でどうぞ。
https://github.com/ka10ryu1/human_delete

学習結果

12枚のそれぞれの画像において、左側が入力画像(人が写っている)で、右側が出力画像(左側にいた人が消えている)。消えた跡は完璧ではないが背景が補完されているのがわかる。

concat.jpg

動作環境

  • Ubuntu 16.04.3 LTS
  • Python 3.5.2
  • chainer 3.5
  • numpy 1.14.2
  • cupy 2.4
  • opencv-python 3.4.0.12

学習済みモデルを使って実行させる

predict.pyを利用すると上記の画像のような結果を得られる。CPU実行でも数秒とかからずに処理が完了する。

human_delete
$ ./predict.py Model/*

一度にたくさん実行したいときはpredict_repeat.pyを利用する。

human_delete
$ ./predict_repeat.py Model/*

入力画像

以下のような画像データセットを準備する。上段が入力画像で下段が正解画像である。

npz2jpg.jpg

このデータセットはGitHubにアップしたソースコードのcreate_dataset.pyを実行することで作成できる。画像はいらすとやから拝借した。画像を重ね合わせる方法については「OpenCVで透過画像を任意の画像に張り付ける」を参照されたい。

human_delete
$ ./create_dataset.py

データセットの中身を確認するためにはTools/npz2jpg.pyを実行すればいい。

human_delete
$ Tools/npz2jpg.py result/test_256x256_000020.npz 

学習

学習を実行するためにはtrain.pyを実行する必要がある。

human_delete
$ ./train.py -g 0

-g 0はGPU(ID:0)を使用するためのオプション引数。ユニット数を増やしたり(-u 32)、エポック数を増やしたり(-e 100)すると削除性能が向上するが、トップの画像のような学習をさせるにはcreate_dataset.pyで生成する画像をもっと増やす必要がある。

ネットワーク層

ネットワーク層はDeep Image Priorをベースに作成している。Deep Image Priorについてはここが詳しい。

Deep Image Priorは以下に示すように、DownSampleBlockUpSampleBlockに分かれている。

DownSampleBlock
class DownSampleBlock(Chain):
    def __init__(self, n_unit, ksize, stride, pad,
                 actfun=None, dropout=0, wd=0.02):

        super(DownSampleBlock, self).__init__()
        with self.init_scope():
            self.cnv = L.Convolution2D(
                None, n_unit, ksize=ksize, stride=stride, pad=pad, initialW=I.Normal(wd)
            )
            self.brn = L.BatchRenormalization(n_unit)

        self.actfun = actfun
        self.dropout_ratio = dropout

    def __call__(self, x):
        h = self.actfun(self.brn(self.cnv(x)))
        if self.dropout_ratio > 0:
            h = F.dropout(h, self.dropout_ratio)

        return h
UpSampleBlock
class UpSampleBlock(Chain):
    def __init__(self, n_unit1, n_unit2, ksize, stride, pad,
                 actfun=None, dropout=0.0, wd=0.02, rate=2):

        super(UpSampleBlock, self).__init__()
        with self.init_scope():
            self.cnv = L.Convolution2D(
                None, n_unit1, ksize=ksize, stride=stride, pad=pad, initialW=I.Normal(wd)
            )
            self.brn = L.BatchRenormalization(n_unit2)

        self.actfun = actfun
        self.dropout_ratio = dropout
        self.rate = rate

    def __call__(self, x):
        h = self.actfun(self.brn(self.PS(self.cnv(x))))
        if self.dropout_ratio > 0:
            h = F.dropout(h, self.dropout_ratio)

        return h

    def PS(self, h):
        """
        "P"ixcel"S"huffler
        Deconvolutionの高速版
        """

        batchsize, in_ch, in_h, in_w = h.shape
        out_ch = int(in_ch / (self.rate ** 2))
        out_h = in_h * self.rate
        out_w = in_w * self.rate
        out = F.reshape(h, (batchsize, self.rate, self.rate, out_ch, in_h, in_w))
        out = F.transpose(out, (0, 3, 4, 1, 5, 2))
        out = F.reshape(out, (batchsize, out_ch, out_h, out_w))
        return out

この二つのブロックを以下のようにして組み合わせてネットワーク層を構成している。

class JC_DDUU(Chain):
    def __init__(self, n_unit=128, n_out=1, rate=4,
                 actfun1=F.relu, actfun2=F.sigmoid,
                 dropout=0.0, view=False):
        """
        [in] n_unit:    中間層のユニット数
        [in] n_out:     出力チャンネル
        [in] actfun1: 活性化関数(Layer A用)
        [in] actfun2: 活性化関数(Layer B用)
        """

        unit1 = n_unit
        unit2 = n_unit * 2
        unit4 = n_unit * 4
        unit8 = n_unit * 8
        nout = (rate**2) * 3

        super(JC_DDUU, self).__init__()
        with self.init_scope():
            # D: n_unit, ksize, stride, pad,
            #    actfun=None, dropout=0, wd=0.02
            self.d1 = DownSampleBlock(unit1, 5, 2, 2, actfun1, dropout)
            self.d2 = DownSampleBlock(unit2, 5, 2, 2, actfun1, dropout)
            self.d3 = DownSampleBlock(unit4, 5, 2, 2, actfun1, dropout)
            self.d4 = DownSampleBlock(unit8, 5, 2, 2, actfun1, dropout)
            self.d5 = DownSampleBlock(unit8, 3, 1, 1, actfun1, dropout)

            # U: n_unit1, n_unit2, ksize, stride, pad,
            #    actfun=None, dropout=0, wd=0.02, rate=2
            self.u1 = UpSampleBlock(unit4, unit1, 5, 1, 2, actfun2, dropout)
            self.u2 = UpSampleBlock(unit4, unit1, 5, 1, 2, actfun2, dropout)
            self.u3 = UpSampleBlock(unit4, unit1, 5, 1, 2, actfun2, dropout)
            self.u4 = UpSampleBlock(unit4, unit1, 5, 1, 2, actfun2, dropout)
            self.u5 = UpSampleBlock(nout, 3, 5, 1, 2, actfun2, 0, 0.02, rate)

        self.view = view
        self.cnt = 0
        self.timer = 0

        print('[Network info]', self.__class__.__name__)
        print('  Unit:\t{0}\n  Out:\t{1}\n  Drop out:\t{2}\nAct Func:\t{3}, {4}'.format(
            n_unit, n_out, dropout, actfun1.__name__, actfun2.__name__)
        )

    def block(self, f, x):
        if self.view:
            print('{0:2}: {1}\t{2:5.3f} s\t{3} '.format(
                self.cnt, f.__class__.__name__, time.time()-self.timer, x.shape))
            self.cnt += 1

        return f(x)

    def __call__(self, x):
        if self.view:
            self.timer = time.time()

        hc = []
        ha = self.block(self.d1, x)
        hb = self.block(self.d2, ha)
        hc = self.block(self.d3, hb)
        hd = self.block(self.d4, hc)
        he = self.block(self.d5, hd)

        h = self.block(self.u1, F.concat([hd, he]))
        h = self.block(self.u2, F.concat([hc, h]))
        h = self.block(self.u3, F.concat([hb, h]))
        h = self.block(self.u4, F.concat([ha, h]))
        y = self.block(self.u5, h)

        if self.view:
            print('Output {0:5.3f} s: {1}'.format(time.time()-self.timer, y.shape))
            exit()
        else:
            return y

学習実行部

train.pyは基本的にchainerのtrain_mnist.py通りだが、以下の部分が少し異なる。

データ取得

predict.py
train, test = GET.imgData(args.in_path)
train = ResizeImgDataset(train, args.shuffle_rate)
test = ResizeImgDataset(test, args.shuffle_rate)

args.in_pathには学習データとテストデータのあるフォルダが格納されており、GET.imgData()の返り値に学習データとテストデータが代入される。

ResizeImgDataset()chainer.dataset.DatasetMixinを親クラスとしているため、trainer.run()中に学習データとテストデータからミニバッチ分だけデータを取得する直前にこれが呼ばれて実行される。ネットワーク層を通したデータは画像サイズが大きくなるように設計されているため、正解データもResizeImgDataset()によって画像サイズを拡大することが目的である。

毎回拡大処理を実行するよりも、学習前にまとめて実行した方が計算効率がいいが、同時に画像データ拡大によるメモリ消費も大きくなるため。このようにしている。

学習時パラメータ保存

predict.py
model_param = {i: getattr(args, i) for i in dir(args) if not '_' in i[0]}
model_param['shape'] = train[0][0].shape

predict.pyなどでは学習時のパラメータが必要になるのでここでmodel_paramに格納している。格納しているのはすべてのオプション引数の値と入力画像の画像サイズ等である。ファイルは拡張子が.jsonのファイルである。

Topの学習時に生成されたjsonファイルは以下の通りである。

param.json
{
    "actfun1": "relu",
    "actfun2": "sigmoid",
    "batchsize": 8,
    "dropout": 0.2,
    "epoch": 500,
    "frequency": 100,
    "gpu_id": 0,
    "in_path": "./result/",
    "lossfun": "mse",
    "only_check": false,
    "optimizer": "adam",
    "out_path": "result/02",
    "plot": true,
    "resume": "result/01/13lkh5jo_300.snapshot",
    "shape": [
        3,
        256,
        256
    ],
    "shuffle_rate": 4,
    "unit": 64
}

以上。