【DCGAN】AIはピクセルアートを描けるようになるのか

管理人

様々な生成AIがありますが、その中から今回は絵を生成できる「DCGAN」を実装してピクセルアートを描かせてみたいと思います。

DCGAN (Deep Convolutional GAN)とは

DCGAN (Deep Convolutional GAN)はUnsupervised Representation Learning with Deep Convolutional Generative Adversarial Networksという論文で提案された生成モデルです。

GANの改良を行なっているものではあり、基本的な考え方はGANと同じです。やっていることは、GeneratiorとDiscriminatorを競わせながらどんどん画像のクオリティをあげるという考え方です。つまり、贋作作家が贋作を作り鑑定士が贋作を見抜く、作家が鑑定士を騙せるような贋作をさらに作る、それを鑑定士がさらに見抜くといった、いたちごっこを続けることでより本物に近づいた贋作ができるようになる贋作作家AIが出来上がるという感じです。

全体像

あくまで実装して、実際にどうなったのかについての結果に早めにいきたいのでこの辺りの構造や特徴については詳しく書きません。簡単にどうっていうのだけ書いておきます。

全体像としては、GANと全く同じになってます。オリジナルのGANとの大きな違いは、それぞれのネットワークで全結合層ではなく、畳み込み層及び転置畳み込み層を使用しているところです。

GANの全体像図

Generator

Generatorのネットワークでは以下の図のように、入力となる100次元のランダムなノイズから64×64サイズの画像を生成します。

Discriminator

画像識別の畳み込みニューラルネットワークの構造をしている。Generatorの逆のようなネットワークになっている。

学習の準備

画像の決定と収集

今回はピクセルアートを作っていきたいと思います。ピクセルアートで再現できたら嬉しいものってなんだろうと考えた結果、個人的に好きなゲームSQUARE ENIXのOCTOPATH TRAVELERのキャラクターを学習データとしてピクセルアートを作ることにします。あ、２も出たよ。

そうと決まったら次はどうやって画像を集めるかだということで、画像を探していたら良さそうなサイトを見つけました。ここからデータをお借りすることにしました。

画像切り抜き

上のような画像を64枚手に入れることができました。ただ、この状態で学習にかけることは難しいので1画像1キャラになるように切り抜いていきたいと思います。もちろん面倒なのでPythonで。

# coding=utf-8
from PIL import Image
import glob
import os
import itertools


def data_create(dir_path):
    files = glob.glob(f'{dir_path}/*')
    base_pos = [left, upper, right, lower]
    span = span

    for f in files:
        print(f'File "{os.path.basename(f)}" Crop Start!')
        im = Image.open(f)

        for i, j in list(itertools.product(range(A), range(B))):
            print(f'crop number is {i}, {j}')

            crop_pos = (base_pos[0] + j * span,
                        base_pos[1] + i * span,
                        base_pos[2] + j * span,
                        base_pos[3] + i * span)

            crop = im.crop(crop_pos)
            crop = crop.convert("RGB")
            crop.save(f'traindata_crop/pixel_art/{os.path.basename(f).replace(".png", "")}-{i}{j}.jpg')


if __name__ == '__main__':
    data_create('train_data')

このプログラムで、下みたいな画像を切り抜いていくと8000枚ほど集まりました。学習には十分かなということで、学習準備はこれでOKとしましょう。

DCGANの実装

PyTorchの公式の記事にいいものがあるのでそれを利用していこうと思います。正直ほぼほぼコピペで作れたので、実装内容をここに書く必要はないんじゃないかと思うレベルでした。一応書いておきますが、、、

あと、いつもAIの実装ってプログラム組むこともそうですけど、データの準備をする方が大変だなと思います。

Generatorクラス

# coding=utf-8
import torch.nn as nn

nc = 3
nz = 100
ngf = 64


