|
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: |
|
|
|
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} |
|
|