Ayomide
56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
from ultralytics import YOLO
|
|
import cv2
|
|
from pathlib import Path
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class MemoryDetector:
|
|
def __init__(self, model_path):
|
|
try:
|
|
self.model = YOLO(model_path)
|
|
logger.info(f"Loaded model from {model_path}")
|
|
except Exception as e:
|
|
logger.error(f"Model loading failed: {str(e)}")
|
|
raise
|
|
|
|
def detect(self, image_path):
|
|
try:
|
|
# Run inference
|
|
results = self.model.predict(image_path, imgsz=416, conf=0.5)
|
|
|
|
# Extract results
|
|
boxes = results[0].boxes.xyxy.cpu().numpy()
|
|
confidences = results[0].boxes.conf.cpu().numpy()
|
|
|
|
# Convert to list of [x1, y1, x2, y2, confidence]
|
|
detections = []
|
|
for box, conf in zip(boxes, confidences):
|
|
detections.append({
|
|
'box': [int(x) for x in box],
|
|
'confidence': float(conf)
|
|
})
|
|
|
|
# Annotate image
|
|
annotated_img = self._draw_boxes(image_path, detections)
|
|
|
|
return {
|
|
'detections': detections,
|
|
'annotated_image': annotated_img
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Detection failed: {str(e)}")
|
|
raise
|
|
|
|
def _draw_boxes(self, image_path, detections):
|
|
img = cv2.imread(str(image_path))
|
|
for det in detections:
|
|
x1, y1, x2, y2 = det['box']
|
|
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
|
cv2.putText(img, f"{det['confidence']:.2f}",
|
|
(x1, y1-10),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 1)
|
|
return img
|