rinogeek's picture
Initial Docker configuration for Hugging Face
1624b73
Raw
History Blame Contribute Delete
1.51 kB
import cv2
import numpy as np
import os
from django.conf import settings
from ultralytics import YOLO
# Global model variable
_model = None
def get_model():
global _model
if _model is None:
MODEL_PATH = os.path.join(settings.BASE_DIR, 'best.pt')
try:
_model = YOLO(MODEL_PATH)
except Exception as e:
print(f"Error loading model: {e}")
return _model
def run_detection(image_path):
"""
Runs YOLO detection on an image file path and returns a list of detections.
"""
model = get_model()
if model is None:
return []
img = cv2.imread(image_path)
if img is None:
return []
results = model(img)
detections = []
for r in results:
boxes = r.boxes
for box in boxes:
# Get coordinates in percentage for the frontend
x_center, y_center, w, h = box.xywhn[0].tolist()
x = (x_center - w/2) * 100
y = (y_center - h/2) * 100
width = w * 100
height = h * 100
conf = float(box.conf[0])
cls = int(box.cls[0])
label = model.names[cls]
detections.append({
"id": len(detections),
"x": x,
"y": y,
"width": width,
"height": height,
"label": label,
"confidence": conf * 100
})
return detections