app.py

import gradio as gr
from simple_inference import predict_location
import os

# 设置默认模型和引擎
DEFAULT_MODEL = "Qwen/Qwen2.5-VL-7B-Instruct"
DEFAULT_ENGINE = "transformers"

def process_image(image, model_name, inference_engine):
    """
    处理上传的图片并返回预测结果
    """
    # 保存上传的图片到临时文件
    temp_path = "temp_image.jpg"
    image.save(temp_path)
    
    try:
        # 调用预测函数
        result = predict_location(
            image_path=temp_path,
            model_name=model_name,
            inference_engine=inference_engine
        )
        
        # 删除临时文件
        os.remove(temp_path)
        
        return result
    except Exception as e:
        # 删除临时文件
        if os.path.exists(temp_path):
            os.remove(temp_path)
        return f"发生错误: {str(e)}"

# 创建Gradio界面
with gr.Blocks(title="位置识别系统") as demo:
    gr.Markdown("# 🏞️ 图片位置识别系统")
    gr.Markdown("上传一张图片，系统将识别图片拍摄的国家和地区。")
    
    with gr.Row():
        with gr.Column():
            # 图片上传组件
            image_input = gr.Image(type="pil", label="上传图片")
            
            # 模型选择
            model_name = gr.Textbox(
                label="模型名称",
                value=DEFAULT_MODEL,
                info="输入Hugging Face模型ID或本地模型路径"
            )
            
            # 推理引擎选择
            inference_engine = gr.Dropdown(
                choices=["transformers", "vllm"],
                value=DEFAULT_ENGINE,
                label="推理引擎"
            )
            
            # 提交按钮
            submit_btn = gr.Button("开始识别", variant="primary")
        
        with gr.Column():
            # 结果显示
            output = gr.Textbox(
                label="识别结果",
                lines=10,
                placeholder="识别结果将显示在这里..."
            )
    
    # 设置提交按钮的回调函数
    submit_btn.click(
        fn=process_image,
        inputs=[image_input, model_name, inference_engine],
        outputs=output
    )
    
    # 添加示例
    gr.Examples(
        examples=[
            ["examples/example1.jpg", DEFAULT_MODEL, DEFAULT_ENGINE],
            ["examples/example2.jpg", DEFAULT_MODEL, DEFAULT_ENGINE],
        ],
        inputs=[image_input, model_name, inference_engine],
        outputs=output,
        fn=process_image,
        cache_examples=True,
    )

# 启动应用
if __name__ == "__main__":
    demo.launch(share=True)