LoginSignup
7
9

More than 5 years have passed since last update.

PyTorchで白黒人物画像に着彩するAIを作る

Last updated at Posted at 2019-04-14

はじめに

 PyTorchでの人工知能制作にはまり、白黒画像に色をつけてみたくなったので、やってみることにしました。細かい説明の前に、最終的な出力結果を紹介します。(人物画像はどれもGoogle画像検索で改変後の非営利目的での再使用が許可された画像です。)
ad.PNG
 パソコンのスペック上低画質の物しか扱えませんが、顔の部分を肌色に塗ることができています。ポイントは、ワイシャツの襟の部分などの白い部分は塗らないようになっているところです。今回はこのAIの仕組みと、簡単なソースコードの解説をしたいと思います。

原理解説

ここからはこのAIの仕組みについて解説します。

学習に使用する画像の読み込み&加工について

 これには画像処理系のライブラリのPython Imaging Library(PIL)を使用します。今回の目的は着彩AIを作ることなので、入力となる白黒画像と、正解データとなるカラー画像が必要になります。白黒画像はPILのconvert('L')から作成することができるため、画像データとして必要になるのは顔写真のカラー画像のみということになります。データセットにはLabeled Faces in the Wildを使用しました。また、これらの画像は読み込み後、切り抜きやリサイズなどの加工を施して最終的に50×50の画像になります。切り抜きについては後に説明します。

モデルについて

 今回は画像処理ということで、定番の畳み込みニューラルネットワーク(CNN)を用いました。と言っても、畳み込み層とプーリング層はそれぞれ一層ずつ、全結合層も一層という簡素な作りになっています。CNNを使う必要性はないかもしれませんが、私が以前作成したAI(これも画像処理系)で、CNNを使う場合と使わない場合で圧倒的にCNNのほうが学習が速かったので、CNNを使います(おそらく畳み込みとプーリングにより画像を圧縮したためだと思っています。今後検証します)。モデルの入力は50×50個のノード、出力は50×50×3個(RGBの3チャンネルなので)のノードとなっています。

学習の際の流れ

 今回のようなニューラルネットワークを用いたAIの学習の際には、

 1. モデルに何らかのデータを入力する
 2. モデルの出力と、正解との誤差を求める
 3. 誤差をもとにモデルを更新する
 4. 1.に戻る

という流れを踏むと思います。今回の場合も基本的に同じです。上の流れを今回作るAIで具体的に言い換えると、

 1. モデルに白黒画像を入力
 2. 出力画像とカラー画像との誤差を求める
 3. 誤差をもとにモデルを更新する
 4. 1.に戻る

という流れになります。入力の白黒画像は各ピクセルごとに0~1の値が格納されている配列です。出力のカラー画像はRGB形式で、これは各ピクセルごとにR,G,Bの各値が0~1の範囲で格納されています(正確には学習が上手くいけば0~1の範囲になります)。この方法で学習してみます。
gf.PNG

 これでも学習自体は成り立っていいるのですが、なかなか鮮明な画像が出力されませんでした(学習する画像の枚数や、エポック数、学習率などの条件はすべて揃えています)。学習を重ねれば鮮明な画像が出力できるようになるのかもしれませんが、僕は待つことが嫌いなのでできるだけ早く学習させるようにしたいものです。ここで、これを解決するための「工夫」をすることにしました。

工夫その1:出力を入力と合成する

 見出しが早速答えとなってしまいましたが、この問題を解決する工夫として、「出力と入力を合成する」ということを思いつきました。これは、このAIの目標である「色を付ける」ということに関してとても合理的な手法だと思います。
 先ほど説明した学習の流れを変更し、誤差算出の前に「出力と入力を合成する」というステップを加えます。これは具体的に何をしているのかというと、「出力のR,B,Gの各値に入力の値をそれぞれ掛ける」という処理をしています。これは画像編集ソフトでレイヤー同士の合成の「乗算モード」と同じです。これにより、入力で色の暗い部分は0に近い値がかけられるため、合成後の画像は0に近い、つまり暗い色が出力されるようになります。これを用いれば、入力画像の明暗の情報を出力画像に継承することができます。この処理を踏まえ、先ほどの学習の流れを書き換えると、

 1. モデルに白黒画像を入力
 2. 出力画像と入力画像を合成
 3. 合成画像とカラー画像との誤差を求める
 4. 誤差をもとにモデルを更新する
 5. 1.に戻る

