File size: 2,412 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 55 56 57 58 59 60 61 62 | """
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."
|