Spaces:
Paused
Paused
File size: 1,030 Bytes
12f2295 |
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 |
# prescription_validation/zone_detector.py
import cv2
import numpy as np
from ultralytics import YOLO
class ZoneDetector:
def __init__(self, model_path: str = "models/signature_model.pt"):
try:
self.model = YOLO(model_path)
print(f"✅ Loaded local YOLO model from '{model_path}'")
except Exception as e:
self.model = None
print(f"❌ Failed to load local YOLO model from '{model_path}'. Please ensure the model file exists.")
def detect(self, img: np.ndarray) -> list[dict]:
if not self.model: return []
rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
results = self.model(rgb_img, verbose=False)
detections = []
for box in results[0].boxes:
x1, y1, x2, y2 = box.xyxy[0].tolist()
class_id = int(box.cls[0])
detections.append({
"label": self.model.names[class_id],
"bbox": tuple(map(int, [x1, y1, x2, y2])),
})
return detections
|