| """ |
| Thin wrapper around Ultralytics YOLO11. |
| |
| Why YOLO here at all, when the VLM can already "see" the image? |
| -> VLMs are known to hallucinate counts ("about 3 people") because they're |
| generating text, not measuring pixels. YOLO gives us *pixel-grounded*, |
| confidence-scored counts and boxes. We feed these counts into the VLM |
| prompt as ground truth context, and use the boxes for drawing on the |
| annotated image. This is a "detector-grounded VLM" pattern — cheap to |
| build, and meaningfully reduces hallucination vs a VLM-only pipeline. |
| """ |
|
|
| from ultralytics import YOLO |
| from .config import YOLO_MODEL_ID, RELEVANT_CLASSES, YOLO_CONF_THRESHOLD |
|
|
|
|
| class DisasterObjectDetector: |
| def __init__(self, model_id: str = YOLO_MODEL_ID): |
| self.model = YOLO(model_id) |
|
|
| def detect(self, image_path: str): |
| """Run detection, return list of dicts: {class, conf, box=[x1,y1,x2,y2]}""" |
| results = self.model.predict( |
| source=image_path, |
| conf=YOLO_CONF_THRESHOLD, |
| verbose=False, |
| )[0] |
|
|
| detections = [] |
| names = results.names |
| for box in results.boxes: |
| cls_id = int(box.cls.item()) |
| cls_name = names[cls_id] |
| if cls_name not in RELEVANT_CLASSES: |
| continue |
| conf = float(box.conf.item()) |
| x1, y1, x2, y2 = [float(v) for v in box.xyxy[0].tolist()] |
| detections.append({ |
| "class": cls_name, |
| "conf": conf, |
| "box": [x1, y1, x2, y2], |
| }) |
| return detections |
|
|
| @staticmethod |
| def summarize(detections): |
| """Turn a list of detections into {'people': 2, 'cars': 1, ...} counts.""" |
| counts = {} |
| for d in detections: |
| label = RELEVANT_CLASSES.get(d["class"], d["class"]) |
| counts[label] = counts.get(label, 0) + 1 |
| return counts |
|
|
| @staticmethod |
| def to_context_string(counts: dict) -> str: |
| """Human-readable string injected into the VLM prompt as grounding context.""" |
| if not counts: |
| return "No standard objects (people/vehicles) were confidently detected by the object detector." |
| parts = [f"{v} {k}" for k, v in counts.items()] |
| return "Object detector confidently found: " + ", ".join(parts) + \ |
| ". Treat these counts as reliable ground truth; do not contradict them." |
|
|
