Spaces:
Paused
Paused
# 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 | |