app.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import gradio as gr
from simple_inference import predict_location
import os

# 设置默认模型和引擎
DEFAULT_MODEL = "TheEighthDay/SeekWorld_RL_PLUS"
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"Error: {str(e)}"

# 创建Gradio界面
demo = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Textbox(label="Model Name", value=DEFAULT_MODEL),
        gr.Dropdown(choices=["transformers", "vllm"], value=DEFAULT_ENGINE, label="Inference Engine")
    ],
    outputs=gr.Textbox(label="Result"),
    title="Geolocation Reasoning App",
    description="Upload an image, and the system will identify the country and region where the image was taken."
)

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