scipy.ndimage.find_objects GPU implementation

🚀 The feature

This function a basic building block of any biomedical or even manufacturing image analysis application.
Given an image containing multiple objects, say ~1000-10,000 objects, it provides the bounding box of each object.

It gives a list of tuple of slices of coordinates of labelled objects/cells within a mask image of dtype Uint16 or Uint8, assuming the image background is 0, and the labelled objects go from 1, 2, ..., max_label.

I was wondering it is possible to implement it in torch C++ using a simple TensorIterator.

Basically the simplest case would be it takes a 2D tensor of size (H, W) as input and
outputs a tensor of slices of size (N, 2) where N is the number of objects, and each row
is [slice(start,end,step), slice(start,end,step)].

which uses Iterators defined here:
https://github.com/scipy/scipy/blob/v1.10.1/scipy/ndimage/src/ni_support.h

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki

Motivation, pitch

This function is the key that allows to embed traditional machine learning tools such as thresholding techniques (such as otsu, triangle, max_entropy, and etc), watershed segmentation, edge filters, dilation/erosion operations into pytorch. Traditional machine learning tools are very helpful and can still be used to achieve great image segmentation. They are backed by numpy and scipy libraries can analyze a single image at a time and can be parallelized across multiple cpus.

However, they are a lot slower compared to pytorch dataloader fetching data for GPU computing.
Platforms such as Nvidia RAPIDS are a failure to achieve this. They are not properly maintained and have lots of issues. They work with Numba and are very tricky to use and install.

Alternatives

It is much faster and much more efficient than torchvision.ops.masks_to_boxes.
The implementation in C numpy can be found here:
https://github.com/scipy/scipy/blob/v1.10.1/scipy/ndimage/src/ni_measure.c

Additional context

Can it also be extended to allow extract objects from a tensor of dimension (B, C, W, H) where B is the batch size, C the number of
channels and W is the width and H is the height.

Start by comparing SciPy's ni_support.h iterators and ni_measure.c implementation, then determine how the proposed torch C++ TensorIterator approach would fit the project. Done means a GPU implementation accepts labelled 2D tensors and returns bounding-box slices for the objects; batched and channel dimensions are an additional requested extension.