How to use the stardist.nms.non_maximum_suppression function in stardist

def test_nms(images, n_rays):
    for img in images:
        prob = edt_prob(img)
        dist = _cpp_star_dist(img, n_rays)
        coord = dist_to_coord(dist)
        nms_a = non_maximum_suppression(coord, prob, prob_thresh=0.4, verbose=False, max_bbox_search=False)
        nms_b = non_maximum_suppression(coord, prob, prob_thresh=0.4, verbose=False, max_bbox_search=True)
        check_similar(nms_a,nms_b)

from utils import prob_dist_image2d
np.random.seed(42)
from scipy.ndimage import zoom
from stardist.geometry.geom2d import dist_to_coord, polygons_to_label
from stardist.nms import non_maximum_suppression
from skimage.morphology import watershed

if __name__ == '__main__':    


    img, prob, dist = prob_dist_image2d()

    grid = (2,2)
    
    coord = dist_to_coord(dist, grid = grid)
    points = non_maximum_suppression(coord, prob, grid = grid, prob_thresh=.4, nms_thresh=.3)

    labels0 = polygons_to_label(coord, prob, points, shape=img.shape)


    
    aff, aff_neg = dist_to_affinity2D(dist,
                                      weights = prob>0.03,
                                      grid = grid,
                                      normed=True, verbose = True);

    factor = tuple(s1/s2 for s1, s2 in zip(img.shape, prob.shape))

    potential = (np.mean(aff,-1))*prob
    potential = zoom(potential, factor, order=1)

    markers = np.zeros(img.shape, np.int32)

def test_nms(images, n_rays):
    for img in images:
        prob = edt_prob(img)
        dist = _cpp_star_dist(img, n_rays)
        coord = dist_to_coord(dist)
        nms_a = non_maximum_suppression(coord, prob, prob_thresh=0.4, verbose=False, max_bbox_search=False)
        nms_b = non_maximum_suppression(coord, prob, prob_thresh=0.4, verbose=False, max_bbox_search=True)
        check_similar(nms_a,nms_b)

def _instances_from_prediction(self, img_shape, prob, dist, prob_thresh=None, nms_thresh=None, overlap_label = None, **nms_kwargs):
        if prob_thresh is None: prob_thresh = self.thresholds.prob
        if nms_thresh  is None: nms_thresh  = self.thresholds.nms
        if overlap_label is not None: raise NotImplementedError("overlap_label not supported for 2D yet!")

        coord = dist_to_coord(dist, grid=self.config.grid)
        points = non_maximum_suppression(coord, prob, grid=self.config.grid,
                                         prob_thresh=prob_thresh, nms_thresh=nms_thresh, **nms_kwargs)
        labels = polygons_to_label(coord, prob, points, shape=img_shape)
        return labels, dict(coord=coord[points[:,0],points[:,1]], points=points, prob=prob[points[:,0],points[:,1]])