という流れになります。図で表すと次のようになります(画像はイメージです)。
キャプチャ.PNG
 これにより先ほどよりも圧倒的に早く成果を上げることができるようになります。こちらが結果です。
キャプチャ.PNG
 はっきりとした画像を作ることができました!顔の部分もしっかりと認識しています!と言いたいところですが、画像中央のみが着彩されています。これではわかりにくいかもしれないので、別の画像を用意しました。
キャプチャ.PNG
 二枚目の画像の人物を少し右にずらしてテストをしました。合成後の画像を見てもわかるように、人物の顔の右上は着色されず、画像中央のみが着彩されていることが分かると思います。これは使用したデータセットに画像中央に人の顔があるものが多いことが原因です。これではAIを使う意味がありません。初めから画面中央をオレンジ色に塗った画像を合成すれば済んでしまうからです。そこでこれを防ぐためにもう一つ「工夫」をすることにしました。

工夫その2:画像をランダムに切り抜く

 この問題の原因は画像中央に人の顔があることが原因です。そのため、画像中央から顔をずらすことでこの問題が解決できると考えました。ランダムなサイズの正方形でランダムな位置を切り抜けば顔の位置をずらすことができます。
agf.PNG

この作業を画像を読み込む際に行うことで、画面中央だけに着彩するということを防ぐことができました。(しかしシャツにも色を付けてしまっています。得意不得意があるのでしょうか。今後学習データを増やして再度挑戦します。)
キャプチャ.PNG
原理の説明は以上です。

実装

 ここからは実際にどのようなプログラムを書いたのかを解説します。

下準備

 プログラムを書く前に、画像データの入ったディレクトリを用意します。まずはD:\PythonPrograms\PaintFaceAI\TrainDataというフォルダを用意します。もちろん別の場所でも構いません。パスを変更した場合はプログラム内のディレクトリ指定用の変数を書き換えてください。このディレクトリには学習に使う訓練データを入れます。face (number).jpgとなるように画像を用意してください。(僕のパソコン(Windows)では画像を全選択→右クリック→名前を変更→faceと入力→Enterとする事で自動に番号付けをすることができました。)すべての画像のサイズは250×250です。また、最低でも8192枚の画像を用意してください(この数は後述の変数TRAIN_NUMの値です)。
キャプチャ.PNG
 同様にD:\PythonPrograms\PaintFaceAI\TestDataも用意します。TrainDataのディレクトリと同様の画像(ただし学習しないので白黒も可)をこちらは5枚用意します(もっと多くの画像をテストしたい場合はTEST_NUMの値を変更してください)。
キャプチャ.PNG
 最後に、モデルの保存先となるフォルダ(:D:\PythonPrograms\PaintFaceAI\Model)を用意してください。もちろんこの場所も自由に決めることができますが、その際は後述の変数のMODEL_LOCを変更してください。出力画像の保存先等はフォルダを自動で初期化するようにしたのですが、モデルは一度消えるとまた数時間学習しなければならないので、めんどくさいですが安全確保のため手動です。フォルダは空のままでOKです。

ライブラリのインポート

 使用するライブラリをインポートします。

Colorize.py
#ライブラリのインポート
import torch
from torch import nn
from torch import optim
import torch.nn.functional as F
import numpy as np
from PIL import Image
import os
import shutil

nnはニューラルネットワークを使うために、optimはPytorchが提供するオプティマイザーを使用するために、nn.functionalは誤差関数を使用するためにインポートします。尚、numpyは画像読み込み時に、osshutilは画像保存先等のディレクトリを整えるときに使います。

