Update src/detection.py

src/detection.py

@@ -1,204 +1,303 @@
-import numpy as np
-from typing import List, Dict, Tuple
-import cv2
-from pathlib import Path
-import yaml
-
-class YOLOv11Detector:
-    """YOLOv11 detector for car damage detection"""
-
-    def __init__(self, config_path: str = "config.yaml"):
-        """Initialize YOLOv11 detector with configuration"""
-        with open(config_path, 'r') as f:
-            self.config = yaml.safe_load(f)
-
-        model_path = self.config['model']['path']
-
-        # Check which model file exists
-        if not Path(model_path).exists():
-            # Try to find available model files
-            model_dir = Path("models")
-            if (model_dir / "best.pt").exists():
-                model_path = str(model_dir / "best.pt")
-                print(f"Using best.pt model from training")
-            elif (model_dir / "last.pt").exists():
-                model_path = str(model_dir / "last.pt")
-                print(f"Using last.pt checkpoint model")
-            elif (model_dir / "best.onnx").exists():
-                model_path = str(model_dir / "best.onnx")
-                print(f"Using best.onnx model")
-            else:
-                raise FileNotFoundError(f"No model files found in models/ directory!")
-
-        self.model_path = model_path
-        self.device = self.config['model']['device']
-        self.confidence = self.config['model']['confidence']
-        self.iou_threshold = self.config['model']['iou_threshold']
-        self.classes = self.config['detection']['classes']
-
-        # Load model based on format
-        if model_path.endswith('.onnx'):
-            self._load_onnx_model()
-        else:  # .pt format
-            self._load_pytorch_model()
-
-    def _load_pytorch_model(self):
-        """Load PyTorch model using Ultralytics"""
-        from ultralytics import YOLO
-        self.model = YOLO(self.model_path)
-
-        # Set model to appropriate device
-        if self.device == 'cuda:0':
-            self.model.to('cuda')
-
-        print(f"Loaded PyTorch model: {self.model_path}")
-
-    def _load_onnx_model(self):
-        """Load ONNX model using OpenCV DNN"""
-        self.net = cv2.dnn.readNet(self.model_path)
-
-        # Set backend based on device
-        if self.device == 'cuda:0':
-            self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
-            self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
-        else:
-            self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
-            self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
-
-        print(f"Loaded ONNX model: {self.model_path}")
-
-    def detect(self, image: np.ndarray) -> Dict:
-        """
-        Perform detection on image
-
-        Args:
-            image: Input image as numpy array (BGR format)
-
-        Returns:
-            Dictionary containing detection results
-        """
-        if self.model_path.endswith('.onnx'):
-            return self._detect_onnx(image)
-        else:
-            return self._detect_pytorch(image)
-
-    def _detect_pytorch(self, image: np.ndarray) -> Dict:
-        """Detection using PyTorch model"""
-        # Run YOLO inference
-        results = self.model(
-            image,
-            conf=self.confidence,
-            iou=self.iou_threshold,
-            device=self.device,
-            verbose=False
-        )
-
-        # Parse results
-        detections = {
-            'boxes': [],
-            'confidences': [],
-            'classes': [],
-            'class_ids': []
-        }
-
-        if len(results) > 0 and results[0].boxes is not None:
-            boxes = results[0].boxes
-
-            for box in boxes:
-                # Get box coordinates (xyxy format)
-                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
-
-                # Get confidence and class
-                conf = float(box.conf[0].cpu().numpy())
-                cls_id = int(box.cls[0].cpu().numpy())
-
-                # Map class ID to class name
-                if cls_id < len(self.classes):
-                    cls_name = self.classes[cls_id]
-                else:
-                    cls_name = f"class_{cls_id}"
-
-                detections['boxes'].append([int(x1), int(y1), int(x2), int(y2)])
-                detections['confidences'].append(conf)
-                detections['classes'].append(cls_name)
-                detections['class_ids'].append(cls_id)
-
-        return detections
-
-    def _detect_onnx(self, image: np.ndarray) -> Dict:
-        """Detection using ONNX model (compatible with original code)"""
-        height, width = image.shape[:2]
-
-        # Preprocess image for ONNX
-        blob = cv2.dnn.blobFromImage(
-            image, 1/255.0, (640, 640),
-            swapRB=True, crop=False
-        )
-
-        self.net.setInput(blob)
-        preds = self.net.forward()
-        preds = preds.transpose((0, 2, 1))
-
-        # Extract outputs
-        detections = self._extract_onnx_output(
-            preds=preds,
-            image_shape=(height, width),
-            input_shape=(640, 640)
-        )
-
-        return detections
-
-    def _extract_onnx_output(self, preds: np.ndarray, image_shape: Tuple[int, int],
-                            input_shape: Tuple[int, int]) -> Dict:
-        """Extract detection results from ONNX model output"""
-        class_ids, confs, boxes = [], [], []
-
-        image_height, image_width = image_shape
-        input_height, input_width = input_shape
-        x_factor = image_width / input_width
-        y_factor = image_height / input_height
-
-        rows = preds[0].shape[0]
-        for i in range(rows):
-            row = preds[0][i]
-            conf = row[4]
-
-            classes_score = row[4:]
-            _, _, _, max_idx = cv2.minMaxLoc(classes_score)
-            class_id = max_idx[1]
-
-            if classes_score[class_id] > self.confidence:
-                confs.append(float(conf))
-                label = self.classes[int(class_id)] if int(class_id) < len(self.classes) else f"class_{class_id}"
-                class_ids.append(label)
-
-                # Extract boxes
-                x, y, w, h = row[0].item(), row[1].item(), row[2].item(), row[3].item()
-                left = int((x - 0.5 * w) * x_factor)
-                top = int((y - 0.5 * h) * y_factor)
-                width = int(w * x_factor)
-                height = int(h * y_factor)
-                box = [left, top, left + width, top + height]
-                boxes.append(box)
-
-        # Apply NMS
-        if len(boxes) > 0:
-            indices = cv2.dnn.NMSBoxes(
-                [[b[0], b[1], b[2]-b[0], b[3]-b[1]] for b in boxes],
-                confs, self.confidence, self.iou_threshold
-            )
-
-            if len(indices) > 0:
-                indices = indices.flatten()
-                return {
-                    'boxes': [boxes[i] for i in indices],
-                    'confidences': [confs[i] for i in indices],
-                    'classes': [class_ids[i] for i in indices],
-                    'class_ids': list(range(len(indices)))
-                }
-
-        return {'boxes': [], 'confidences': [], 'classes': [], 'class_ids': []}
-
-    def detect_batch(self, images: List[np.ndarray]) -> List[Dict]:
-        """Detect on multiple images"""
-        return [self.detect(img) for img in images]
+import numpy as np
+from typing import List, Dict, Tuple
+import cv2
+from pathlib import Path
+import yaml
+import torch
+import random
+import os
+
+class YOLOv11Detector:
+    """YOLOv11 detector for car damage detection with deterministic inference"""
+
+    def __init__(self, config_path: str = "config.yaml", deterministic: bool = True):
+        """Initialize YOLOv11 detector with configuration"""
+        
+        # Enable deterministic behavior
+        if deterministic:
+            self._set_deterministic()
+            
+        with open(config_path, 'r') as f:
+            self.config = yaml.safe_load(f)
+
+        model_path = self.config['model']['path']
+
+        # Check which model file exists
+        if not Path(model_path).exists():
+            # Try to find available model files
+            model_dir = Path("models")
+            if (model_dir / "best.pt").exists():
+                model_path = str(model_dir / "best.pt")
+                print(f"Using best.pt model from training")
+            elif (model_dir / "last.pt").exists():
+                model_path = str(model_dir / "last.pt")
+                print(f"Using last.pt checkpoint model")
+            elif (model_dir / "best.onnx").exists():
+                model_path = str(model_dir / "best.onnx")
+                print(f"Using best.onnx model")
+            else:
+                raise FileNotFoundError(f"No model files found in models/ directory!")
+
+        self.model_path = model_path
+        self.device = self.config['model']['device']
+        self.confidence = self.config['model']['confidence']
+        self.iou_threshold = self.config['model']['iou_threshold']
+        self.classes = self.config['detection']['classes']
+        self.deterministic = deterministic
+
+        # Load model based on format
+        if model_path.endswith('.onnx'):
+            self._load_onnx_model()
+        else:  # .pt format
+            self._load_pytorch_model()
+    
+    def _set_deterministic(self, seed: int = 42):
+        """Set deterministic behavior for reproducible results"""
+        print(f"Setting deterministic mode with seed: {seed}")
+        
+        # Set random seeds
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        
+        # Set CUDA deterministic settings
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed(seed)
+            torch.cuda.manual_seed_all(seed)
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+            
+            # Additional CUDA deterministic settings
+            os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
+            os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
+        
+        # Set PyTorch deterministic operations
+        torch.use_deterministic_algorithms(True, warn_only=True)
+        
+        # Set OpenCV random seed for ONNX inference
+        cv2.setRNGSeed(seed)
+
+    def _load_pytorch_model(self):
+        """Load PyTorch model using Ultralytics"""
+        from ultralytics import YOLO
+        
+        # Ensure deterministic loading
+        if self.deterministic:
+            torch.manual_seed(42)
+            
+        self.model = YOLO(self.model_path)
+
+        # Set model to appropriate device
+        if self.device == 'cuda:0' and torch.cuda.is_available():
+            self.model.to('cuda')
+        else:
+            self.model.to('cpu')
+            
+        # Set model to evaluation mode for consistent inference
+        if hasattr(self.model.model, 'eval'):
+            self.model.model.eval()
+
+        print(f"Loaded PyTorch model: {self.model_path}")
+        print(f"Model device: {next(self.model.model.parameters()).device}")
+
+    def _load_onnx_model(self):
+        """Load ONNX model using OpenCV DNN"""
+        self.net = cv2.dnn.readNet(self.model_path)
+
+        # Set backend based on device
+        if self.device == 'cuda:0':
+            self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
+            self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
+        else:
+            self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
+            self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
+
+        print(f"Loaded ONNX model: {self.model_path}")
+
+    def detect(self, image: np.ndarray) -> Dict:
+        """
+        Perform detection on image with deterministic behavior
+
+        Args:
+            image: Input image as numpy array (BGR format)
+
+        Returns:
+            Dictionary containing detection results
+        """
+        # Ensure deterministic preprocessing
+        if self.deterministic:
+            # Reset random seeds before each inference for consistency
+            if hasattr(torch, 'manual_seed'):
+                torch.manual_seed(42)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed(42)
+        
+        if self.model_path.endswith('.onnx'):
+            return self._detect_onnx(image)
+        else:
+            return self._detect_pytorch(image)
+
+    def _detect_pytorch(self, image: np.ndarray) -> Dict:
+        """Detection using PyTorch model with deterministic settings"""
+        
+        # Ensure model is in eval mode
+        if hasattr(self.model.model, 'eval'):
+            self.model.model.eval()
+        
+        # Disable gradients for inference
+        with torch.no_grad():
+            # Run YOLO inference with deterministic settings
+            results = self.model(
+                image,
+                conf=self.confidence,
+                iou=self.iou_threshold,
+                device=self.device,
+                verbose=False,
+                # Add deterministic parameters
+                augment=False,  # Disable test-time augmentation
+                half=False,     # Disable FP16 for consistency
+                max_det=1000    # Set consistent max detections
+            )
+
+        # Parse results
+        detections = {
+            'boxes': [],
+            'confidences': [],
+            'classes': [],
+            'class_ids': []
+        }
+
+        if len(results) > 0 and results[0].boxes is not None:
+            boxes = results[0].boxes
+
+            for box in boxes:
+                # Get box coordinates (xyxy format)
+                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
+
+                # Get confidence and class
+                conf = float(box.conf[0].cpu().numpy())
+                cls_id = int(box.cls[0].cpu().numpy())
+
+                # Map class ID to class name
+                if cls_id < len(self.classes):
+                    cls_name = self.classes[cls_id]
+                else:
+                    cls_name = f"class_{cls_id}"
+
+                detections['boxes'].append([int(x1), int(y1), int(x2), int(y2)])
+                detections['confidences'].append(conf)
+                detections['classes'].append(cls_name)
+                detections['class_ids'].append(cls_id)
+
+        # Sort results by confidence for consistency
+        if len(detections['boxes']) > 0:
+            # Create indices sorted by confidence (descending)
+            sorted_indices = sorted(range(len(detections['confidences'])), 
+                                  key=lambda i: detections['confidences'][i], reverse=True)
+            
+            # Reorder all detection lists
+            detections['boxes'] = [detections['boxes'][i] for i in sorted_indices]
+            detections['confidences'] = [detections['confidences'][i] for i in sorted_indices]
+            detections['classes'] = [detections['classes'][i] for i in sorted_indices]
+            detections['class_ids'] = [detections['class_ids'][i] for i in sorted_indices]
+
+        return detections
+
+    def _detect_onnx(self, image: np.ndarray) -> Dict:
+        """Detection using ONNX model with deterministic preprocessing"""
+        height, width = image.shape[:2]
+
+        # Deterministic preprocessing for ONNX
+        blob = cv2.dnn.blobFromImage(
+            image, 1/255.0, (640, 640),
+            swapRB=True, crop=False
+        )
+
+        self.net.setInput(blob)
+        preds = self.net.forward()
+        preds = preds.transpose((0, 2, 1))
+
+        # Extract outputs
+        detections = self._extract_onnx_output(
+            preds=preds,
+            image_shape=(height, width),
+            input_shape=(640, 640)
+        )
+
+        return detections
+
+    def _extract_onnx_output(self, preds: np.ndarray, image_shape: Tuple[int, int],
+                            input_shape: Tuple[int, int]) -> Dict:
+        """Extract detection results from ONNX model output"""
+        class_ids, confs, boxes = [], [], []
+
+        image_height, image_width = image_shape
+        input_height, input_width = input_shape
+        x_factor = image_width / input_width
+        y_factor = image_height / input_height
+
+        rows = preds[0].shape[0]
+        for i in range(rows):
+            row = preds[0][i]
+            conf = row[4]
+
+            classes_score = row[4:]
+            _, _, _, max_idx = cv2.minMaxLoc(classes_score)
+            class_id = max_idx[1]
+
+            if classes_score[class_id] > self.confidence:
+                confs.append(float(conf))
+                label = self.classes[int(class_id)] if int(class_id) < len(self.classes) else f"class_{class_id}"
+                class_ids.append(label)
+
+                # Extract boxes
+                x, y, w, h = row[0].item(), row[1].item(), row[2].item(), row[3].item()
+                left = int((x - 0.5 * w) * x_factor)
+                top = int((y - 0.5 * h) * y_factor)
+                width = int(w * x_factor)
+                height = int(h * y_factor)
+                box = [left, top, left + width, top + height]
+                boxes.append(box)
+
+        # Apply NMS with deterministic ordering
+        if len(boxes) > 0:
+            # Convert to proper format for NMS
+            nms_boxes = [[b[0], b[1], b[2]-b[0], b[3]-b[1]] for b in boxes]
+            
+            indices = cv2.dnn.NMSBoxes(
+                nms_boxes,
+                confs, 
+                self.confidence, 
+                self.iou_threshold
+            )
+
+            if len(indices) > 0:
+                indices = indices.flatten()
+                
+                # Create detection results
+                final_boxes = [boxes[i] for i in indices]
+                final_confs = [confs[i] for i in indices]
+                final_classes = [class_ids[i] for i in indices]
+                
+                # Sort by confidence for consistency
+                sorted_data = sorted(zip(final_boxes, final_confs, final_classes, range(len(indices))), 
+                                   key=lambda x: x[1], reverse=True)
+                
+                return {
+                    'boxes': [item[0] for item in sorted_data],
+                    'confidences': [item[1] for item in sorted_data],
+                    'classes': [item[2] for item in sorted_data],
+                    'class_ids': [item[3] for item in sorted_data]
+                }
+
+        return {'boxes': [], 'confidences': [], 'classes': [], 'class_ids': []}
+
+    def detect_batch(self, images: List[np.ndarray]) -> List[Dict]:
+        """Detect on multiple images with consistent ordering"""
+        return [self.detect(img) for img in images]
+    
+    def reset_deterministic_state(self):
+        """Reset deterministic state - call this between different sessions"""
+        if self.deterministic:
+            self._set_deterministic(42)
+            print("Deterministic state reset")