import gradio as gr
from image_debug import analysis_image, compare
import numpy as np
import base64
from PIL import Image
import io


def numpy_to_base64(img_array):
    # 将 numpy array 转换为 PIL Image
    img_pil = Image.fromarray(img_array)

    # 创建一个字节流
    buffer = io.BytesIO()
    # 将图片保存到字节流
    img_pil.save(buffer, format="PNG")
    # 获取 base64 字符串
    img_str = base64.b64encode(buffer.getvalue()).decode()

    return img_str


def combine_analysis(result1, result2):
    if not result1 or not result2:
        return "请先完成两张图片的分析"
    result = compare(result1, result2)
    return result

def process_image1(image1):
    if image1 is None:
        return "请上传图片1"
    image1 = numpy_to_base64(image1)
    return analysis_image(image1)


def process_image2(image2):
    if image2 is None:
        return "请上传图片2"
    image2 = numpy_to_base64(image2)
    return analysis_image(image2)


with gr.Blocks(title="图片分析演示") as demo:
    gr.Markdown("# 双图片分析系统")

    with gr.Row():
        with gr.Column():
            image1_input = gr.Image(label="第一张图片")
            result1 = gr.Markdown(label="第一张图片分析结果")
            # 删除 btn1

        with gr.Column():
            image2_input = gr.Image(label="第二张图片")
            result2 = gr.Markdown(label="第二张图片分析结果")
            # 删除 btn2

    with gr.Row():
        combine_btn = gr.Button("综合分析", variant="primary")
    with gr.Row():
        final_result = gr.Markdown(label="综合分析结果")

    # 替换 btn1.click 为 image1_input.change
    image1_input.change(
        fn=process_image1,
        inputs=image1_input,
        outputs=result1
    )

    # 替换 btn2.click 为 image2_input.change
    image2_input.change(
        fn=process_image2,
        inputs=image2_input,
        outputs=result2
    )

    combine_btn.click(
        fn=combine_analysis,
        inputs=[result1, result2],
        outputs=final_result
    )

if __name__ == "__main__":
    demo.launch()