import gradio as gr
import cv2
import numpy as np
from openvino.runtime import Core
from pathlib import Path
import tempfile

# ---------------- OpenVINO Upscale Function ----------------
def upscale_image_openvino(image, model_xml="RealESRGAN_x4plus_dynamic.xml", device="CPU"):
    """
    使用 OpenVINO 版本的 Real-ESRGAN 放大圖片，傳入 numpy array
    """
    # 建立 OpenVINO Core
    core = Core()
    compiled_model = core.compile_model(model_xml, device)
    input_layer = compiled_model.input(0)
    output_layer = compiled_model.output(0)

    # image 是 numpy array (H,W,C) RGB
    img = image.astype(np.float32) / 255.0
    img = np.transpose(img, (2, 0, 1))[np.newaxis, :]  # NCHW

    # 推論
    result = compiled_model([img])[output_layer]
    result = np.clip(result[0].transpose(1, 2, 0), 0, 1)
    result = (result * 255).astype(np.uint8)
    return result  # numpy array (H,W,C) RGB

# ---------------- Gradio 接口 ----------------
def gradio_upscale(image):
    if image is None:
        return None
    result = upscale_image_openvino(image)
    return result

# 建立 Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## OpenVINO Real-ESRGAN x4 放大器")

    with gr.Row():
        input_img = gr.Image(label="上傳圖片", type="numpy")
        output_img = gr.Image(label="放大後圖片", type="numpy")

    upscale_btn = gr.Button("realESRGANx4")
    upscale_btn.click(gradio_upscale, inputs=input_img, outputs=output_img)

# 啟動
demo.launch(share=True)