Ayomide
137 lines
4.8 KiB
Python
137 lines
4.8 KiB
Python
import cv2
|
|
import numpy as np
|
|
import logging
|
|
from pathlib import Path
|
|
from ultralytics import YOLO
|
|
from exceptions import ModelLoadError, DetectionError, ImageProcessingError
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class MemoryDetector:
|
|
def __init__(self, model_path, confidence_threshold=0.3, image_size=416):
|
|
self.model_path = model_path
|
|
self.confidence_threshold = confidence_threshold
|
|
self.image_size = image_size
|
|
self.model = None
|
|
self._load_model()
|
|
|
|
def _load_model(self):
|
|
"""Load YOLO model with error handling."""
|
|
try:
|
|
if not Path(self.model_path).exists():
|
|
raise FileNotFoundError(f"Model file not found: {self.model_path}")
|
|
|
|
logger.info(f"Loading model from {self.model_path}")
|
|
self.model = YOLO(self.model_path)
|
|
logger.info("Model loaded successfully")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to load model: {e}")
|
|
raise ModelLoadError(f"Model loading failed: {str(e)}")
|
|
|
|
def detect_from_bytes(self, image_bytes):
|
|
"""Detect memory modules from image bytes."""
|
|
try:
|
|
# Decode image
|
|
nparr = np.frombuffer(image_bytes, np.uint8)
|
|
image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
|
|
|
if image is None:
|
|
raise ImageProcessingError("Could not decode image")
|
|
|
|
return self._perform_detection(image)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Detection from bytes failed: {e}")
|
|
raise DetectionError(f"Detection failed: {str(e)}")
|
|
|
|
def detect_from_file(self, file_path):
|
|
"""Detect memory modules from image file."""
|
|
try:
|
|
if not Path(file_path).exists():
|
|
raise FileNotFoundError(f"Image file not found: {file_path}")
|
|
|
|
image = cv2.imread(str(file_path))
|
|
if image is None:
|
|
raise ImageProcessingError("Could not load image file")
|
|
|
|
return self._perform_detection(image)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Detection from file failed: {e}")
|
|
raise DetectionError(f"Detection failed: {str(e)}")
|
|
|
|
def _perform_detection(self, image):
|
|
"""Perform detection on image."""
|
|
try:
|
|
logger.info("Running detection")
|
|
|
|
# Run inference
|
|
results = self.model.predict(
|
|
image,
|
|
imgsz=self.image_size,
|
|
conf=self.confidence_threshold,
|
|
verbose=False
|
|
)
|
|
|
|
# Extract detections
|
|
detections = self._extract_detections(results[0])
|
|
|
|
# Create annotated image
|
|
annotated_image = self._draw_boxes(image, detections)
|
|
|
|
logger.info(f"Detection completed: {len(detections)} objects found")
|
|
|
|
return {
|
|
'detections': detections,
|
|
'annotated_image': annotated_image,
|
|
'detection_count': len(detections)
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Detection processing failed: {e}")
|
|
raise DetectionError(f"Detection processing failed: {str(e)}")
|
|
|
|
def _extract_detections(self, result):
|
|
"""Extract detection results."""
|
|
detections = []
|
|
|
|
if result.boxes is not None:
|
|
boxes = result.boxes.xyxy.cpu().numpy()
|
|
confidences = result.boxes.conf.cpu().numpy()
|
|
classes = result.boxes.cls.cpu().numpy() if result.boxes.cls is not None else None
|
|
|
|
for i, (box, conf) in enumerate(zip(boxes, confidences)):
|
|
detection = {
|
|
'box': [float(coord) for coord in box], # [x1, y1, x2, y2]
|
|
'confidence': float(conf),
|
|
'class': int(classes[i]) if classes is not None else 0
|
|
}
|
|
detections.append(detection)
|
|
|
|
return detections
|
|
|
|
def _draw_boxes(self, image, detections):
|
|
"""Draw bounding boxes on image."""
|
|
annotated = image.copy()
|
|
|
|
for detection in detections:
|
|
box = detection['box']
|
|
confidence = detection['confidence']
|
|
|
|
# Extract coordinates
|
|
x1, y1, x2, y2 = map(int, box)
|
|
|
|
# Draw bounding box
|
|
cv2.rectangle(annotated, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
|
|
|
# Draw confidence score
|
|
label = f"Memory: {confidence:.2f}"
|
|
cv2.putText(
|
|
annotated, label, (x1, y1 - 10),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1
|
|
)
|
|
|
|
return annotated
|