Configurable iou threshold (#24)

## Problem
In certain projects, non-max suppression (NMS) is causing detection issues. By default, points are merged during inference at an IoU of 0.45. However, relevant objects may be only some px apart, preventing the model from detecting all objects.

Simply reducing point size isn't viable as it would exclude small references during training.

## Changes
- Made NMS IoU threshold configurable in both Trainer and Detector
- Fixed bug in TensorRT-based YOLOv5 inference that prevented detection of boxes directly at the top or left image edges
  - Note: This edge detection bug did not affect point detection

Author: denniswittich

README.md

@@ -35,6 +35,17 @@ There are two variants of the detector:
 - to be deployed on a regular Linux computer, e.g. running Ubuntu (referred to as cloud-detectors)
 - to be deployed on a Jetson Nano running Linux4Tegra (L4T)
 
+Mandatory parameters are those described in [Zauberzeug Learning Loop Node Library](https://github.com/zauberzeug/learning_loop_node).
+Besides, the following parameters may be set:
+
+| Name           | Purpose                                   | Value                     | Default | Required only with ./docker.sh |
+| -------------- | ----------------------------------------- | ------------------------- | ------- | ------------------------------ |
+| LINKLL         | Link the node library into the container? | TRUE or FALSE             | FALSE   | Yes                            |
+| DETECTOR_NAME  | Will be the name of the container         | String                    | -       | Yes                            |
+| WEIGHT_TYPE    | Data type to convert weights to           | String [FP32, FP16, INT8] | FP16    | No                             |
+| IOU_THRESHOLD  | IoU threshold for NMS                     | Float                     | 0.45    | No                             |
+| CONF_THRESHOLD | Confidence threshold for NMS              | Float                     | 0.2     | No                             |
+
 ### Cloud-Detector
 
 New images can be pulled with `docker pull zauberzeug/yolov5-detector:nlvX.Y.Z-cloud`, where `X.Y.Z` is the version of the node-lib used.
@@ -43,13 +54,6 @@ Legacy images can be pulled with `docker pull zauberzeug/yolov5-detector:cloud`.
 Pulled images can be run with the `docker.sh` script by calling `./docker.sh run-image`.
 Local builds can be run with `./docker.sh run`.
 If the container does not use the GPU, try `./docker.sh d`.
-Mandatory parameters are those described in [Zauberzeug Learning Loop Node Library](https://github.com/zauberzeug/learning_loop_node). Besides, the following parameters may be set:
-
-| Name          | Purpose                                   | Value                     | Default | Required only with ./docker.sh |
-| ------------- | ----------------------------------------- | ------------------------- | ------- | ------------------------------ |
-| LINKLL        | Link the node library into the container? | TRUE or FALSE             | FALSE   | Yes                            |
-| DETECTOR_NAME | Will be the name of the container         | String                    | -       | Yes                            |
-| WEIGHT_TYPE   | Data type to convert weights to           | String [FP32, FP16, INT8] | FP16    | No                             |
 
 ### L4T-Detector
 

detector/docker.sh

@@ -28,7 +28,7 @@ fi
 
 # ========================== BUILD CONFIGURATION / IMAGE SELECTION =======================
 
-SEMANTIC_VERSION=0.1.11
+SEMANTIC_VERSION=0.1.12
 NODE_LIB_VERSION=0.14.0
 build_args=" --build-arg NODE_LIB_VERSION=$NODE_LIB_VERSION"
 

detector/yolov5.py

@@ -11,12 +11,12 @@
 
 import cv2
 import numpy as np
-import pycuda.driver as cuda
-import tensorrt as trt
-from pycuda._driver import Error as CudaError
+import pycuda.driver as cuda  # type: ignore # pylint: disable=import-error
+import tensorrt as trt  # type: ignore # pylint: disable=import-error
+from pycuda._driver import (  # type: ignore # pylint: disable=import-error
+    Error as CudaError,
+)
 
-CONF_THRESH = 0.2
-IOU_THRESHOLD = 0.4
 LEN_ALL_RESULT = 38001
 LEN_ONE_RESULT = 38
 
@@ -40,17 +40,24 @@ class YoLov5TRT():
     description: A YOLOv5 class that warps TensorRT ops, preprocess and postprocess ops.
     """
 
-    def __init__(self, engine_file_path: str):
+    def __init__(self, engine_file_path: str, iou_threshold: float, conf_threshold: float):
+        logging.info('Initializing YOLOv5 TRT engine with iou_threshold: %s, conf_threshold: %s',
+                     iou_threshold, conf_threshold)
         # Create a Context on this device,
         try:
             cuda.init()
-        except CudaError as e:
-            logging.exception('cuda init error:', e)
+        except CudaError:
+            logging.exception('cuda init error:')
             self.cuda_init_error = True
             return
 
         self.cuda_init_error = False
 
+        self.iou_threshold = iou_threshold
+        """a iou threshold to filter detections during nms"""
+        self.conf_threshold = conf_threshold
+        """a confidence threshold to filter detections during nms"""
+
         self.ctx = cuda.Device(0).make_context()
         stream = cuda.Stream()
         TRT_LOGGER = trt.Logger(trt.Logger.INFO)
@@ -107,7 +114,7 @@ def check_cuda_init_error(self):
     def infer(self, image_raw):
         self.check_cuda_init_error()
 
-        threading.Thread.__init__(self)
+        threading.Thread.__init__(self)  # type: ignore
         # Make self the active context, pushing it on top of the context stack.
         self.ctx.push()
         # Restore
@@ -259,8 +266,7 @@ def _post_process(self, output, origin_h, origin_w):
         pred = np.reshape(output[1:], (-1, LEN_ONE_RESULT))[:num, :]
         pred = pred[:, :6]
         # Do nms
-        boxes = self._non_max_suppression(
-            pred, origin_h, origin_w, conf_thres=CONF_THRESH, nms_thres=IOU_THRESHOLD)
+        boxes = self._non_max_suppression(pred, origin_h, origin_w)
         result_boxes = boxes[:, :4] if len(boxes) else np.array([])
         result_scores = boxes[:, 4] if len(boxes) else np.array([])
         result_classid = boxes[:, 5] if len(boxes) else np.array([])
@@ -309,19 +315,21 @@ def bbox_iou(self, box1, box2, x1y1x2y2=True):
 
         return iou
 
-    def _non_max_suppression(self, prediction, origin_h, origin_w, conf_thres=0.5, nms_thres=0.4):
+    def _non_max_suppression(self, prediction, origin_h, origin_w):
         """
         description: Removes detections with lower object confidence score than 'conf_thres' and performs
         Non-Maximum Suppression to further filter detections.
         param:
             prediction: detections, (x1, y1, x2, y2, conf, cls_id)
             origin_h: original image height
             origin_w: original image width
-            conf_thres: a confidence threshold to filter detections
-            nms_thres: a iou threshold to filter detections
         return:
             boxes: output after nms with the shape (x1, y1, x2, y2, conf, cls_id)
         """
+
+        conf_thres = self.conf_threshold
+        nms_thres = self.iou_threshold
+
         # Get the boxes that score > CONF_THRESH
         boxes = prediction[prediction[:, 4] >= conf_thres]
         # Trandform bbox from [center_x, center_y, w, h] to [x1, y1, x2, y2]

detector/yolov5_detector.py

@@ -23,6 +23,8 @@ def __init__(self) -> None:
         assert self.weight_type in ['FP16', 'FP32', 'INT8'], 'WEIGHT_TYPE must be one of FP16, FP32, INT8'
         self.log = logging.getLogger('Yolov5Detector')
         self.log.setLevel(logging.INFO)
+        self.iou_threshold = float(os.getenv('IOU_THRESHOLD', '0.45'))
+        self.conf_threshold = float(os.getenv('CONF_THRESHOLD', '0.2'))
 
     def init(self) -> None:
         assert self.model_info is not None, 'model_info must be set before calling init()'
@@ -37,30 +39,40 @@ def init(self) -> None:
             self.yolov5 = None
             self.log.info('destroyed old yolov5 instance')
 
-        self.yolov5 = yolov5.YoLov5TRT(engine_file)
+        self.yolov5 = yolov5.YoLov5TRT(engine_file, self.iou_threshold, self.conf_threshold)
         for _ in range(3):
             warmup = yolov5.warmUpThread(self.yolov5)
             warmup.start()
             warmup.join()
 
     @staticmethod
     def clip_box(
-            x: float, y: float, width: float, height: float, img_width: int, img_height: int) -> Tuple[
-            float, float, float, float]:
-        '''make sure the box is within the image
-            x,y is the center of the box
+            x1: float, y1: float, width: float, height: float, img_width: int, img_height: int) -> Tuple[
+            int, int, int, int]:
+        '''Clips a box defined by top-left corner (x1, y1), width, and height
+           to stay within image boundaries (img_width, img_height).
+           Returns the clipped (x1, y1, width, height) as ints.
         '''
-        left = max(0, x - 0.5 * width)
-        top = max(0, y - 0.5 * height)
-        right = min(img_width, x + 0.5 * width)
-        bottom = min(img_height, y + 0.5 * height)
+        x2 = x1 + width
+        y2 = y1 + height
 
-        x = 0.5 * (left + right)
-        y = 0.5 * (top + bottom)
-        width = right - left
-        height = bottom - top
+        # Clip coordinates
+        clipped_x1 = round(max(0.0, x1))
+        clipped_y1 = round(max(0.0, y1))
+        clipped_x2 = round(min(float(img_width), x2))
+        clipped_y2 = round(min(float(img_height), y2))
 
-        return x, y, width, height
+        # Recalculate dimensions
+        clipped_width = clipped_x2 - clipped_x1
+        clipped_height = clipped_y2 - clipped_y1
+
+        # Ensure width and height are non-negative
+        if clipped_width < 0:
+            clipped_width = 0
+        if clipped_height < 0:
+            clipped_height = 0
+
+        return clipped_x1, clipped_y1, clipped_width, clipped_height
 
     @staticmethod
     def clip_point(x: float, y: float, img_width: int, img_height: int) -> Tuple[float, float]:
@@ -87,13 +99,13 @@ def evaluate(self, image: bytes) -> ImageMetadata:
                     skipped_detections.append((category.name, detection))
                     continue
                 if category.type == CategoryType.Box:
-                    x, y, w, h = self.clip_box(x, y, w, h, im_width, im_height)
+                    clipped_x1, clipped_y1, clipped_w, clipped_h = self.clip_box(x, y, w, h, im_width, im_height)
                     image_metadata.box_detections.append(
                         BoxDetection(category_name=category.name,
-                                     x=round(x),
-                                     y=round(y),
-                                     width=round(x+w)-round(x),
-                                     height=round(y+h)-round(y),
+                                     x=clipped_x1,
+                                     y=clipped_y1,
+                                     width=clipped_w,
+                                     height=clipped_h,
                                      category_id=category.id,
                                      model_name=self.model_info.version,
                                      confidence=probability))