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