import gradio as gr
import requests
import json
import base64

def call_api(image, text_input):
    # Check if image is provided
    if image is None:
        return "请上传图片"

    # Handle image processing
    try:
        # Read image file
        if isinstance(image, str):
            # If it's a filepath
            with open(image, "rb") as f:
                image_data = f.read()
        else:
            # If it's already in memory (Gradio might pass different types)
            image_data = image
        
        # Convert to base64
        image_base64 = base64.b64encode(image_data).decode("utf-8")
        image_url = f"data:image/jpeg;base64,{image_base64}"

        # Construct request payload
        payload = {
            "model": "/data1/models/PKUAgri/qwen2_vl_lora_sft",
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "image_url", "image_url": {"url": image_url}},
                        {"type": "text", "text": text_input}
                    ]
                }
            ]
        }

        headers = {
            "Content-Type": "application/json"
        }

        # Make API call
        response = requests.post(
            "http://10.1.10.115:4001/v1/chat/completions",
            data=json.dumps(payload),
            headers=headers,
            timeout=30
        )
        
        if response.status_code == 200:
            res = response.json()
            return res["choices"][0]["message"]["content"]
        else:
            return f"错误: {response.status_code} - {response.text}"
            
    except Exception as e:
        return f"处理请求时出错: {str(e)}"


# Build Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## 🌄 病虫害识别模型")

    with gr.Row():
        image_input = gr.Image(type="filepath", label="上传图片")
        text_input = gr.Textbox(label="请输入你的问题", placeholder="例如：图中有什么？")

    submit_btn = gr.Button("提交")
    output = gr.Textbox(label="回答", lines=5)

    submit_btn.click(fn=call_api, inputs=[image_input, text_input], outputs=output)

# Launch the app
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port= 40011, share=True)