Add sigmoid smoothing to instance segmentation post-processing of YOLO models family in inference-models (#2217)

* Add sigmoid smoothing to instance segmentation post-processing of YOLO models family in inference-models

* Adjust tests cases

* Make linters happy

* Adjust tests cases

* Adjust tests cases

* Update docs

* Bump version of inference requirements

Author: PawelPeczek-Roboflow

inference_models/docs/changelog.md

@@ -1,5 +1,25 @@
 # Changelog
 
+## `0.24.3`
+
+### Changed
+
+- Added `sigmoid` smoothing for instance-segmentation masks in YOLOv8, YOLOv11, YOLOv12 models family.
+Smoothing can be enabled / disabled via `masks_smoothing_enabled` parameter of `post_process(...)` method
+(which can be passed as `**kwarg` to `forward(...)`) with default set with 
+`INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_SMOOTHING_ENABLED` (set to `True`). Additionally, the binarization 
+threshold for masks can be controlled via `masks_binarization_threshold` parameter - default to be 
+controlled with `INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_BINARIZATION_THRESHOLD` (set to `0.5` or `0.0` 
+depending on `INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_SMOOTHING_ENABLED`).
+
+!!! warning "Instance-segmentation masks will change"
+
+    Due to smoothing, there is slight change to segmentation masks expected - mainly regarding edges 
+    of predictions which should be smoother now. Change is dictated by alignment to old `inference` versions
+    behaviour, effectively drifting from `ultralytics` post-processing.
+
+---
+
 ## `0.24.2`
 
 ### Fixed

inference_models/inference_models/configuration.py

@@ -350,6 +350,18 @@
     variable_name="INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_CLASS_AGNOSTIC_NMS",
     default=INFERENCE_MODELS_DEFAULT_CLASS_AGNOSTIC_NMS,
 )
+INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_SMOOTHING_ENABLED = get_boolean_from_env(
+    variable_name="INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_SMOOTHING_ENABLED",
+    default=True,
+)
+INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_BINARIZATION_THRESHOLD = get_float_from_env(
+    variable_name="INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_BINARIZATION_THRESHOLD",
+    default=(
+        0.5
+        if INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_SMOOTHING_ENABLED
+        else 0.0
+    ),
+)
 INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_KEY_POINTS_THRESHOLD = get_float_from_env(
     variable_name="INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_KEY_POINTS_THRESHOLD",
     default=0.0,

inference_models/inference_models/models/auto_loaders/auto_resolution_cache.py

@@ -125,6 +125,8 @@ def invalidate(self, auto_negotiation_hash: str) -> None:
 
 
 def generate_auto_resolution_cache_path(auto_negotiation_hash: str) -> str:
-    return os.path.abspath(os.path.join(
-        INFERENCE_HOME, "auto-resolution-cache", f"{auto_negotiation_hash}.json"
-    ))
+    return os.path.abspath(
+        os.path.join(
+            INFERENCE_HOME, "auto-resolution-cache", f"{auto_negotiation_hash}.json"
+        )
+    )

inference_models/inference_models/models/auto_loaders/core.py

@@ -1510,7 +1510,9 @@ def generate_shared_blobs_path() -> str:
 def generate_model_package_cache_path(model_id: str, package_id: str) -> str:
     ensure_package_id_is_os_safe(model_id=model_id, package_id=package_id)
     model_id_slug = slugify_model_id_to_os_safe_format(model_id=model_id)
-    return os.path.abspath(os.path.join(INFERENCE_HOME, "models-cache", model_id_slug, package_id))
+    return os.path.abspath(
+        os.path.join(INFERENCE_HOME, "models-cache", model_id_slug, package_id)
+    )
 
 
 def ensure_package_id_is_os_safe(model_id: str, package_id: str) -> None:

inference_models/inference_models/models/common/roboflow/post_processing.py

@@ -314,7 +314,7 @@ def crop_masks_to_boxes(
     scaling: float = 0.25,
 ) -> torch.Tensor:
     n, h, w = masks.shape
-    scaled_boxes = boxes * scaling
+    scaled_boxes = torch.round(boxes * scaling)
     x1, y1, x2, y2 = (
         scaled_boxes[:, 0][:, None, None],
         scaled_boxes[:, 1][:, None, None],

inference_models/inference_models/models/yololite/yololite_object_detection_onnx.py

@@ -200,11 +200,18 @@ def post_process(
         **kwargs,
     ) -> List[Detections]:
         # Backward compatibility: earlier model packages have no post_processing config — always unfused 3-tensor output
-        if self._inference_config.post_processing and self._inference_config.post_processing.fused:
+        if (
+            self._inference_config.post_processing
+            and self._inference_config.post_processing.fused
+        ):
             nms_results = self._post_process_fused(model_results, confidence)
         else:
             nms_results = self._post_process_unfused(
-                model_results, confidence, iou_threshold, max_detections, class_agnostic_nms,
+                model_results,
+                confidence,
+                iou_threshold,
+                max_detections,
+                class_agnostic_nms,
             )
         rescaled_results = rescale_detections(
             detections=nms_results,
@@ -228,7 +235,9 @@ def _post_process_fused(
     ) -> List[torch.Tensor]:
         # Single output tensor [B, max_det, 6]: x1, y1, x2, y2, conf, class_id
         output = model_results[0]
-        return post_process_nms_fused_model_output(output=output, conf_thresh=confidence)
+        return post_process_nms_fused_model_output(
+            output=output, conf_thresh=confidence
+        )
 
     def _post_process_unfused(
         self,
@@ -240,7 +249,9 @@ def _post_process_unfused(
     ) -> List[torch.Tensor]:
         # Decoded outputs without fused NMS: boxes_xyxy [B,N,4], obj_logits [B,N,1], cls_logits [B,N,C]
         boxes_xyxy, obj_logits, cls_logits = (
-            model_results[0], model_results[1], model_results[2],
+            model_results[0],
+            model_results[1],
+            model_results[2],
         )
         obj_conf = torch.sigmoid(obj_logits)
         cls_conf = torch.sigmoid(cls_logits)

inference_models/inference_models/models/yolov8/yolov8_instance_segmentation_onnx.py

@@ -14,6 +14,8 @@
     INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_CLASS_AGNOSTIC_NMS,
     INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_CONFIDENCE,
     INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_IOU_THRESHOLD,
+    INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_BINARIZATION_THRESHOLD,
+    INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_SMOOTHING_ENABLED,
     INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MAX_DETECTIONS,
 )
 from inference_models.entities import ColorFormat
@@ -207,6 +209,8 @@ def post_process(
         iou_threshold: float = INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_IOU_THRESHOLD,
         max_detections: int = INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MAX_DETECTIONS,
         class_agnostic_nms: bool = INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_CLASS_AGNOSTIC_NMS,
+        masks_smoothing_enabled: bool = INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_SMOOTHING_ENABLED,
+        masks_binarization_threshold: float = INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_BINARIZATION_THRESHOLD,
         **kwargs,
     ) -> List[InstanceDetections]:
         instances, protos = model_results
@@ -230,6 +234,8 @@ def post_process(
                 protos=image_protos,
                 masks_in=image_bboxes[:, 6:],
             )
+            if masks_smoothing_enabled:
+                pre_processed_masks = torch.nn.functional.sigmoid(pre_processed_masks)
             cropped_masks = crop_masks_to_boxes(
                 image_bboxes[:, :4], pre_processed_masks
             )
@@ -249,6 +255,7 @@ def post_process(
                 size_after_pre_processing=image_meta.size_after_pre_processing,
                 inference_size=image_meta.inference_size,
                 static_crop_offset=image_meta.static_crop_offset,
+                binarization_threshold=masks_binarization_threshold,
             )
             final_results.append(
                 InstanceDetections(

inference_models/inference_models/models/yolov8/yolov8_instance_segmentation_torch_script.py

@@ -15,6 +15,8 @@
     INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_CLASS_AGNOSTIC_NMS,
     INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_CONFIDENCE,
     INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_IOU_THRESHOLD,
+    INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_BINARIZATION_THRESHOLD,
+    INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_SMOOTHING_ENABLED,
     INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MAX_DETECTIONS,
 )
 from inference_models.entities import ColorFormat
@@ -170,6 +172,8 @@ def post_process(
         iou_threshold: float = INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_IOU_THRESHOLD,
         max_detections: int = INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MAX_DETECTIONS,
         class_agnostic_nms: bool = INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_CLASS_AGNOSTIC_NMS,
+        masks_smoothing_enabled: bool = INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_SMOOTHING_ENABLED,
+        masks_binarization_threshold: float = INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_BINARIZATION_THRESHOLD,
         **kwargs,
     ) -> List[InstanceDetections]:
         instances, protos = model_results
@@ -193,6 +197,8 @@ def post_process(
                 protos=image_protos,
                 masks_in=image_bboxes[:, 6:],
             )
+            if masks_smoothing_enabled:
+                pre_processed_masks = torch.nn.functional.sigmoid(pre_processed_masks)
             cropped_masks = crop_masks_to_boxes(
                 image_bboxes[:, :4], pre_processed_masks
             )
@@ -212,6 +218,7 @@ def post_process(
                 size_after_pre_processing=image_meta.size_after_pre_processing,
                 inference_size=image_meta.inference_size,
                 static_crop_offset=image_meta.static_crop_offset,
+                binarization_threshold=masks_binarization_threshold,
             )
             final_results.append(
                 InstanceDetections(

inference_models/inference_models/models/yolov8/yolov8_instance_segmentation_trt.py

@@ -15,6 +15,8 @@
     INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_CLASS_AGNOSTIC_NMS,
     INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_CONFIDENCE,
     INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_IOU_THRESHOLD,
+    INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_BINARIZATION_THRESHOLD,
+    INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_SMOOTHING_ENABLED,
     INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MAX_DETECTIONS,
 )
 from inference_models.entities import ColorFormat
@@ -265,6 +267,8 @@ def post_process(
         iou_threshold: float = INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_IOU_THRESHOLD,
         max_detections: int = INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MAX_DETECTIONS,
         class_agnostic_nms: bool = INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_CLASS_AGNOSTIC_NMS,
+        masks_smoothing_enabled: bool = INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_SMOOTHING_ENABLED,
+        masks_binarization_threshold: float = INFERENCE_MODELS_YOLO_ULTRALYTICS_DEFAULT_MASKS_BINARIZATION_THRESHOLD,
         **kwargs,
     ) -> List[InstanceDetections]:
         with torch.cuda.stream(self._post_process_stream):
@@ -291,6 +295,10 @@ def post_process(
                     protos=image_protos,
                     masks_in=image_bboxes[:, 6:],
                 )
+                if masks_smoothing_enabled:
+                    pre_processed_masks = torch.nn.functional.sigmoid(
+                        pre_processed_masks
+                    )
                 cropped_masks = crop_masks_to_boxes(
                     image_bboxes[:, :4], pre_processed_masks
                 )
@@ -310,6 +318,7 @@ def post_process(
                     size_after_pre_processing=image_meta.size_after_pre_processing,
                     inference_size=image_meta.inference_size,
                     static_crop_offset=image_meta.static_crop_offset,
+                    binarization_threshold=masks_binarization_threshold,
                 )
                 final_results.append(
                     InstanceDetections(

inference_models/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "inference-models"
-version = "0.24.2"
+version = "0.24.3"
 description = "The new inference engine for Computer Vision models"
 readme = "README.md"
 requires-python = ">=3.10,<3.13"