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