File size: 647 Bytes
4efaadb | 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 | import torch
from ultralytics import YOLO
import cv2
import numpy as np
# Load the model
model = YOLO("best.pt")
def predict(image_path):
# Load image
image = cv2.imread(image_path)
# Inference
results = model(image)
# Parse results into CVAT format (adjust as needed)
annotations = []
for box in results[0].boxes.data:
x1, y1, x2, y2, confidence, class_id = box[:6]
annotations.append({
"x1": int(x1), "y1": int(y1),
"x2": int(x2), "y2": int(y2),
"confidence": float(confidence),
"class": int(class_id)
})
return annotations
|