概要

作成したAIのモデルを簡単にデモできるGradioと言うライブラリを見つけたので、それを利用して画像を高解像度にするWebページを作成した。利用するモデルは、Tensorhubに公開されている、画像の解像度を4倍に拡大するGANの手法であるESRGANを利用する。

作成されるWeb画面は以下のような感じ。

スクリーンショット 2020-11-23 18.18.08.png

実行環境

Ubuntu18.04
Python3.8.0

すぐ始めたい人

以下のコマンドを実行後、http://127.0.0.1:7860/ にアクセスする。

git clone https://github.com/sey323/gradio-esrgan.git
cd gradio-esrgan
pip install -r requirements.txt
python gradio-esrgan.py

実装

0. 全体像

ソースコードの全体像

import gradio as gr
import tensorflow as tf
import tensorflow_hub as hub

model = hub.load("https://tfhub.dev/captain-pool/esrgan-tf2/1")


def predict(
    inp,
):
    """
    ESRGANを用いて、入力画像を4倍の解像度に変換する。
    https://tfhub.dev/captain-pool/esrgan-tf2/1
    inp :
        入力画像
    """
    h, w, c = inp.shape  # 入力画像から画像サイズを取得する
    inp = inp.reshape((-1, h, w, c))
    inp = tf.cast(inp, tf.float32)
    output = model(inp)
    output = tf.cast(tf.clip_by_value(output[0], 0, 255), tf.uint8)
    return output


def gradio_gui():
    """
    gradioのGui画面を定義する
    """
    image = gr.inputs.Image(label="Input Image",)
    output = gr.outputs.Image(label="Output Image", type="numpy")

    interface = gr.Interface(fn=predict, inputs=image, outputs=output)
    interface.launch()


if __name__ == "__main__":
    gradio_gui()

1. Tensorflow hub, gradioのインストール

Tensorflow hubのインストール

公式サイトを参考にTensorflow hubをインストールする。

pip install "tensorflow>=2.0.0"
pip install --upgrade tensorflow-hub

gradioのインストール

公式サイトを参考にインストールする。

pip install gradio

2. Tensorflow hubのモデルを利用する

TensorFlow Hubの公式サイトを参考に、利用するESRGANのモデルのダウンロードと、プログラムの作成を行う。 Gradioで実行できるように、入力画像をTensorflowで学習できる形式に変換し、高解像度の画像を返すmethodを作成した。

import tensorflow as tf
import tensorflow_hub as hub

model = hub.load("https://tfhub.dev/captain-pool/esrgan-tf2/1")


def predict(
    inp,
):
    """
    ESRGANを用いて、入力画像を4倍の解像度に変換する。
    https://tfhub.dev/captain-pool/esrgan-tf2/1
    inp :
        入力画像
    """
    h, w, c = inp.shape  # 入力画像から画像サイズを取得する
    inp = inp.reshape((-1, h, w, c))
    inp = tf.cast(inp, tf.float32)
    output = model(inp)
    output = tf.cast(tf.clip_by_value(output[0], 0, 255), tf.uint8)
    return output

3. gradioを用いたWeb画面の定義

gradioは、大きく分けて3つのクラスで構成されている。

クラス	概要
gr.inputs	モデルの入力(Image, Textbox, Audioなど)
gr.outputs	モデルの出力(Image, Textbox, Audioなど)
gr.Interface	Gradioで実行する関数。入力エリア、出力エリアを定義し、Web画面を描画する。

inputsとoutputsには、画像以外にもテキストや音声、チェックボックス、スライダーといった物も選択できるため、利用するモデルに応じて柔軟に対応できる。 ESRGANは、入力出力ともに画像しか利用しないので、gr.inputs.Imageを選択する。 Imageはpillowの形式で与えられる。

Gradio | Docs

import gradio as gr

~~~~省略~~~~~

def gradio_gui():
    """
    gradioのGui画面を定義する
    """
    image = gr.inputs.Image(label="Input Image",)
    output = gr.outputs.Image(label="Output Image", type="numpy")

    interface = gr.Interface(fn=predict, inputs=image, outputs=output)
    interface.launch()


if __name__ == "__main__":
    gradio_gui()

動作確認

実行は以下のコマンドで行う。

$ python gradio-esrgan.py 
2020-11-23 01:39:34.566267: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x7f9d83e7e520 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2020-11-23 01:39:34.566291: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Host, Default Version
Running locally at: http://127.0.0.1:7860/
To get a public link for a hosted model, set Share=True
Interface loading below...
<IPython.lib.display.IFrame object at 0x147170040>

実行後、URLが表示されるのでアクセスする。 * http://127.0.0.1:7860/

Web画面にアクセスすると、以上のような画面が表示される。左のパネルにgr.inputsに指定した、モデルに入力する画像を読み込ませるエリア、右のパネルにgr.outputsに指定した出力されるものを表示するエリアが描画されている。

スクリーンショット 2020-11-23 18.18.08.png

試しに、マンドリルの画像を拡大してみる。入力する画像は、ドラックアンドドロップでinputエリアに持っていく。数秒待つと「OUTPUT IMAGE」のエリアに、処理済みの画像が表示される。右クリックから画像を保存を選択し、単純に同じサイズに拡大した画像と比較してみる。

左がFinderで同じサイズに拡大した画像、右がESRGANにより生成された画像である。 ESRGANで拡大したものの方が、毛並みの線や輪郭線がかなりシャープに表現されていることが確認できる。

終わりに

Gradioを使ってTensorflowのモデルを、簡単に公開できるWeb画面を作成できた。 Web画面の定義は、３つのモジュールに関する記述をするだけで良いので、慣れてきたら10分もかからずWebページが作成できそう。

GradioHubには、様々なサンプルがあるので、テキストや音声を利用した場合にどうなるかは、こちらも参照いただきたい。

Gradio | Hub

いちろう’s blog

すーぱーえんじにあ

Gradio+ESRGANを利用した画像を高解像度にするWeb画面の作成

概要