設定用の変数の定義

 良いパラメータを探す際や、様々な実験を行う際に、プログラムのあちこちに散らばっている設定を一つの箇所にまとめておけば楽なので、あらかじめ設定用の変数を用意し、プログラム内でそれらを使用します。

Colorize.py
#設定
SOURCE_IMAGE_SIZE = 250         #画像の縦と横のピクセル数
IMAGE_SCALE = 0.2               #元画像の何倍の大きさの画像を扱うのか
IMAGE_SIZE = int(SOURCE_IMAGE_SIZE*IMAGE_SCALE) #扱う画像のサイズ
TRAIN_NUM = 8192                #学習に使う画像の数
EPOCH_NUM = 16                  #同じデータセットを繰り返し学習する数
TEST_NUM = 5                    #テストする画像の数
LR = 0.0001                     #学習率
BATCH_SIZE = 128                #バッチサイズ
BASIC_COLOR_SPACE = "RGB"       #色空間の指定
#デバイスの指定(グラボが使えるなら使う)
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
LEARN = True    #学習を行うのか
TEST = True     #学習中に1エポックごとにテストを行うのか
TEST_VISUALIZE = True #テストで画像を出力するのかどうか
MODEL_NAME = "Model"#モデル保存時の名前

#使用するディレクトリのパス(適宜指定してください)

#学習用の画像が保存されたディレクトリ
TRAIN_DATA_LOC = r"D:\PythonPrograms\PaintFaceAI\TrainData"

#テスト用の画像が保存されたディレクトリ
TEST_DATA_LOC = r"D:\PythonPrograms\PaintFaceAI\TestData"

#最終的な出力画像の保存先
OUT_LOC = r"D:\PythonPrograms\PaintFaceAI\Output"

#入力画像の保存先
IN_LOC = r"D:\PythonPrograms\PaintFaceAI\Input"

#元画像(学習用の画像を切り抜き、リサイズ処理したもの)の保存先
SOURCE_LOC = r"D:\PythonPrograms\PaintFaceAI\Source"

#色情報(ネットワーク出力を入力と合成する前の画像)の保存先
COLOR_LOC = r"D:\PythonPrograms\PaintFaceAI\Color"

#モデルの保存先
MODEL_LOC = r"D:\PythonPrograms\PaintFaceAI\Model"

画像読み込み用クラス

 画像を読み込むための機能を持つクラスを用意します。
 このクラスはふたつのメソッドを持ちます。一つ目はload_batchメソッドです。これは学習時に使用します。指定された枚数の画像を読み込み、切り抜き、リサイズの処理を行った後にそれを白黒化します。また、カラー画像と白黒画像共に各ピクセルごとの色の情報を配列に格納し、それを返します。二つ目のload_test_imageでは、リサイズと白黒化のみを行い、指定枚数の白黒画像の情報の入った配列を返します。基本的にload_batchメソッドと同じです。

