add more logic to process image

config.yaml

@@ -1,6 +1,6 @@
 model:
   path: 'VietCat/GTSRB-Model/models/GTSRB.pt' # Path to the YOLO model on Hugging Face Hub (will be downloaded automatically)
-  confidence_threshold: 0.25 # Minimum confidence for detections
+  confidence_threshold: 0.15 # Minimum confidence for detections
 
 inference:
   box_color: (128, 0, 128) # Purple color for bounding boxes (BGR format)

model.py

@@ -39,6 +39,50 @@ class TrafficSignDetector:
         self.thickness = config['inference']['thickness']
         self.classes = config['classes']
     
+    def _ensure_square(self, image, target_size=640):
+        """
+        Adjust image to square while maintaining aspect ratio.
+        - If image is smaller: pad to target_size x target_size
+        - If image is larger: resize down to target_size x target_size
+        Letterbox padding is added to preserve aspect ratio.
+        :param image: input image (numpy array)
+        :param target_size: target size (default 640x640)
+        :return: square image (target_size x target_size)
+        """
+        height, width = image.shape[:2]
+        max_dim = max(width, height)
+        
+        # Scale to fit target while maintaining aspect ratio
+        scale = target_size / max_dim
+        
+        # Calculate new dimensions
+        new_width = int(width * scale)
+        new_height = int(height * scale)
+        
+        # Resize image
+        resized = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_LINEAR)
+        
+        # Create canvas and place resized image
+        canvas = np.full((target_size, target_size, 3), (114, 114, 114), dtype=np.uint8)
+        pad_x = (target_size - new_width) // 2
+        pad_y = (target_size - new_height) // 2
+        canvas[pad_y:pad_y + new_height, pad_x:pad_x + new_width] = resized
+        
+        print(f"Original: {image.shape} → Scale: {scale:.3f} → Resized: {resized.shape} → Final: {canvas.shape}")
+        
+        return canvas, scale, pad_x, pad_y
+    
+    def _preprocess(self, image):
+        """
+        Preprocess image: normalize pixel values to [0, 1] range.
+        :param image: input image (numpy array, uint8)
+        :return: normalized image (float32)
+        """
+        # Normalize pixel values from [0, 255] to [0, 1]
+        image = image.astype(np.float32) / 255.0
+        print(f"Image normalized - Min: {image.min():.3f}, Max: {image.max():.3f}, Mean: {image.mean():.3f}")
+        return image
+
     def detect(self, image):
         """
         Perform inference on the image and draw bounding boxes.
@@ -46,24 +90,54 @@ class TrafficSignDetector:
         :return: image with drawn bounding boxes
         """
         print(f"Input image shape: {image.shape}")
-        results = self.model(image, conf=self.conf_threshold)
+        
+        # Store original image for drawing (uint8)
+        original_image = image.copy()
+        
+        # Apply letterbox preprocessing to ensure 640x640 matching training size
+        # Returns both processed image and transformation info
+        image, scale, pad_x, pad_y = self._ensure_square(image, target_size=640)
+        
+        # Normalize pixel values for inference
+        image = self._preprocess(image)
+        
+        # Use imgsz=640 to match training size
+        results = self.model(image, conf=self.conf_threshold, imgsz=640)
         print(f"Number of results: {len(results)}")
 
+        # Get original dimensions for coordinate transformation
+        orig_h, orig_w = original_image.shape[:2]
+        
         for result in results:
             boxes = result.boxes
             print(f"Number of boxes in this result: {len(boxes)}")
+            
+            # Debug: print all detection confidences
+            if len(boxes) > 0:
+                confidences = [float(box.conf[0]) for box in boxes]
+                print(f"Detected confidences: {confidences}")
+            else:
+                print(f"No detections above threshold {self.conf_threshold}")
+            
             for box in boxes:
-                # Get bounding box coordinates
+                # Get bounding box coordinates from letterboxed image
                 x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
+                
+                # Convert coordinates back to original image space
+                x1 = max(0, int((x1 - pad_x) / scale))
+                y1 = max(0, int((y1 - pad_y) / scale))
+                x2 = min(orig_w, int((x2 - pad_x) / scale))
+                y2 = min(orig_h, int((y2 - pad_y) / scale))
+                
                 conf = box.conf[0].cpu().numpy()
                 cls = int(box.cls[0].cpu().numpy())
                 print(f"Detected: {self.classes[cls]} with conf {conf:.2f} at ({x1},{y1})-({x2},{y2})")
 
-                # Draw bounding box
-                cv2.rectangle(image, (x1, y1), (x2, y2), self.box_color, self.thickness)
+                # Draw bounding box on original image
+                cv2.rectangle(original_image, (x1, y1), (x2, y2), self.box_color, self.thickness)
 
                 # Draw label
                 label = f"{self.classes[cls]}: {conf:.2f}"
-                cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, self.text_color, 2)
+                cv2.putText(original_image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, self.text_color, 2)
 
-        return image
+        return original_image