Upload TensorRTBBoxDetector.py

TensorRTBBoxDetector.py

@@ -173,12 +173,13 @@ class TRTYOLOBBoxDetector:
                 continue
 
             # ------------------------------------------------------------------
-            # Create full-image mask from bbox (uint8 0/255)
+            # Create full-image mask from bbox as float32 in [0, 1]
             # ------------------------------------------------------------------
-            mask = np.zeros((h, w), dtype=np.uint8)
-            mask[y1_i:y2_i, x1_i:x2_i] = 255
+            mask = np.zeros((h, w), dtype=np.float32)
+            mask[y1_i:y2_i, x1_i:x2_i] = 1.0
 
-            # Optional dilation / erosion via GPU-aware helper
+            # Optional dilation / erosion via GPU-aware helper.
+            # IMPORTANT: dilate_mask must return float32 [0,1] as well.
             if dilation:
                 mask = dilate_mask(mask, dilation)
 
@@ -196,7 +197,7 @@ class TRTYOLOBBoxDetector:
             # Crop image + mask
             # ------------------------------------------------------------------
             cropped_image = crop_image(image, crop_region)   # torch [1, h', w', C]
-            cropped_mask = crop_ndarray2(mask, crop_region)  # np.uint8 [h', w']
+            cropped_mask = crop_ndarray2(mask, crop_region)  # np.float32 [h', w'] in [0,1]
 
             # Build SEG object for this detection
             seg = SEG(
@@ -241,6 +242,8 @@ class TRTYOLOBBoxDetector:
         return core.segs_to_combined_mask((shape, seg_list))
 
 
+
+
 # -------------------------------------------------------------------------
 # NODE 1: TRTYOLOEngineBuilder
 # - Builds a TensorRT engine from a .pt file in the node folder.