| """ |
| Vision-Language reasoning engine, powered by Qwen2.5-VL-3B-Instruct. |
| |
| This is the "brain" of the app: it doesn't just caption the image, it is |
| prompted to reason like a triage officer and return a structured JSON |
| judgment (risk score + recommended actions), grounded with YOLO's object |
| counts to reduce hallucination. |
| """ |
|
|
| import json |
| import re |
| import os |
| from pathlib import Path |
|
|
| import torch |
| from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor |
| from qwen_vl_utils import process_vision_info |
|
|
| from .config import VLM_MODEL_ID, DEVICE, DTYPE, VLM_MAX_NEW_TOKENS |
|
|
| PROMPT_PATH = Path(__file__).parent.parent / "prompts" / "triage_prompt.txt" |
|
|
|
|
| class DisasterVLM: |
| def __init__(self, model_id: str = VLM_MODEL_ID): |
| self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained( |
| model_id, |
| torch_dtype=DTYPE, |
| device_map="auto" if DEVICE == "cuda" else None, |
| ) |
| if DEVICE != "cuda": |
| self.model.to(DEVICE) |
| self.processor = AutoProcessor.from_pretrained(model_id) |
| self.prompt_template = PROMPT_PATH.read_text() |
|
|
| def _build_prompt(self, yolo_context: str) -> str: |
| return self.prompt_template.format(yolo_context=yolo_context) |
|
|
| def analyze(self, image_path: str, yolo_context: str = ""): |
| """Returns (parsed_dict, raw_text). parsed_dict is {} if JSON parsing |
| fails completely (caller falls back to safe defaults).""" |
| prompt_text = self._build_prompt(yolo_context) |
|
|
| messages = [ |
| { |
| "role": "user", |
| "content": [ |
| {"type": "image", "image": image_path}, |
| {"type": "text", "text": prompt_text}, |
| ], |
| } |
| ] |
|
|
| chat_prompt = self.processor.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=True |
| ) |
| image_inputs, video_inputs = process_vision_info(messages) |
| inputs = self.processor( |
| text=[chat_prompt], |
| images=image_inputs, |
| videos=video_inputs, |
| padding=True, |
| return_tensors="pt", |
| ).to(self.model.device) |
|
|
| with torch.no_grad(): |
| generated_ids = self.model.generate( |
| **inputs, max_new_tokens=VLM_MAX_NEW_TOKENS, do_sample=False |
| ) |
| trimmed = [ |
| out_ids[len(in_ids):] |
| for in_ids, out_ids in zip(inputs.input_ids, generated_ids) |
| ] |
| raw_text = self.processor.batch_decode( |
| trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False |
| )[0] |
|
|
| return self._safe_parse_json(raw_text), raw_text |
|
|
| @staticmethod |
| def _safe_parse_json(text: str) -> dict: |
| """Models occasionally wrap JSON in ```json fences or add stray text. |
| Strip fences first, then fall back to extracting the first {...} block.""" |
| cleaned = text.strip() |
| cleaned = re.sub(r"^```(json)?", "", cleaned).strip() |
| cleaned = re.sub(r"```$", "", cleaned).strip() |
| try: |
| return json.loads(cleaned) |
| except json.JSONDecodeError: |
| pass |
|
|
| match = re.search(r"\{.*\}", cleaned, re.DOTALL) |
| if match: |
| try: |
| return json.loads(match.group(0)) |
| except json.JSONDecodeError: |
| pass |
| return {} |
|
|