Colorize.py
#画像読み込み用クラス
class image_loder():
    #ミニバッチ学習に使用する1ミニバッチ分の画像を取得する
    def load_batch(self,image_number,batch_size,
        random_crop = True, crop_max = 200,message = ''):

        #画像のピクセル数と同じ数だけの要素を持つTensor型の配列を作る
        color_image = torch.Tensor(batch_size,3,IMAGE_SIZE,IMAGE_SIZE)
        greyscale_image = torch.Tensor(batch_size,1,IMAGE_SIZE,IMAGE_SIZE)

        #バッチサイズの回数繰り返す
        for num in range(batch_size):
            #ファイル名
            image_name = "face (" + str(int(image_number+1+num)) + ").jpg"

            #画像読み込みのついでに進捗を確認できるようなメッセージを出力
            print("Loading Image :",image_name,"  |  ",message)

            #画像を読み込む
            image = Image.open(TRAIN_DATA_LOC + "\\" + image_name)

            #ランダムな場所で切り抜く
            if random_crop == True:
                #切り抜く幅の設定
                crop_size = np.random.randint(1,crop_max)
                #切りぬく正方形の左上の座標の指定
                x = np.random.randint(0,crop_size)
                y = np.random.randint(0,crop_size)
                #切り抜く
                image = image.crop((x,y,x+SOURCE_IMAGE_SIZE-crop_size,y+SOURCE_IMAGE_SIZE-crop_size))

            #指定の大きさにリサイズ
            image = image.resize((IMAGE_SIZE,IMAGE_SIZE)).convert(BASIC_COLOR_SPACE)

            #配列の各要素に読み込んだ画像の各ピクセルの色を格納する
            for x in range(IMAGE_SIZE):
                for y in range(IMAGE_SIZE):
                    pixcel = image.getpixel((x,y))
                    #白黒なのでRGBの平均を取る。その後0~1の範囲に変換する
                    color_image[num,0,x,y]=pixcel[0]/255
                    color_image[num,1,x,y]=pixcel[1]/255
                    color_image[num,2,x,y]=pixcel[2]/255

            #白黒に変換
            image = image.convert('L')

            #配列の各要素に読み込んだ画像の各ピクセルの色を格納する
            for x in range(IMAGE_SIZE):
                for y in range(IMAGE_SIZE):
                    pixcel = image.getpixel((x,y))
                    greyscale_image[num,0,x,y]=pixcel/255

        #配列の形状を整える
        color_image = color_image.view(batch_size,
                                        3,
                                        IMAGE_SIZE,
                                        IMAGE_SIZE).to(DEVICE)
        greyscale_image = greyscale_image.view(batch_size,
                                        1,
                                        IMAGE_SIZE,
                                        IMAGE_SIZE).to(DEVICE)

        #画像データの配列を返す
        return [color_image,greyscale_image]

    #テスト用に使用する画像の取得。取得枚数は別指定。
    #仕組みはほぼload_batchと同じなのでコメントは省略
    def load_test_image(self,test_num,message = ''):
        greyscale_image = torch.Tensor(test_num,1,IMAGE_SIZE,IMAGE_SIZE)
        for num in range(test_num):
            image_name = "face (" + str(int(num+1)) + ").jpg"
            print("Loading Image :",image_name,"  |  ",message)
            image = Image.open(TEST_DATA_LOC + "\\" + image_name)
            image = image.resize((IMAGE_SIZE,IMAGE_SIZE)).convert(BASIC_COLOR_SPACE)
            image = image.convert('L')
            for x in range(IMAGE_SIZE):
                for y in range(IMAGE_SIZE):
                    pixcel = image.getpixel((x,y))
                    greyscale_image[num,0,x,y]=pixcel/255
        greyscale_image = greyscale_image.view(test_num,
                                        1,
                                        IMAGE_SIZE,
                                        IMAGE_SIZE).to(DEVICE)
        return greyscale_image

load_batchメソッドの引数について

引数が多いので簡単に説明します。

引数 説明
image_number 画像読み込み時に、何番目の画像から読み込むのかを指定します。
batch_size バッチサイズを指定します。このメソッドではバッチサイズ個の画像を返します
random_crop ランダムな位置で切り抜くのか否かを指定します。ランダムに切り抜くことで画像中央ばかりに人の顔があることがなくなります
crop_max 切り抜く幅の最大値を指定します。例えばcrop_maxが10で、元画像のサイズが50×50の場合、40×40~50×50の大きさの画像が作られます
message コンソールに表示するメッセージを指定します。画像読み込みが一番時間がかかるため、このついでに進捗を確認できるので便利です。必ず必要というわけではありません。

load_test_imageメソッドの引数について

こちらも簡単に説明します。

引数 説明
test_num 何枚の画像をテストするのかを指定します。
messege コンソールに表示するメッセージを指定します。

ニューラルネットワークを構築するクラス

 今回はnn.Moduleクラスを継承するという形でネットワークを定義します。

