Waifu2xで画像を一括リサイズ[Google Colab & Google Drive]

Zipファイルに保管した画像を一括でリサイズしたかったので。
Google Driveにzipファイルを保管し、ドライブ上に変換後のZipファイルを生成します。

GoogleDriveのルートディレクトリに[StableDiffusion]フォルダを作成していただき、一括リサイズしたいzipファイルを設置してください。

ランタイムの実行後、リサイズするフォルダを選択するフォームが表示されるので、該当する番号を入力するとファイルが出力されます。

# -*- coding: utf-8 -*-
"""Waifu2x

Automatically generated by Colaboratory.

"""

!apt-get install -y cmake libopencv-dev
!git clone https://github.com/DeadSix27/waifu2x-converter-cpp.git
!cd waifu2x-converter-cpp && mkdir out && cd out && cmake .. && make

!/content/waifu2x-converter-cpp/out/waifu2x-converter-cpp --help

from google.colab import drive, files
import os
import subprocess
from tqdm.notebook import tqdm
import zipfile
import shutil

# Google Driveをマウント
drive.mount('/content/drive')

# StableDiffusionフォルダの存在確認、なければ作成
stable_diffusion_path = '/content/drive/MyDrive/StableDiffusion'
if not os.path.exists(stable_diffusion_path):
    os.makedirs(stable_diffusion_path)
    print(f"StableDiffusionフォルダが作成されました。ZIPファイルをこのフォルダにアップロードしてください。")
    files.upload()
else:
    zip_files = [f for f in os.listdir(stable_diffusion_path) if f.endswith('.zip')]

    # ユーザーにZIPファイルを選択させる
    print("以下のZIPファイルが見つかりました。選択してください：")
    for i, zip_file in enumerate(zip_files):
        print(f"{i+1}. {zip_file}")

    selected_index = int(input("選択したいZIPファイルの番号を入力してください：")) - 1
    selected_zip_file = zip_files[selected_index]

# 入出力パス設定
drive_path = os.path.join(stable_diffusion_path, selected_zip_file)
input_dir = '/content/input_images'
output_dir = '/content/output_images'
waifu2x_path = '/content/waifu2x-converter-cpp/out/waifu2x-converter-cpp'

os.makedirs(input_dir, exist_ok=True)
os.makedirs(output_dir, exist_ok=True)

# ZIPを解凍
with zipfile.ZipFile(drive_path, 'r') as zip_ref:
    zip_ref.extractall(input_dir)

# 画像を拡大
for image_file in tqdm(os.listdir(input_dir), desc='Processing images'):
    input_path = os.path.join(input_dir, image_file)
    output_path = os.path.join(output_dir, image_file)

    model_dir = "/content/waifu2x-converter-cpp/models_rgb/"  # 実際のモデルファイルのパスに置き換えてください

    # waifu2xで拡大
    result = subprocess.run(
        [waifu2x_path, '-i', input_path, '-o', output_path, '--scale-ratio', '2', '--model-dir', model_dir],
        capture_output=True,
        text=True
    )

    # デバッグ情報を表示したいときはコメントアウトを外す
    # if result.returncode != 0:
    #     print(f"Error processing {image_file}: {result.stderr}")
    # else:
    #     print(f"Successfully processed {image_file}. Output: {result.stdout}")


# 処理が完了したら、出力ディレクトリ内の画像をZIPにまとめる
output_zip_path = os.path.join(stable_diffusion_path, "enlarged_images.zip")
with zipfile.ZipFile(output_zip_path, 'w') as zipf:
    for root, _, files in os.walk(output_dir):
        for file in files:
            zipf.write(os.path.join(root, file), file)

# ZIPファイル作成後、一時ディレクトリを削除
shutil.rmtree(input_dir)
shutil.rmtree(output_dir)

print(f"画像の拡大が完了しました。出力ZIPファイルはGoogle DriveのStableDiffusionフォルダに保存されています：{output_zip_path}")