import gradio as gr
from PIL import Image
from detector import FakeImageDetector  

FIXED_THRESHOLD = 0.4999

print("正在初始化检测器，请稍候...")
try:
    detector = FakeImageDetector()
    
    print("检测器初始化完成，Web 服务准备就绪。")
    models_loaded = True
except Exception as e:
    print(f"模型加载失败: {e}")
    models_loaded = False
    detector = None

def predict_image(input_image_numpy):
    """
    接收 Gradio 的输入 (numpy array)，调用检测器，并返回结果。
    """
    if not models_loaded or detector is None:
        return "错误：模型未能成功加载，请检查后台日志。", None

    pil_image = Image.fromarray(input_image_numpy)
    
    result_text, score = detector.detect(pil_image, FIXED_THRESHOLD)
    
    label_color = "red" if score > FIXED_THRESHOLD else "green"
    
    return result_text, gr.Label(value=f"{score:.10f}", label=label_color)


with gr.Blocks(title="伪造图像检测器", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 伪造图像检测器 (Fake Image Detector)
        上传一张图片，模型将判断其为 **真实的 (Real)** 还是 **AI 生成的伪造图像 (Fake)**。
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            # 输入组件
            image_input = gr.Image(type="numpy", label="上传图片", height=300)
            # threshold_slider = gr.Slider(
            #     minimum=0.495, maximum=0.55, value=0.499892068, step=0.0001, 
            #     label="检测门限 (Threshold)",
            #     info="得分低于此门限的图片被认为是伪造的"
            # )
            submit_btn = gr.Button("开始检测", variant="primary")
        
        with gr.Column(scale=1):
            # 输出组件
            result_output_text = gr.Textbox(label="检测结论", lines=2)
            # 这里我们用一个临时的 Label 来显示带颜色的分数
            result_output_score = gr.Label(label="模型原始得分")

    submit_btn.click(
        fn=predict_image, 
        inputs=[image_input], 
        outputs=[result_output_text, result_output_score]
    )

if not models_loaded:
    print("\n由于模型加载失败，Gradio Web服务无法启动。")
else:
    print("正在启动 Gradio 服务...")
    
    demo.launch(server_name="0.0.0.0")