app.py

import os
os.environ["USE_HF"] = "1"

import json
import tempfile
from typing import Optional

import gradio as gr
from PIL import Image
from swift.llm import (
    PtEngine,
    RequestConfig,
    get_model_tokenizer,
    get_template,
    InferRequest,
)

MODEL_ID = "ghost233lism/GeoAgent"

engine = None
template = None


def system_prompt():
    return """You are an expert with rich experience in the field of geolocation, skilled at accurately locating the geographic location of images through various clues in the images, such as traffic signs, architectural styles, natural landscapes, etc.
At the same time, you are also a mentor in building the chain of thought, able to organize complex ideas into clear and standardized patterns.
You possess knowledge in multiple disciplines such as geography, cartography, transportation, and architecture, and are able to identify the characteristics of different countries, regions, and locations. At the same time, you have the ability to analyze logic and construct a chain of thought.
Task: Output the thought chain and final answer based on the image input by the user. The thought chain includes:
Country Identification/Regional Guess/Precise Localization.

Possible clues include:
National clues: (Example: traffic sign shape/color, language and text, driving direction, architectural style, vegetation and climate characteristics, etc.)
Regional clues: (logo/enterprise, topography, vegetation type, regional traffic signs, dialect/spelling, license plate style, area code/postal code, infrastructure features, etc.)
Accurate positioning: (road sign text, street name, house number, landmark building, river and lake water system, place attributes such as park/city/commercial district, shop name and storefront, etc.)

Do not output objects that do not exist in the image.

Output strictly in JSON format:
{
  "ChainOfThought": {
    "CountryIdentification": {
      "Clues": [],
      "Reasoning": "",
      "Conclusion": "",
      "Uncertainty": ""
    },
    "RegionalGuess": {
      "Clues": [],
      "Reasoning": "",
      "Conclusion": "",
      "Uncertainty": ""
    },
    "PreciseLocalization": {
      "Clues": [],
      "Reasoning": "",
      "Conclusion": "",
      "Uncertainty": ""
    }
  },
  "FinalAnswer": "Country; Region; Specific Location"
}"""


def load_model():
    global engine, template
    if engine is not None:
        return

    model, tokenizer = get_model_tokenizer(
        MODEL_ID,
        model_type="qwen2_5_vl",
    )
    template_type = model.model_meta.template
    template = get_template(
        template_type,
        tokenizer,
        default_system=system_prompt(),
    )
    engine = PtEngine.from_model_template(
        model,
        template,
        max_batch_size=1,
    )


def run_geoagent(image: Optional[Image.Image], prompt: str, max_new_tokens: int):
    if image is None:
        return "请先上传图片。"

    load_model()

    if not prompt.strip():
        prompt = "Based on the image, tell me the specific location and your thinking process"

    tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
    image = image.convert("RGB")
    image.save(tmp.name, format="JPEG")
    image_path = tmp.name
    tmp.close()

    try:
        infer_request = InferRequest(
            messages=[
                {
                    "role": "user",
                    "content": prompt,
                }
            ],
            images=[image_path],
        )

        request_config = RequestConfig(
            max_tokens=int(max_new_tokens),
            temperature=0,
        )

        resp = engine.infer([infer_request], request_config)[0]

        text = ""
        if hasattr(resp, "choices") and resp.choices:
            choice = resp.choices[0]
            if hasattr(choice, "message") and hasattr(choice.message, "content"):
                text = choice.message.content
            elif hasattr(choice, "text"):
                text = choice.text

        if not text:
            text = str(resp)

        # 如果返回的是 JSON 字符串，尽量格式化一下
        try:
            parsed = json.loads(text)
            text = json.dumps(parsed, ensure_ascii=False, indent=2)
        except Exception:
            pass

        return text
    except Exception as e:
        return f"推理报错：{str(e)}"
    finally:
        if os.path.exists(image_path):
            os.remove(image_path)


with gr.Blocks() as demo:
    gr.Markdown("# GeoAgent API Demo")
    gr.Markdown("上传一张图片，调用 GeoAgent 做地理定位推理。")

    with gr.Row():
        image_in = gr.Image(type="pil", label="输入图片")
        output = gr.Textbox(label="输出结果", lines=24)

    prompt_in = gr.Textbox(
        label="提示词",
        value="Based on the image, tell me the specific location and your thinking process",
    )
    max_tokens_in = gr.Slider(
        minimum=128,
        maximum=4096,
        value=2048,
        step=128,
        label="max_new_tokens",
    )

    btn = gr.Button("开始推理")
    btn.click(
        fn=run_geoagent,
        inputs=[image_in, prompt_in, max_tokens_in],
        outputs=output,
        api_name="predict",
    )

demo.launch()