src/pipeline.py

"""
TriagePipeline: the single entrypoint the UI (or any batch script / API)
calls. Wires together:

    image -> YOLO11 (grounding)
          -> Qwen2.5-VL (reasoning, grounded by YOLO context)
          -> TriageResult (validated, clamped, defaulted)
          -> OpenCV annotated image

Models are loaded once (lazy singletons) so repeated Gradio calls are fast.
"""

import cv2
import numpy as np

from .object_detector import DisasterObjectDetector
from .vlm_engine import DisasterVLM
from .annotator import annotate
from .schema import TriageResult

_detector = None
_vlm = None


def get_detector() -> DisasterObjectDetector:
    global _detector
    if _detector is None:
        _detector = DisasterObjectDetector()
    return _detector


def get_vlm() -> DisasterVLM:
    global _vlm
    if _vlm is None:
        _vlm = DisasterVLM()
    return _vlm


def run_triage(image_path: str):
    """
    Returns:
        annotated_image_rgb (np.ndarray, HxWx3, RGB) - for Gradio display
        result (TriageResult)
    """
    detector = get_detector()
    vlm = get_vlm()

    # 1. Grounded object detection
    detections = detector.detect(image_path)
    counts = detector.summarize(detections)
    yolo_context = detector.to_context_string(counts)

    # 2. VLM scene reasoning, grounded with YOLO context
    parsed_json, raw_text = vlm.analyze(image_path, yolo_context=yolo_context)
    result = TriageResult.from_dict(parsed_json, yolo_detections=counts, raw_model_output=raw_text)

    # If the VLM failed to report people count but YOLO found people, trust YOLO.
    if result.people_visible_count == 0 and counts.get("people", 0) > 0:
        result.people_visible_count = counts["people"]

    # 3. Annotate image (boxes + risk badge)
    image_bgr = cv2.imread(image_path)
    annotated_bgr = annotate(
        image_bgr, detections, result.risk_score, result.risk_level, result.risk_color
    )
    annotated_rgb = cv2.cvtColor(annotated_bgr, cv2.COLOR_BGR2RGB)

    return annotated_rgb, result