src/config.py

"""
Central configuration for the Disaster Triage pipeline.
Change model IDs / thresholds here — nothing else in the codebase
should hardcode these values.
"""

import torch

# ---------------------------------------------------------------------------
# Models
# ---------------------------------------------------------------------------
VLM_MODEL_ID = "Qwen/Qwen2.5-VL-3B-Instruct"   # swap to -7B-Instruct if you have >=16GB VRAM
YOLO_MODEL_ID = "yolo11s.pt"                    # nano="yolo11n.pt" for speed, small="yolo11s.pt" for accuracy

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32

# Max new tokens for the VLM's JSON response. Keep tight -> faster + less
# chance of the model rambling outside the JSON schema.
VLM_MAX_NEW_TOKENS = 700

# ---------------------------------------------------------------------------
# YOLO -> disaster-relevant classes (subset of COCO-80)
# These are the classes we bother reporting to the operator / feeding to VLM.
# ---------------------------------------------------------------------------
RELEVANT_CLASSES = {
    "person": "people",
    "car": "cars",
    "truck": "trucks",
    "bus": "buses",
    "motorcycle": "motorcycles",
    "bicycle": "bicycles",
    "boat": "boats",
    "traffic light": "traffic lights",
    "fire hydrant": "fire hydrants",
}
YOLO_CONF_THRESHOLD = 0.35

# ---------------------------------------------------------------------------
# Risk scoring -> level mapping (used if the VLM omits risk_level)
# ---------------------------------------------------------------------------
RISK_LEVELS = [
    (0, 25, "LOW", "#2e7d32"),        # green
    (25, 50, "MODERATE", "#f9a825"),  # amber
    (50, 75, "HIGH", "#ef6c00"),      # orange
    (75, 101, "CRITICAL", "#c62828"), # red
]

def risk_level_from_score(score: int):
    for lo, hi, label, color in RISK_LEVELS:
        if lo <= score < hi:
            return label, color
    return "UNKNOWN", "#616161"