[cycleGAN]ゾンビランドサガの水野愛を生き返らせる #Python

ゾンビランドサガというアニメの水野愛(以下、愛ちゃん)がクッソ可愛い。
ゾンビの姿の愛ちゃんも可愛らしいですが、やっぱり生前の愛ちゃんの方が可愛い。
そこで、cycleGANを使って、ゾンビ愛ちゃんと生前愛ちゃんのデータセットから変換してみようと思います。

docker + pytorchの作成済みモデルを利用してお手軽に実装します。
また、ローカルにnvidia-dockerの環境が構築されている前提です。
データに関しては、１話〜６話までの画像から愛ちゃんを手作業で検出し、データセットを作成しました。
データセットの公開予定はありません。

画像数

ep	ゾンビ	生前
opening	3	84
1	34	369
2	312	503
3	261	643
4	97	490
5	104	455
6	115	398
Sum	926	2942

手順

1. pytorch/pytorchのイメージからコンテナを作成

docker pull pytorch/pytorch
nvidia-docker run -it --name aichanmajiangel -p 8097:8097 pytorch/pytorch

GPUを使うためnvidia-dockerを使う。
visdom.serverがデフォルトで8097に立つのでポートを開放しておく。
再接続する場合はnvidia-docker exec -it aichanmajiangel bash

2. pytorch-CycleGAN-and-pix2pixをcloneする

git clone https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix.git
cd pytorch-CycleGAN-and-pix2pix
pip install -r requirements.txt

3. データセットの用意

a. 気合と根気で集めます。

import kiai
import konki
from os import path
from PIL import Image

images_path = 'aichan_image/'
for i in range(konki.max()):
    image = kiai.get('水野愛')
    if zombie:
        image.save(path.join(images_path,
                     'zombie',
                     'zombie_{}.png'.format(str(i).zfill(5))))
    else:
        image.save(path.join(images_path,
                     'aichan',
                     'aichan_{}.png'.format(str(i).zfill(5))))

b. trainとtestにスプリットする。

import os
import shutil
import random

images_path = 'aichan_image/'
zombie = os.listdir(path.join(images_path, 'zombie'))
aichan = os.listdir(path.join(images_path, 'aichan'))

# random shuffle
zombie = random.sample(zombie, len(zombie))
aichan = random.sample(aichan, len(aichan))
# zombe[0]   -> 'xxxx.png'
# len(zombi) ->  917

# split 9:1
trainA = zombie[:int(len(zombie)/10*9)]
testA = zombie[int(len(zombie)/10*9):]
trainB = aichan[:int(len(aichan)/10*9)]
testB = aichan[int(len(aichan)/10*9):]

# create dataset
if not os.path.exists('dataset'):
    os.makedirs('dataset')
    os.makedirs('dataset/trainA')
    os.makedirs('dataset/trainB')
    os.makedirs('dataset/testA')
    os.makedirs('dataset/testB')

for c, n in zip([trainA, testA, trainB, testB], ['trainA', 'testA', 'trainB', 'testB']):
    for i,image in enumerate(c):
        shutil.copy(image, path.join('dataset',n))