Skip to content

[InferenceSlicer] - add segmentation models support #678

@SkalskiP

Description

@SkalskiP

Description

Currently, sv.InferenceSlicer supports only object detection models. Adding support for instance segmentation would require the following changes:

  • The sv.InferenceSlicer uses Non-Max Suppression (NMS) to sift out duplicate detections at the tile intersection. At the moment, Supervision only has a box-based NMS. A segmentation-based NMS would be almost ideal, the only change would be to replace the box_iou_batch with a new mask_iou_batch.
  • A segmentation-based NMS must be plugged into sv.InferenceSlicer. At this point, we would have to check whether the detections have masks. And if so, use the new NMS.

API

# create
def mask_iou_batch(boxes_true: np.ndarray, boxes_detection: np.ndarray) -> np.ndarray:
    pass

# rename non_max_suppression -> box_non_max_suppression

# create
def mask_non_max_suppression(predictions: np.ndarray, iou_threshold: float = 0.5) -> np.ndarray:
    pass

# change InferenceSlicer

Usage example

import cv2
import supervision as sv
from ultralytics import YOLO

image = cv2.image = cv2.imread(<SOURCE_IMAGEPATH>)
model = YOLO("yolov8x-seg.pt")

def callback(image_slice: np.ndarray) -> sv.Detections:
    result = model(image_slice)[0]
    return sv.Detections.from_ultralytics(result)

slicer = sv.InferenceSlicer(
    callback=callback,
    slice_wh=(512, 512),
    iou_threshold=0.5,
)

detections = slicer(image)

Additional

  • Note: Please share a Google Colab with minimal code to test the new feature. We know it's additional work, but it will definitely speed up the review process. Each change must be tested by the reviewer. Setting up a local environment to do this is time-consuming. Please ensure that Google Colab can be accessed without any issues (make it public). Thank you! 🙏🏻

Metadata

Metadata

Assignees

Labels

Q2.2024Tasks planned for execution in Q2 2024.enhancementNew feature or request

Type

No type

Projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions