レガシ-AIシリーズ(３)　Anime2Sketch アニメ画を線画に変換

2024年8月13日 22:07

レガシ-AIシリーズ(2)はアニメ画をスケッチ（線画）に変換するAIです。今では同様の機能が画像生成AIのGUIに組み込まれてしまいましたが、今回は単独で動くAIとして、コードを書き直しました。
参照元は以下の通り。こちらにはデモもあります。

環境設定

仮想環境作成

python -m venv a2s
source a2s/bin/activate

PyTorchのインストール(CPU版）

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

gradioのインストール

pip install gradio

今回はGPUを使わないので簡単です。

ソースコードとウエイトを配置

a2sディレクトリが作成されているはずなので、以下の2つのソースコードを配置します。
weightsディレクトリを作成し、以下からウエイトをダウンロードし配置します。

a2sディレクイトリ構成

ソースコード本体 (ファイル名　as_gui.py)

初期化
必要なモジュールをimport
モデルのロード（今回はCPUで動かします）

def process_image(image, load_size)
Unetを作成（中程）
イメージをロード
スケッチに変換
gradio表示用にPILに変換

def tensor_to_img(input_image, imtype=np.uint8)
tensorをndartrey形式に変換

with gr.Blocks() as demo:
gradio_guiのコード

import gradio as gr
import os
from PIL import Image
import torchvision.transforms as transforms
import numpy as np
import torch
from model_u_a2s import create_model

# create model
wg='./weights/netG.pth'
device = torch.device('cpu')
model = create_model(wg).to(device)      # create a model given opt.model and other options
model.eval()

def process_image(image, load_size):
    transform_list = []
    transform_list += [transforms.ToTensor()]
    transform_list += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
    transform =transforms.Compose(transform_list)
    image = transform(image).unsqueeze(0) 
    
    img_skh = model(image.to(device))
    out_img =tensor_to_img(img_skh)
    out_img = np.repeat(out_img, 3, axis=-1)
    return Image.fromarray(out_img)

def tensor_to_img(input_image, imtype=np.uint8): #input_image (tensor)-  the input image tensor array /imtype (type)-  the desired type of the converted numpy array
    image_tensor = input_image.data
    image_numpy = image_tensor[0].cpu().float().numpy()  # convert it into a numpy array
    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0  # post-processing: tranpose and scaling
    return image_numpy.astype(imtype)

# Gradioインターフェースの設定
with gr.Blocks() as demo:
    gr.Markdown("## Image Processing with UnetGenerator")
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Input Image")
            load_size = gr.Number(label="Load Size (Set to 0 for no resize)", value=512)
            process_button = gr.Button("Process Image")
        with gr.Column():
            output_image = gr.Image(label="Processed Image")
    # 画像処理ボタンのイベント
    process_button.click(
        fn=process_image, 
        inputs=[image_input, load_size], 
        outputs=output_image
    )
#アプリケーションを起動
demo.launch()

モデルロードジェネレーションクラス (フィル名　model_u_a2s.py

class UnetGenerator(nn.Module):　ジェネレータ
Unet-based generatorです。torchで記述されています

def create_model(weight):　モデルをロード
モデルをロードしnetを返します

アプリの実行

 python as_gui.py

http://127.0.0.1:7860をブラウザで開きます

初期画面

変換したいイメージをアップロード

変換実行

Process Imageボタンで変換できます。

まとめ

今回はスケッチ変換AIを動かしました。とても軽いのでCPUで動きます。イメージを線画にする方法は色々とあると思いますが、これはとても手軽です。ネットワーク定義やモデルロードもベタに書かれているので、画像処理AIの理解にも役立つかもしれません。レガシーなAIですが、まだ3年前のこと、この後生成AIが爆発的に進歩して、このあたりで発表されたAIは忘れ去られていますが、それなりに使えるかもですね。

レガシ-AIシリーズ(３) Anime2Sketch アニメ画を線画に変換