Create predict.py

predict.py

@@ -0,0 +1,27 @@
+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