class Generator(nn.Module):
    def __init__(self, ngpu):
        super(Generator, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            nn.ConvTranspose2d( nz, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            nn.ConvTranspose2d( ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            nn.ConvTranspose2d( ngf * 2, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            nn.ConvTranspose2d( ngf, nc, 4, 2, 1, bias=False),
            nn.Tanh()
            # state size. (nc) x 64 x 64
        )

    def forward(self, input):
        return self.main(input)

Discriminatorクラス

# coding=utf-8
import torch.nn as nn

nc = 3
nz = 100
ndf = 64


class Discriminator(nn.Module):
    def __init__(self, ngpu):
        super(Discriminator, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )

    def forward(self, input):
        return self.main(input)

学習

# coding=utf-8
from __future__ import print_function
import argparse
import os
import random
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from IPython.display import HTML

manualSeed = random.randint(1, 10000)
print("Random Seed: ", manualSeed)
random.seed(manualSeed)
torch.manual_seed(manualSeed)

dataroot = "/dataset/traindata_crop"
workers = 2
batch_size = 128
image_size = 64
nc = 3
nz = 100
ngf = 64
ndf = 64
num_epochs = 600
lr = 0.0002
beta1 = 0.5
ngpu = 1

device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")


def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)


def train(
        netG,
        optimizerG,
        netD,
        optimizerD,
        dataloader,
        criterion,
        fixed_noise,
        real_label,
        fake_label,
        num_epochs
):
    img_list = []
    G_losses = []
    D_losses = []
    iters = 0

    print("Starting Training Loop...")
    for epoch in range(num_epochs):
        for i, data in enumerate(dataloader, 0):

            netD.zero_grad()
            real_cpu = data[0].to(device)
            b_size = real_cpu.size(0)
            label = torch.full((b_size,), real_label, device=device)
            output = netD(real_cpu).view(-1)
            errD_real = criterion(output, label)
            errD_real.backward()
            D_x = output.mean().item()

            noise = torch.randn(b_size, nz, 1, 1, device=device)
            fake = netG(noise)
            label.fill_(fake_label)
            output = netD(fake.detach()).view(-1)
            errD_fake = criterion(output, label)
            errD_fake.backward()
            D_G_z1 = output.mean().item()
            errD = errD_real + errD_fake
            optimizerD.step()

            netG.zero_grad()
            label.fill_(real_label)
            output = netD(fake).view(-1)
            errG = criterion(output, label)
            errG.backward()
            D_G_z2 = output.mean().item()
            optimizerG.step()

            if i % 50 == 0:
                print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f'
                      % (epoch, num_epochs, i, len(dataloader),
                         errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))

            G_losses.append(errG.item())
            D_losses.append(errD.item())

            with torch.no_grad():
                fake = netG(fixed_noise).detach().cpu()
            img_list.append(vutils.make_grid(fake, padding=2, normalize=True))

            iters += 1

        real_batch = next(iter(dataloader))

        plt.figure(figsize=(15, 15))
        plt.subplot(1, 2, 1)
        plt.axis("off")
        plt.title("Real Images")
        plt.imshow(
            np.transpose(vutils.make_grid(real_batch[0].to(device)[:64], padding=5, normalize=True).cpu(), (1, 2, 0)))

        plt.subplot(1, 2, 2)
        plt.axis("off")
        plt.title("Fake Images")
        plt.imshow(np.transpose(img_list[-1], (1, 2, 0)))
        plt.savefig(f"epoch_{epoch}.jpg")
        img_list.clear()
        
    torch.save(netG.state_dict(), 'modelG.pth')
    torch.save(netD.state_dict(), 'modelD.pth')
    return img_list, G_losses, D_losses


def main():
    dataset = dset.ImageFolder(root=dataroot,
                               transform=transforms.Compose([
                                   transforms.Resize(image_size),
                                   transforms.CenterCrop(image_size),
                                   transforms.ToTensor(),
                                   transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                               ]))
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
                                             shuffle=True, num_workers=workers)

    device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")

    real_batch = next(iter(dataloader))
    plt.figure(figsize=(8, 8))
    plt.axis("off")
    plt.title("Training Images")
    plt.imshow(
        np.transpose(vutils.make_grid(real_batch[0].to(device)[:64],
                                      padding=2, normalize=True).cpu(), (1, 2, 0)))
    plt.show()

    netG = Generator(ngpu).to(device)
    if (device.type == 'cuda') and (ngpu > 1):
        netG = nn.DataParallel(netG, list(range(ngpu)))
    netG.apply(weights_init)

    netD = Discriminator(ngpu).to(device)
    if (device.type == 'cuda') and (ngpu > 1):
        netD = nn.DataParallel(netD, list(range(ngpu)))
    netD.apply(weights_init)

    criterion = nn.BCELoss()
    fixed_noise = torch.randn(64, nz, 1, 1, device=device)

    real_label = 1.0
    fake_label = 0.0

    optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))
    optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))

    print(netG)
    print(netD)


    img_list, G_losses, D_losses = train(netG, optimizerG, netD, optimizerD, dataloader,
          criterion, fixed_noise, real_label, fake_label, num_epochs)

    plt.figure(figsize=(10, 5))
    plt.title("Generator and Discriminator Loss During Training")
    plt.plot(G_losses, label="G")
    plt.plot(D_losses, label="D")
    plt.xlabel("iterations")
    plt.ylabel("Loss")
    plt.legend()
    plt.show()


if __name__ == '__main__':
    main()

学習結果

それでは早速学習回数に伴ってどうなってるか見ていこうと思います。64枚1セットの絵を2枚並べてます。左が参考程度に置いている学習元の画像集で右側がAIが作成した画像集です。

Epoch1

やっぱり最初は、全然ダメですね。

Epoch3

まだまだではありますが、真ん中に何かあればいいんだな？ってことぐらいは学んできましたね。

Epoch16

お？少しキャラらしい何かが見えてきました。

Epoch102

いろんなキャラっぽく見えるようになってきてます。ここまでいけるとは正直思わなかった。写真みたいなものでやったことはありましたが、ピクセルアートできる限り情報を落としたものなので学習ちゃんとできるか不安でしたが安心した。

Epoch139

かなりいい感じにキャラ出来上がってるけど、完璧とは言えない。

一個抜き出して拡大してみてみます。どうでしょう？それなりにいいんじゃないでしょうか？

Epoch158

あれ？なんかノイズ増えてる。

Epoch424

ノイズ増えてきたのでもう過学習になってきてるのか？と思ったので一気に学習回数増やしてどうなるかみてみます。

……..？

..？

なんで全部同じキャラに？？？ここについてはなんかわかりませんねぇ。どっかでこのキャラに良さを見つけちゃったのかなAI君。

まとめ

意外と早くピクセルアートらしいものを作ることはできましたが、これで描けたと言えたのかは疑問ではあります。が、ここからCycle GANなどのアーキテクチャを試してもっといい感じなAIを作ってみたいなと思います。