eisei-ai-space / test.py
vumichien's picture
Initial commit
87a77d3 verified
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)