File size: 2,009 Bytes
4bab068 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | """
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"
|