initial commit

backend/detector.py

@@ -0,0 +1,55 @@
+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