Colorize.py
#ニューラルネットワークの設定
class network(nn.Module):
    #nn.Moduleを継承してモデルを用意する
    def __init__(self):
        super(network,self).__init__()

        #畳み込み層(画像を圧縮して学習速度が高まった)
        self.features = nn.Sequential(
            nn.Conv2d(1,8,kernel_size=2,stride=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2,stride=1),
        )

        #全結合層
        self.classifier = nn.Sequential(
            nn.Linear(18432,(IMAGE_SIZE**2)*3),
            nn.ReLU(inplace=True),
            nn.Linear((IMAGE_SIZE**2)*3,(IMAGE_SIZE**2)*3)
        )

    #インプットからアウトプットまでの流れ
    def forward(self,x):
        x = self.features(x)
        x = x.view(x.size(0),-1)
        x = self.classifier(x)
        return x

可視化用のクラス

 プログラム内での画像データは配列になっているので、これを可視化する必要があります。そのためのクラスを作ります。

Colorize.py
#可視化用のクラス
class visualizer():
    def __init__(self):
        return None

    #画像に変換
    def visualize(self,image,size,name,show_or_not,color_space):
        #画像の縦横のピクセル数を用意
        width,height = size

        #指定サイズの画像を用意
        image_v = Image.new(color_space,(width,height))

        #白黒画像でない場合
        if color_space != "L":
            for x in range(width):
                for y in range(height):
                    #各ピクセルの色を入力された配列を参照して決定する。配列では0~1だが、0~255に変換してから出力する
                    value=image[:,x,y]
                    image_v.putpixel((x,y),(int((value[0])*255),int((value[1])*255),int((value[2])*255)))

        #白黒画像の場合
        else:
            for x in range(width):
                for y in range(height):
                    #各ピクセルの色を入力された配列を参照して決定する。配列では0~1だが、0~255に変換してから出力する
                    value=image[:,x,y]
                    image_v.putpixel((x,y),(int(value[0]*255)))

        #画像の表示
        if show_or_not == True:
            image_v.show()

        #画像の保存
        image_v.convert("RGB").save(name+".bmp")

着彩用のクラス

 このプログラムの本体となるクラスです。かなりややこしく仕上がっております。今回はメソッドごとに説明します。

__init__メソッド

 このクラスの初期設定です。ここではモデルの作成と最適化方法の選択をします。今回は最適化方法はAdamを使用しました。

Colorize.py

#色を塗るAIのクラス
class painter():
    def __init__(self):
        #ネットワークの作成  
        self.model = network().to(DEVICE)
        print("Created Model:\n",self.model) 

        #最適化手法の設定
        self.optimizer = optim.Adam(self.model.parameters(),lr = LR)

        #続く...

multiplyメソッド(合成処理)

 先に合成処理用のメソッドを作っておきます。これを一つのメソッドにする必要はないかもしれませんが、合成法を変える場合に便利なのでメソッドにしました。

Colorize.py

    #白黒画像とネットワークの出力を合成
    def multiply(self,a,b):
        return a * b
        #続く...

paint_trainメソッド

 このメソッドでは学習時の画像処理関連の流れをまとめています。白黒画像を受け取ると、それをネットワークに入力し、出力を得ます。さらにその出力を入力と合成し、合成されたものを返します。また、学習の進捗を確認するために画像を可視化し保存しています。

