Update src/detection.py

src/detection.py

@@ -3,12 +3,9 @@ from typing import List, Dict, Tuple
 import cv2
 from pathlib import Path
 import yaml
-import gc
-import torch
-from datetime import datetime
 
 class YOLOv11Detector:
-    """YOLOv11 detector for car damage detection with memory management"""
+    """YOLOv11 detector for car damage detection"""
 
     def __init__(self, config_path: str = "config.yaml"):
         """Initialize YOLOv11 detector with configuration"""
@@ -38,54 +35,12 @@ class YOLOv11Detector:
         self.confidence = self.config['model']['confidence']
         self.iou_threshold = self.config['model']['iou_threshold']
         self.classes = self.config['detection']['classes']
-        
-        # Thêm tracking cho model reload
-        self.model = None
-        self.net = None
-        self.inference_count = 0
-        self.last_reload_time = datetime.now()
-        self.max_inferences_before_reload = 100  # Reload sau 100 lần inference
-        self.max_hours_before_reload = 2  # Reload sau 2 giờ
-
-        # Load model lần đầu
-        self._load_model()
-
-    def _load_model(self):
-        """Load hoặc reload model"""
-        try:
-            # Clear existing model
-            if hasattr(self, 'model') and self.model is not None:
-                del self.model
-            if hasattr(self, 'net') and self.net is not None:
-                del self.net
-            
-            # Force garbage collection
-            gc.collect()
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-            
-            # Load model based on format
-            if self.model_path.endswith('.onnx'):
-                self._load_onnx_model()
-            else:  # .pt format
-                self._load_pytorch_model()
-                
-            # Reset counters
-            self.inference_count = 0
-            self.last_reload_time = datetime.now()
-            
-            print(f"Model (re)loaded successfully at {self.last_reload_time}")
-            
-        except Exception as e:
-            print(f"Error loading model: {e}")
-            raise
-
-    def _should_reload_model(self) -> bool:
-        """Kiểm tra xem có cần reload model không"""
-        time_diff = (datetime.now() - self.last_reload_time).total_seconds() / 3600
-        
-        return (self.inference_count >= self.max_inferences_before_reload or 
-                time_diff >= self.max_hours_before_reload)
+
+        # 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"""
@@ -93,27 +48,17 @@ class YOLOv11Detector:
         self.model = YOLO(self.model_path)
 
         # Set model to appropriate device
-        if self.device == 'cuda:0' and torch.cuda.is_available():
+        if self.device == 'cuda:0':
             self.model.to('cuda')
-        else:
-            self.model.to('cpu')
-
-        # QUAN TRỌNG: Set model to evaluation mode
-        self.model.model.eval()  # Ultralytics YOLO có nested model
-        
-        # Freeze BatchNorm và Dropout
-        for module in self.model.model.modules():
-            if isinstance(module, (torch.nn.BatchNorm2d, torch.nn.Dropout)):
-                module.eval()
 
-        print(f"Loaded PyTorch model: {self.model_path} (eval mode)")
+        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' and cv2.cuda.getCudaEnabledDeviceCount() > 0:
+        if self.device == 'cuda:0':
             self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
             self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
         else:
@@ -124,7 +69,7 @@ class YOLOv11Detector:
 
     def detect(self, image: np.ndarray) -> Dict:
         """
-        Perform detection on image with automatic model reload
+        Perform detection on image
 
         Args:
             image: Input image as numpy array (BGR format)
@@ -132,84 +77,53 @@ class YOLOv11Detector:
         Returns:
             Dictionary containing detection results
         """
-        # Kiểm tra và reload model nếu cần
-        if self._should_reload_model():
-            print("Reloading model due to usage limit/time limit...")
-            self._load_model()
-        
-        # Increment counter
-        self.inference_count += 1
-        
-        try:
-            if self.model_path.endswith('.onnx'):
-                return self._detect_onnx(image)
-            else:
-                return self._detect_pytorch(image)
-        except Exception as e:
-            print(f"Detection failed, attempting model reload: {e}")
-            # Thử reload model và detect lại
-            self._load_model()
-            if self.model_path.endswith('.onnx'):
-                return self._detect_onnx(image)
-            else:
-                return self._detect_pytorch(image)
+        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"""
-        try:
-            # QUAN TRỌNG: Sử dụng no_grad để tắt gradient computation
-            with torch.no_grad():
-                # Đảm bảo model ở eval mode
-                self.model.model.eval()
-                
-                # 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': []
-                }
+        # Run YOLO inference
+        results = self.model(
+            image,
+            conf=self.confidence,
+            iou=self.iou_threshold,
+            device=self.device,
+            verbose=False
+        )
 
-                if len(results) > 0 and results[0].boxes is not None:
-                    boxes = results[0].boxes
+        # Parse results
+        detections = {
+            'boxes': [],
+            'confidences': [],
+            'classes': [],
+            'class_ids': []
+        }
 
-                    for box in boxes:
-                        # Get box coordinates (xyxy format)
-                        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
+        if len(results) > 0 and results[0].boxes is not None:
+            boxes = results[0].boxes
 
-                        # Get confidence and class
-                        conf = float(box.conf[0].cpu().numpy())
-                        cls_id = int(box.cls[0].cpu().numpy())
+            for box in boxes:
+                # Get box coordinates (xyxy format)
+                x1, y1, x2, y2 = box.xyxy[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}"
+                # Get confidence and class
+                conf = float(box.conf[0].cpu().numpy())
+                cls_id = int(box.cls[0].cpu().numpy())
 
-                        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)
+                # 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}"
 
-            # Clear GPU cache sau mỗi inference
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
+                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
-            
-        except Exception as e:
-            print(f"PyTorch detection error: {e}")
-            raise
+        return detections
 
     def _detect_onnx(self, image: np.ndarray) -> Dict:
         """Detection using ONNX model (compatible with original code)"""
@@ -287,24 +201,4 @@ class YOLOv11Detector:
 
     def detect_batch(self, images: List[np.ndarray]) -> List[Dict]:
         """Detect on multiple images"""
-        return [self.detect(img) for img in images]
-
-    def get_model_stats(self) -> Dict:
-        """Trả về thống kê về model"""
-        time_diff = (datetime.now() - self.last_reload_time).total_seconds() / 3600
-        
-        # Kiểm tra model mode
-        model_mode = "unknown"
-        if hasattr(self, 'model') and self.model is not None:
-            if hasattr(self.model, 'model'):
-                model_mode = "eval" if not self.model.model.training else "train"
-        
-        return {
-            "inference_count": self.inference_count,
-            "hours_since_reload": round(time_diff, 2),
-            "last_reload_time": self.last_reload_time.isoformat(),
-            "model_path": self.model_path,
-            "device": self.device,
-            "model_mode": model_mode,  # Thêm thông tin model mode
-            "torch_no_grad_enabled": not torch.is_grad_enabled()
-        }
+        return [self.detect(img) for img in images]