LoginSignup
217
194

More than 3 years have passed since last update.

@Takacie

機械学習で画像の高画質化を試みる（備忘録）- Python3

Last updated at Posted at 2020-02-21

1.はじめに

最近、Twitterで「謎の技術で高画質化された画像」なるものがタイムラインにいくつか流れてきて興味が湧いたので、機械学習の勉強がてら画像の高画質化の方法を、僕のように「理屈無しで手っ取り早く機械学習に触れたい！」という人に向けて備忘録としてここに残しておくことにしました。


2.実行環境の構築

今回、この手の機械学習でよく用いられる「pix2pix」を使用しました。pix2pixはGANを用いた画像生成アルゴリズムで、２枚の対になった画像から相互間の関係を学習し、１枚の元画像に対して学習結果に基づいて画像を生成することができます。※詳細: [https://phillipi.github.io/pix2pix/]
pix2pixはPyTorchやKerasなどのライブラリを用いて実行することができますが、今回はtensorflowを用いた実行環境を構築しました。以下に僕が今回使用したライブラリなどをリスト化しておきます。

・Python 3.7.6
・pix2pix-tensorflow [https://github.com/affinelayer/pix2pix-tensorflow]
・tensorflow 1.14.0（最新バージョンだと実行できないので注意）
・OpenCV 4.2.0.32

tensorflowの最新バージョンではpix2pixを実行できないので、pipでインストールする際に以下のようにバージョンを指定します。
※（20/11/29）現在はtensorflow1.4.1まで対応しているそうです。以下のコマンドを適宜変更してください。

pip install tensorflow==1.14.0

OpenCVはpix2pixの実行とは直接関係しないのですが、後述の学習データの作成で使用します。今の所最新バージョンで問題なさそうですが、もし実行できないのであれば同じようにバージョンを指定してインストールを行ってください。

3.学習データの作成

pix2pixには予め学習用データセットが用意されていますが、今回の目的は画像の高画質化なので、それに適した学習データを自作することにします。
前述したとおり、pix2pixは以下のような２枚の対になった画像を繋げて１枚にしたものを学習データとして用います。（256px × 256pxを２枚繋げた512px × 256pxの画像）
052_2_1.jpg

この画像を作成するためには

1. 大きな画像を256px × 256pxに切り分ける
2. 切り分けたそれぞれの画像にぼかしを入れる
3. 切り分けた画像とぼかしを入れた画像を結合する

の３つの処理を行う必要があります。そのためのコード（pix2pix_util.py）を作成したのでコピペして使用してください。
※追記にこの３つの処理を一気にやってくれるコードを載せています

pix2pix_util.py
import cv2
import numpy as np
from argparse import ArgumentParser
import glob
import os

if __name__ == "__main__":
    arg = ArgumentParser()
    arg.add_argument('-md', '--mode', help='select mode [split|blur|margin]', required=True)
    arg.add_argument('-id', '--inputdir', help='select input directory', required=True)
    arg.add_argument('-id2', '--inputdir2', help='select second input directory(margin)')
    arg.add_argument('-od', '--outputdir', help='select output directory', required=True)
    args = arg.parse_args()

    #画像を切り分ける
    if args.mode == "split":
        os.makedirs(args.outputdir, exist_ok=True) #アウトプット先が存在しなければ作成
        if os.path.exists(args.inputdir) == True:
            files = glob.glob(args.inputdir+"\\*") #ファイル一覧をリスト形式で返す
            print("SPLIT: Processing...")
            for i in files:
                file_name = os.path.basename(i) #この関数でファイル名部分だけを返すことができる
                img = cv2.imread(i) #画像の読み込み
                if type(img) == np.ndarray: #OpenCVは画像をnumpy.ndarrayの型で読み込んでいるのでこの条件文
                    height, width = img.shape[:2]
                    for j in range(int(height / 256)):
                        for k in range(int(width / 256)):
                            splited_image = img[256 * j : 256 * j + 256, 256 * k : 256 * k + 256]
                            cv2.imwrite("{0}\\{1}_{2}_{3}.jpg".format(args.outputdir, file_name[:-4], str(j), str(k)), splited_image) #画像の書き出し
                else:
                    print("ERROR: {0} is not supported type".format(file_name)) #ファイルが画像じゃなかった場合のエラー
            print("SPLIT: Completed") #完了のメッセージ
        else:
            print("ERROR: No such directory \"{0}\"".format(args.inputdir)) #インプットディレクトリが存在しない場合のエラー
    #画像をぼかす
    elif args.mode == "blur":
        os.makedirs(args.outputdir, exist_ok=True)
        if os.path.exists(args.inputdir) == True:
            files = glob.glob(args.inputdir+"\\*")
            print("BLUR: Processing...")
            for i in files:
                file_name = os.path.basename(i)
                img = cv2.imread(i)
                if type(img) == np.ndarray:
                    blur_image = cv2.blur(img, (10, 10)) #OpenCVのBlur関数
                    cv2.imwrite(args.outputdir+"\\"+file_name, blur_image)
                else:
                    print("ERROR: {0} is not supported type".format(file_name))
            print("BLUR: Completed")
        else:
            print("ERROR: No such directory \"{0}\"".format(args.inputdir))
    #画像を繋げる
    elif args.mode == "margin":
        if args.inputdir2 != None:
            os.makedirs(args.outputdir, exist_ok=True)
            if os.path.exists(args.inputdir) == True and os.path.exists(args.inputdir2) == True:
                files1 = glob.glob(args.inputdir+"\\*")
                files2 = glob.glob(args.inputdir2+"\\*")
                print("MARGIN: Processing...")
                for i in files1:
                    file_name = os.path.basename(i)
                    if args.inputdir2+"\\"+file_name in files2:
                        image_1, image_2 = cv2.imread(i), cv2.imread(args.inputdir2+"\\"+file_name)
                        if type(image_1) == np.ndarray and type(image_2) == np.ndarray:
                            margined_image = cv2.hconcat([image_1, image_2]) #OpenCVの画像結合関数
                            cv2.imwrite(args.outputdir+"\\"+file_name, margined_image)
                        else:
                            print("ERROR: {0} is not supported type".format(file_name))
                    else:
                        print("ERROR: \"{0}\" is not in \"{1}\"".format(file_name, args.inputdir2))
                print("MARGIN: Completed")
            else:
                print("ERROR: No such directory \"{0}\" or \"{1}\"".format(args.inputdir, args.inputdir2))
        else:
            print("ERROR: Select second inputdir (-id2/--inputdir2)")
    else:
        print("ERROR: Select a mode from \"split\", \"blur\", \"margin\"")

使用方法

pix2pix_util.pyはコマンドライン上で実行します。

まず、大きな画像を切り分ける（ここで切り分ける画像は幅、高さともに256px以上である必要があります。また、形式はOpenCVが対応している形式なら問題ないと思いますが、ファイル名に日本語が入っているとエラーが発生するようです。）

$ python pix2pix_util.py --mode "split" --inputdir <切り分けたい画像があるディレクトリ> --outputdir <切り分けた画像の保存先>

次に、切り分けられた画像にぼかしを入れる

$ python pix2pix_util.py --mode "blur" --inputdir <切り分けた画像があるディレクトリ> --outputdir <ぼかしを入れた画像の保存先>

最後に、切り分けられた画像とそれにぼかしを入れた画像を結合する

$ python pix2pix_util.py --mode "margin" --inputdir <切り分けた画像があるディレクトリ> --inputdir2 <ぼかしを入れた画像があるディレクトリ> --outputdir <切り分けた画像の保存先>

この結合された画像を用いて学習を行いますが、学習後に学習結果のテストも行うため、学習データとテストデータを２つのフォルダにに分けておきましょう。（一応pix2pix-tensorflow内にあるsplit.pyで同じことができますが、手動でも問題ないと思います。）

4.実際に学習してみる

pix2pix-tensorflowの中に「pix2pix.py」というファイルがあり、これを実行することによって学習やテストを行うことができます。学習を行うにはコマンドライン上で以下のコマンドを実行します。

$ python pix2pix.py --mode train --output_dir <学習結果の保存先> --max_epochs <最大の世代数> --input_dir <学習データがあるディレクトリ> --which_direction BtoA

max_epochsで最大の世代を指定することができますが、あまりこの値が大きすぎると（GPUを搭載していないパソコンでは特に）膨大な時間を要するので、この値を小さくするか、学習データを減らす必要があります。ここはそれぞれの環境によってうまく調節しましょう。

学習が終了したら学習で使用した画像とは別の画像（テストデータ）を用いて学習結果のテストを行います。

$ python pix2pix.py --mode test --output_dir <生成データの保存先> --input_dir <テストデータがあるディレクトリ> --checkpoint <学習結果のディレクトリ>

output_dirで指定したディレクトリに生成されたデータが保存されます。checkpointでは学習時に指定した学習結果を保存したディレクトリを指定してください。

5.学習結果を見てみる

338a4062f9ea56be0a1a1fc4cc564971.png
4055a421c6c24270032f3432c8b8becf.png
500枚の学習データを第５世代まで学習させた結果です。（左からぼかし画像、生成画像、元画像）
パソコンの性能と時間の都合上それほど学習量は多くありませんが、それでもぼかし画像と比べて輪郭がくっきりしているのが確認できます。学習データの数や世代数をもっと大きくすればさらに良い結果を期待できると思われます。

6.まとめ

今回はpix2pixを用いて機械学習によって画像の高画質化を試みました。自分自身、機械学習に実際に触れるのは初めてだったので、どんなものなのかと試行錯誤しながらこのテーマに望みましたが、改めて機械学習の潜在能力の高さに感銘を受けました。
（Qiitaでの投稿も初めてだけど上手くできただろうか...）

追記

学習データ作成で３つの処理を行いましたが、「手っ取り早く」という点でこのコードは向いていないと思い、先ほど作成したpix2pix_util.pyを改変しました。これで１回のコマンド実行で一気に学習データの画像まで作成してくれます。容量も圧迫するのでね...

pix2pix_util.py
import cv2
import numpy as np
from argparse import ArgumentParser
import glob
import os

if __name__ == "__main__":
    arg = ArgumentParser()
    arg.add_argument('-id', '--inputdir', help='select input directory', required=True)
    arg.add_argument('-od', '--outputdir', help='select output directory', required=True)
    args = arg.parse_args()

    os.makedirs(args.outputdir, exist_ok=True) #アウトプット先が存在しなければ作成
    if os.path.exists(args.inputdir) == True:
        files = glob.glob(args.inputdir+"/*") #ファイル一覧をリスト形式で返す
        print("PIX2PIX_UTIL: Processing...")
        for i in files:
            file_name = os.path.basename(i) #この関数でファイル名部分だけを返すことができる
            img = cv2.imread(i) #画像の読み込み
            if type(img) == np.ndarray: #OpenCVは画像をnumpy.ndarrayの型で読み込んでいるのでこの条件文
                height, width = img.shape[:2]
                for j in range(int(height / 256)):
                    for k in range(int(width / 256)):
                        splited_image = img[256 * j : 256 * j + 256, 256 * k : 256 * k + 256]
                        blurred_image = cv2.blur(splited_image, (10, 10))
                        margined_image = cv2.hconcat([splited_image, blurred_image]) #OpenCVの画像結合関数
                        cv2.imwrite("{0}/{1}_{2}_{3}.jpg".format(args.outputdir, file_name[:-4], str(j), str(k)), margined_image) #画像の書き出し
            else:
                print("ERROR: {0} is not supported type".format(file_name)) #ファイルが画像じゃなかった場合のエラー
        print("PIX2PIX_UTIL: Completed") #完了のメッセージ
    else:
        print("ERROR: No such directory \"{0}\"".format(args.inputdir)) #インプットディレクトリが存在しない場合のエラー

$ python pix2pix_util.py -id <画像があるディレクトリ> -od <保存先>

前と同じく、元画像は幅、高さともに256px以上のファイル名は半角英数字でお願いします。

217
194
4

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
Sign upLogin
217
194