File size: 838 Bytes
811ffe8 |
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 |
from ultralytics import YOLO
detection_model_id = "yolo11n.pt"
detection_model = YOLO(detection_model_id)
def run_detection(image_path: str, is_visualize: bool = False):
"""YOLOv11: return list of {box, label, score} for a single image."""
results = detection_model(image_path)
r = results[0]
detections = []
for box in r.boxes:
# box.xyxy is a 1×4 tensor, box.conf is a 1-element tensor, box.cls likewise
coords = box.xyxy.cpu().numpy().flatten().tolist()
score = float(box.conf.cpu().numpy().item())
cls_id = int(box.cls.cpu().numpy().item())
detections.append({
"box": coords,
"label": r.names[cls_id],
"score": score,
})
if is_visualize:
r.save()
r.show()
return {"detections": detections}
|