handler.py

import os
import json
import base64
import tempfile
from io import BytesIO

import requests
from PIL import Image
from swift.llm import PtEngine, RequestConfig, get_model_tokenizer, get_template, InferRequest

MODEL_ID = "ghost233lism/GeoAgent"


class EndpointHandler:
    def __init__(self, path=""):
        self.model_path = MODEL_ID
        self.max_new_tokens = 2048
        self.engine = None
        self.template = None
        self._load_model()

    def _load_model(self):
        model, tokenizer = get_model_tokenizer(
            self.model_path,
            model_type="qwen2_5_vl",
        )
        template_type = model.model_meta.template
        self.template = get_template(
            template_type,
            tokenizer,
            default_system=self._system_prompt(),
        )
        self.engine = PtEngine.from_model_template(
            model,
            self.template,
            max_batch_size=1,
        )

    @staticmethod
    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 _save_temp_image(self, image_url=None, image_base64=None):
        if image_url:
            resp = requests.get(image_url, timeout=30)
            resp.raise_for_status()
            img = Image.open(BytesIO(resp.content)).convert("RGB")
        elif image_base64:
            if image_base64.startswith("data:image"):
                _, image_base64 = image_base64.split(",", 1)
            img = Image.open(BytesIO(base64.b64decode(image_base64))).convert("RGB")
        else:
            raise ValueError("Please provide `inputs.image_url` or `inputs.image_base64`.")

        tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
        img.save(tmp.name, format="JPEG")
        tmp.close()
        return tmp.name

    def __call__(self, data):
        payload = data.get("inputs", {}) or {}

        image_url = payload.get("image_url")
        image_base64 = payload.get("image_base64")
        prompt = payload.get(
            "prompt",
            "Based on the image, tell me the specific location and your thinking process",
        )
        max_new_tokens = int(payload.get("max_new_tokens", self.max_new_tokens))

        image_path = self._save_temp_image(image_url=image_url, image_base64=image_base64)

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

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

            resp = self.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

            return {
                "generated_text": text,
                "raw_response": str(resp),
            }
        finally:
            if os.path.exists(image_path):
                os.remove(image_path)