Colorize.py

    #学習時の画像処理の担当
    def paint_train(self,epoch,num,color_image,greyscale_image,visualize = True):

        #白黒画像をニューラルネットワークに入れて計算させる
        output = self.model(greyscale_image).view(BATCH_SIZE,3,IMAGE_SIZE,IMAGE_SIZE).to(DEVICE)

        #ニューラルネットワークの出力をcolor_dataとして保存
        color_data = output

        #画像を合成
        output = self.multiply(output,greyscale_image)

        #バッチごとに画像を保存する(とりあえず配列の0番目の画像を出力する)
        if visualize == True:
            print("Visualizing Source Image")
            vi.visualize(color_image[0].view(3,int(IMAGE_SIZE),int(IMAGE_SIZE)),
                        [int(IMAGE_SIZE),int(IMAGE_SIZE)],
                        SOURCE_LOC + "\source_" + str(epoch+1) + "-" + str(num),
                        False,BASIC_COLOR_SPACE)

            print("Visualizing Input Image")
            vi.visualize(greyscale_image[0].view(1,int(IMAGE_SIZE),int(IMAGE_SIZE)),
                        [int(IMAGE_SIZE),int(IMAGE_SIZE)],
                        IN_LOC + "\imput_" + str(epoch+1) + "-" + str(num),
                        False,"L")

            print("Visualizing Color Data")
            vi.visualize(color_data[0].view(3,IMAGE_SIZE,IMAGE_SIZE),
                        [IMAGE_SIZE,IMAGE_SIZE],
                        COLOR_LOC + "\color_" + str(epoch+1) + "-" + str(num),
                        False,BASIC_COLOR_SPACE)

            print("Visualizing Output Image")   
            vi.visualize(output[0].view(3,IMAGE_SIZE,IMAGE_SIZE),
                        [IMAGE_SIZE,IMAGE_SIZE],
                        OUT_LOC + "\output_" + str(epoch+1) + "-" + str(num),
                        False,BASIC_COLOR_SPACE)

        return output
        #続く...

updateメソッド

 モデルの更新をします。引数に本来のカラー画像(正解データのようなもの)とAIの出力を持ち、これらの誤差を計算したのちに誤差伝搬します。

Colorize.py
    #モデルの更新
    def update(self,color_image,output):
        print("Updating Model")

        #誤差を計算する(誤差関数)
        loss = F.smooth_l1_loss(color_image,output)

        #誤差関数に依存しない誤差の計算(差の合計)
        real_loss = torch.abs(color_image-output).sum()

        #誤差伝搬
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        print('Loss:',str(float(real_loss)))
        #続く...

trainメソッド

 このメソッドでは、学習全体の流れを記述します。

Colorize.py

    #学習
    def train(self):
        print("Train")

        for epoch in range(EPOCH_NUM):
            #テストするならテストする
            if TEST:
                print("Test")
                self.test(epoch)

            #学習本体
            for num in range(int(TRAIN_NUM/BATCH_SIZE)):
                #進捗確認用のメッセージ
                progress = "epoch : "+str(epoch+1)+" - batch : "+str(num+1)+" / "+str(int(TRAIN_NUM/BATCH_SIZE)) + " - Batch size = "+str(BATCH_SIZE)+" - lr = "+str(LR)
                print("\n",progress)

                #テストデータを着彩#num番目の画像を読み込む
                [color_image,greyscale_image] = il.load_batch(num*BATCH_SIZE,BATCH_SIZE,message=progress,random_crop = True)

                #画像処理
                output = self.paint_train(epoch,num,color_image,greyscale_image,visualize = False)

                #モデルの更新
                self.update(color_image,output)

            print("Saving Model")
            #モデルの保存
            torch.save(self.model.state_dict(),MODEL_LOC+"/"+MODEL_NAME+str(epoch)+".pth")
            #続く...

paint_testメソッド

 テスト時の画像処理を担当します。基本的にpaint_trainメソッドと同じです。

