import os

def replace_import_line(file_path):
    # 置換対象の行
    old_line = "from torchvision.transforms.functional_tensor import rgb_to_grayscale"
    # 新しい行
    new_line = "from torchvision.transforms.functional import rgb_to_grayscale"

    # ファイルが存在するかチェック
    if not os.path.exists(file_path):
        print(f"ファイルが見つかりません: {file_path}")
        return

    # ファイルを読み込んで内容を置換
    with open(file_path, 'r') as file:
        lines = file.readlines()

    # 置換処理
    updated_lines = [line.replace(old_line, new_line) if old_line in line else line for line in lines]

    # ファイルに書き込み
    with open(file_path, 'w') as file:
        file.writelines(updated_lines)

    print(f"ファイルの置換が完了しました: {file_path}")

# 書き換えるファイルのパスを指定
file_path = "/usr/local/lib/python3.10/site-packages/basicsr/data/degradations.py"
replace_import_line(file_path)
import gradio as gr
import cv2
import numpy as np
from gfpgan import GFPGANer

# GFPGANerの初期化
restorer = GFPGANer(model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.4.pth',
                    upscale=2, arch='clean', channel_multiplier=2)

# 画像補正関数
def restore_faces(input_img):
    # 入力画像を処理
    img = np.array(input_img)
    
    # GFPGANのenhanceメソッドで顔を補正
    _, restored_img, _ = restorer.enhance(img, has_aligned=False, only_center_face=False)
    
    # 補正された画像 (restored_img) を返す
    return restored_img

# Gradioインターフェースの作成
iface = gr.Interface(fn=restore_faces, 
                     inputs=gr.Image(type="numpy"), 
                     outputs="image",
                     title="AI生成顔をGFPGANで補正",
                     description="AIで生成された顔画像をGFPGANを使って補正します。")

# アプリケーションを起動
iface.launch()