Spaces:
Sleeping
Sleeping
from ultralytics import YOLO | |
import supervision as sv | |
def parse_detection(detections): | |
parsed_rows = [] | |
for i in range(len(detections.xyxy)): | |
x_min = float(detections.xyxy[i][0]) | |
y_min = float(detections.xyxy[i][1]) | |
x_max = float(detections.xyxy[i][2]) | |
y_max = float(detections.xyxy[i][3]) | |
width = int(x_max - x_min) | |
height = int(y_max - y_min) | |
row = { | |
"x": int(y_min), | |
"y": int(x_min), | |
"width": width, | |
"height": height, | |
"class_id": "" | |
if detections.class_id is None | |
else int(detections.class_id[i]), | |
"confidence": "" | |
if detections.confidence is None | |
else float(detections.confidence[i]), | |
"tracker_id": "" | |
if detections.tracker_id is None | |
else int(detections.tracker_id[i]), | |
} | |
if hasattr(detections, "data"): | |
for key, value in detections.data.items(): | |
if key == "class_name": | |
key = "class" | |
row[key] = ( | |
str(value[i]) | |
if hasattr(value, "__getitem__") and value.ndim != 0 | |
else str(value) | |
) | |
parsed_rows.append(row) | |
return parsed_rows | |
model = YOLO("models/best_v2.pt", task="detect") | |
results = model(["data/IMG_0050.jpg"])[0] | |
width, height = results.orig_shape[1], results.orig_shape[0] | |
print(results.orig_shape) | |
print(results.speed) | |
output = sv.Detections.from_ultralytics(results) | |
output = parse_detection(output) | |
parse_result = {'predictions': output, 'image': {'width': width, 'height': height}} | |
print(parse_result) |