Colorize.py
    #テストでの画像処理を担当(基本的にpaint_trainと同じ。
    #paint_testでは出力する画像の枚数を指定できるようにした。詳細なコメントは省略)
    def paint_test(self,epoch,greyscale_image,visualize = True,visualize_num = 1):
        output = self.model(greyscale_image).view(TEST_NUM,3,IMAGE_SIZE,IMAGE_SIZE).to(DEVICE)
        color = output
        output = self.multiply(output,greyscale_image)
        if visualize == True:
            for v_num in range(visualize_num):
                print("Visualizing Input Image")
                vi.visualize(greyscale_image[v_num].view(1,int(IMAGE_SIZE),int(IMAGE_SIZE)),
                            [int(IMAGE_SIZE),int(IMAGE_SIZE)],
                            IN_LOC + "\Test_imput_" + str(epoch) + "-" +str(v_num+1),
                            False,"L")
                print("Visualizing Color Data")
                vi.visualize(color[v_num].view(3,IMAGE_SIZE,IMAGE_SIZE),
                            [IMAGE_SIZE,IMAGE_SIZE],
                            COLOR_LOC + "\Test_color_" + str(epoch) + "-" +str(v_num+1),
                            False,BASIC_COLOR_SPACE)  
                print("Visualizing Output Image")   
                vi.visualize(output[v_num].view(3,IMAGE_SIZE,IMAGE_SIZE),
                            [IMAGE_SIZE,IMAGE_SIZE],
                            OUT_LOC + "\Test_output_" + str(epoch) + "-" +str(v_num+1),
                            False,BASIC_COLOR_SPACE)
        return output
        #続く...

testメソッド

 テスト時の全体の流れを記述します。これもtrainメソッドとほぼ同じですが、誤差の計算やモデルの更新はしません。

Colorize.py
    #テスト
    def test(self,text):
        print("TestStart")

        #テストデータを着彩#num番目の画像を読み込む
        greyscale_image = il.load_test_image(TEST_NUM)

        self.paint_test(text,greyscale_image,visualize = TEST_VISUALIZE,visualize_num = TEST_NUM)

        print('Test End')

ディレクトリを初期化するメソッド

 画像の出力先のフォルダをいったん削除し、再びフォルダを作ります。手動でフォルダ内の画像を削除する手間が省けるのででスムーズに実験が進みます。

Colorize.py
#画像保存用のディレクトリを用意する
def directory_setup():
    #フォルダがない状態でフォルダを削除しようとするとエラーになるのでエラーになったら削除しない
    try:
        #出力先のフォルダの消去
        shutil.rmtree(SOURCE_LOC)
    except:
        pass
    try:
        #出力先のフォルダの消去
        shutil.rmtree(IN_LOC)
    except:
        pass
    try:
        #出力先のフォルダの消去
        shutil.rmtree(OUT_LOC)
    except:
        pass
    try:
        #出力先のフォルダの消去
        shutil.rmtree(COLOR_LOC)
    except:
        pass

    os.mkdir(SOURCE_LOC)
    os.mkdir(IN_LOC)
    os.mkdir(OUT_LOC)
    os.mkdir(COLOR_LOC)

実行

 最後に作成したクラスのオブジェクトを作成し、学習を開始するコードを書きます。

Colorize.py
#クラスの作成
il = image_loder()
p = painter()
vi = visualizer()

#ディレクトリを整理
directory_setup()

#学習
if LEARN:
    p.train()

p.test("Result")

print("Finish")

これを実行すると、画像保存用のフォルダが作られ、学習が進みます。気長に待ちましょう。ちなみに実行画面はこんな感じで、進捗度合いなどが表示されます。
キャプチャ.PNG
実装の説明は以上です。

おわりに

 今回作成したAIの最大の成果は、ほとんどの画像でシャツの白い部分などを塗らないことを学習したことだと考えています。背景色まで肌色に塗ってしまっていることに関しては、背景色を適切に塗るには膨大な画像データが必要となるため解決は難しいと考えています。しかし、瞳の色や唇の色など、人間の顔である程度共通する箇所は適切な色を塗れるようにしたいため、高解像度の画像を扱うようにする、データの量を増やすなどして改善していきたいと考えています。

主な参考文献(2019/04/14時点)

  1. PyTorchニューラルネットワーク実装ハンドブック(書籍)
  2. torch.Tensorの作成と基本操作
  3. PIL/Pillow チートシート
  4. 初めてのPython画像処理
  5. Python, Pillowで画像の一部をトリミング(切り出し/切り抜き)
  6. Python, Pillowで画像を一括リサイズ(拡大・縮小)
7
9
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
7
9