【Python】AIを使った写真のノイズ除去ツールを自作しよう

       
import os
import sys
sys.path.append(os.path.join(os.path. dirname(__file__), '../libs'))
from create_ui import CreateUI
import torch
from denise import *

# ボタンが押された時の処理
def func(param):
    input_folder = param["input_folder"]
    output_folder = param["output_folder"]
    model_combo = param["model_combo"]

    model_path = ""
    if(model_combo == "質感優先") :
        model_path = "scunet_color_real_gan"
    elif(model_combo == "ディテール優先") :
        model_path = "scunet_color_real_psnr"
    elif(model_combo == "ノイズ少") :
        model_path = "scunet_color_15"
    elif(model_combo == "ノイズ中") :
        model_path = "scunet_color_25"
    elif(model_combo == "ノイズ多") :
        model_path = "scunet_color_50"  
         
    denoise_image(model_path, input_folder, output_folder)

# インスタンス生成
ui = CreateUI(func)


# 画面の表示
ui.show()

import os
import time
import argparse
import numpy as np
import torch
from utils import utils_model, utils_image as util
from models.network_scunet import SCUNet as net

def load_model(model_name, model_zoo, device, n_channels=3):
    model_path = os.path.join(model_zoo, model_name + '.pth')
    model = net(in_nc=n_channels, config=[4, 4, 4, 4, 4, 4, 4], dim=64)
    model.load_state_dict(torch.load(model_path), strict=True)
    model.eval()
    
    for _, param in model.named_parameters():
        param.requires_grad = False

    return model.to(device), model_path

def process_images(model, input_folder, output_folder, device):
    util.mkdir(output_folder)

    img_paths = util.get_image_paths(input_folder)
    total_images = len(img_paths)

    print(f"処理開始: {total_images} 枚の画像を処理します。")

    start_time = time.time()

    for idx, img_path in enumerate(img_paths):
        img_name, ext = os.path.splitext(os.path.basename(img_path))
        img_L = util.imread_uint(img_path, n_channels=3)
        img_L = util.uint2single(img_L)

        img_L = util.single2tensor4(img_L).to(device)

        # 画像処理の開始時間を計測
        image_start_time = time.time()
        
        # メモリ管理を改善するため test_mode() を使用
        img_E = utils_model.test_mode(model, img_L, refield=64, mode=5)
        img_E = util.tensor2uint(img_E)

        output_path = os.path.join(output_folder, img_name + ext)
        util.imsave(img_E, output_path)
        
        # 画像処理の終了時間を計測
        image_end_time = time.time()
        processing_time = image_end_time - image_start_time

        print(f"{idx + 1}/{total_images} 処理完了: {img_name + ext} (処理時間: {processing_time:.2f} 秒)")

        # メモリ解放
        del img_L, img_E
        torch.cuda.empty_cache()

    end_time = time.time()
    total_processing_time = end_time - start_time

    print(f"全画像処理完了！総処理時間: {total_processing_time:.2f} 秒")

def denoise_image(model_name='scunet_color_real_gan', input_folder='', output_folder='', model_zoo='model_zoo'):
    """ Python プログラム内から直接 denoise を実行する関数 """
    if not input_folder or not output_folder:
        raise ValueError("input_folder と output_folder は指定する必要があります。")

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model, model_path = load_model(model_name, model_zoo, device)

    print(f"モデル使用: {model_name} (パス: {model_path})")
    process_images(model, input_folder, output_folder, device)

def denoise_batch():
    """ コマンドラインからバッチで denoise を実行する関数 """
    parser = argparse.ArgumentParser(description="バッチで画像を処理")
    parser.add_argument('--model_name', type=str, default='scunet_color_real_gan', help="使用するモデルの名前")
    parser.add_argument('--input_folder', type=str, required=True, help="入力画像のフォルダ")
    parser.add_argument('--output_folder', type=str, required=True, help="処理結果の保存先フォルダ")
    parser.add_argument('--model_zoo', type=str, default='model_zoo', help="モデルの保存場所")

    args = parser.parse_args()
    denoise(args.model_name, args.input_folder, args.output_folder, args.model_zoo)

if __name__ == '__main__':
    denoise_batch()

title:写真ノイズ除去ツール
mode:light
max:1024,300
size:600,300
min:600,300
input_folder:変換元フォルダ,フォルダ選択
{
    label:　　　画質選択
    model_combo:質感優先,ディテール優先,ノイズ少,ノイズ中,ノイズ多
}
output_folder:出力先フォルダ,フォルダ選択

exec_button:ノイズ除去の実行