import gradio as gr
import onnxruntime as ort
import cv2
import numpy as np
from PIL import Image
import os
import uuid

# 啟動 ONNX 模型 Session
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if ort.get_device() == 'GPU' else ['CPUExecutionProvider']
session = ort.InferenceSession('AnimeGANv3_Hayao_STYLE_36.onnx', providers=providers)

# 圖片處理
def process_image(img):
    h, w = img.shape[:2]
    def to_8s(x): return 256 if x < 256 else x - x % 8
    img = cv2.resize(img, (to_8s(w), to_8s(h)))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 127.5 - 1.0
    return img

# 圖片轉換 + 下載功能
def cartoonize_image(input_image):
    img = np.array(input_image)
    h, w = img.shape[:2]
    processed = process_image(img)
    processed = np.expand_dims(processed, axis=0)

    input_name = session.get_inputs()[0].name
    output = session.run(None, {input_name: processed})[0]
    result = (np.squeeze(output) + 1.) / 2 * 255
    result = np.clip(result, 0, 255).astype(np.uint8)
    result = cv2.resize(result, (w, h))
    result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)

    result_image = Image.fromarray(result)

    # 保存下載檔案
    os.makedirs("outputs", exist_ok=True)
    filename = f"outputs/{uuid.uuid4().hex}.png"
    result_image.save(filename)

    return result_image, filename

# Gradio UI
interface = gr.Interface(
    fn=cartoonize_image,
    inputs=gr.Image(type="pil", label="上傳圖片"),
    outputs=[
        gr.Image(type="pil", label="宮崎駿風格轉換結果"),
        gr.File(label="⬇️ 下載圖片")
    ],
    title="🎨 AnimeGANv3 - 宮崎駿風格轉換器",
    description="上傳圖片，將自動轉換為宮崎駿動畫風格！",
    allow_flagging="never"
)

interface.launch()