Edited at

CycleGANに自分のデータセットを用いて学習させる

More than 1 year has passed since last update.


0. 概要

ドメインの異なる画像を学習できるCycleGANを動かす。

このネットワークは何ができるかというと、「馬」と「しまうま」の画像両方を学習させると、どうしたら互いの画像を相互変換可能かを学習でき、最終的に馬の画像を入れると形はそのままで色合いをしまうまに変換してしまうなんてことができるそう。

他にも線分のみのキャラクターを入れると、2次元のかわいらしい絵に変換してくれたりだとか、使い道は多そう!

image.png

このようにテクスチャーの変更が得意なネットワークです。一方、形の変換までは苦手だそう。

以下、実装

https://github.com/xhujoy/CycleGAN-tensorflow


1. 動作環境

GeForce GTX 1080 Ti

TensorFlow 1.8
CUDA 9.0
CUDNN 7.5
Python 3.6


2. デフォルトで動かしてみる

学習:

まずはソースコードを持ってくる

git clone https://github.com/xhujoy/CycleGAN-tensorflow

cd CycleGAN-tensorflow

次に学習データを持ってくる。とりあえず王道の馬としまうまのデータセット。

bash ./download_dataset.sh horse2zebra

学習開始:

python main.py --dataset_dir=horse2zebra

使うGPUを選択するときは:

$ nvidia-smi -L

→ GPU一覧を確認

# GPUを指定
export CUDA_VISIBLE_DEVICES="0"

学習の可視化

tensorboard --logdir=./logs

ネットワーク構造が表示された後に、スムーズに学習が始まっていった。

image.png

テスト:

まずは馬からしまうまへ。

python main.py --dataset_dir=horse2zebra --phase=test --which_direction=AtoB

--which_direction=AtoBこいつをBtoAにしてあげれば逆も可能。

実行したらmain.py直下のフォルダtest/に結果が格納される。

これが入力画像

n02381460_20.jpg

そしてこちらが出力画像

AtoB_n02381460_20.jpg

3分程度学習させた程度なので、上手く変換できていない・・・


3. オリジナル画像で学習させる

まずデータ構造を理解する。

学習データ及びテストデータはmain.py直下のフォルダdatasets/に保存する。

datasets/以下に例えばcarというフォルダを作る。

内部のフォルダ構成は以下のようになる。

*datasets

+ horse2zebra
+ car
- trainA
- trainB
- testA
- testB

ここで各学習データには以下を入れる。

trainAには普通乗用車の画像を入れる

trainBにはタクシーの画像を入れる

testAには普通乗用車のテスト画像を入れる

testBにはタクシーのテスト画像を入れる

このように画像を配置したら以下で学習を実行できる。

python main.py --dataset_dir=car

ものすごく簡単である。

なお、細かいが画像サイズは256 x 256が初期設定である。

このため、サイズを変更する場合はソースコードを変えましょう。

以上!