Repository: Project-MONAI/MONAI
File: monai/transforms/utils.py (Python)

Code before:
def convert_points_to_disc(
    image_size: Sequence[int], point: Tensor, point_label: Tensor, radius: int = 2, disc: bool = False
):
    """
    Convert a 3D point coordinates into image mask. The returned mask has the same spatial
    size as `image_size` while the batch dimension is the same as 'point' batch dimension.
    The point is converted to a mask ball with radius defined by `radius`. The output
    contains two channels each for negative (first channel) and positive points.

    Args:
        image_size: The output size of the converted mask. It should be a
        point: [b, N, 3]
        point_label: [b, N], 0 or 2 means negative points, 1 or 3 means postive points.
        radius: disc ball radius size
        disc: If true, use regular disc other other use gaussian.
    """
    masks = torch.zeros([point.shape[0], 2, image_size[0], image_size[1], image_size[2]], device=point.device)
    _array = [
        torch.arange(start=0, end=image_size[i], step=1, dtype=torch.float32, device=point.device) for i in range(3)
    ]
    coord_rows, coord_cols, coord_z = torch.meshgrid(_array[2], _array[1], _array[0])
    # [1, 3, h, w, d] -> [b, 2, 3, h, w, d]
    coords = unsqueeze_left(torch.stack((coord_rows, coord_cols, coord_z), dim=0), 6)
    coords = coords.repeat(point.shape[0], 2, 1, 1, 1, 1)
    for b in range(point.shape[0]):
        for n in range(point.shape[1]):
            point_bn = unsqueeze_right(point[b, n], 6)
            if point_label[b, n] > -1:
                channel = 0 if (point_label[b, n] == 0 or point_label[b, n] == 2) else 1
                pow_diff = torch.pow(coords[b, channel] - point_bn[b, n], 2)
                if disc:
                    masks[b, channel] += pow_diff.sum(0) < radius**2
                else:
                    masks[b, channel] += torch.exp(-pow_diff.sum(0) / (2 * radius**2))
    return masks


def remove_small_objects(
    img: NdarrayTensor,
    min_size: int = 64,
    connectivity: int = 1,
    independent_channels: bool = True,
    by_measure: bool = False,
    pixdim: Sequence[float] | float | np.ndarray | None = None,
) -> NdarrayTensor:
    """
    Use `skimage.morphology.remove_small_objects` to remove small objects from images.
    See: https://scikit-image.org/docs/dev/api/skimage.morphology.html#remove-small-objects.

    Data should be one-hotted.


Review comment (suggestion):
```suggestion
    for b, n in np.ndindex(*point.shape[:2]):
```

Code after:


def get_largest_connected_component_mask_point(
    img_pos: NdarrayTensor,
    img_neg: NdarrayTensor,
    point_coords: NdarrayTensor,
    point_labels: NdarrayTensor,
    pos_val: Sequence[int] = (1, 3),
    neg_val: Sequence[int] = (0, 2),
    margins: int = 3,
) -> NdarrayTensor:
    """
    Gets the connected component of img_pos and img_neg that include the positive points and
    negative points separately. The function is used for combining automatic results with interactive
    results in VISTA3D.

    Args:
        img_pos: bool type tensor, shape [B, 1, H, W, D], where B means the foreground masks from a single 3D image.
        img_neg: same format as img_pos but corresponds to negative points.
        pos_val: positive point label values.
        neg_val: negative point label values.
        point_coords: the coordinates of each point, shape [B, N, 3], where N means the number of points.
        point_labels: the label of each point, shape [B, N].
    """

    cucim_skimage, has_cucim = optional_import("cucim.skimage")

    use_cp = has_cp and has_cucim and isinstance(img_pos, torch.Tensor) and img_pos.device != torch.device("cpu")
    if use_cp:
        img_pos_ = convert_to_cupy(img_pos.short())  # type: ignore
        img_neg_ = convert_to_cupy(img_neg.short())  # type: ignore
        label = cucim_skimage.measure.label
        lib = cp
    else:
        if not has_measure:
            raise RuntimeError("skimage.measure required.")
        img_pos_, *_ = convert_data_type(img_pos, np.ndarray)
        img_neg_, *_ = convert_data_type(img_neg, np.ndarray)
        # for skimage.measure.label, the input must be bool type
        if img_pos_.dtype != bool or img_neg_.dtype != bool:
            raise ValueError("img_pos and img_neg must be bool type.")
        label = measure.label
        lib = np

    features_pos, _ = label(img_pos_, connectivity=3, return_num=True)
    features_neg, _ = label(img_neg_, connectivity=3, return_num=True)

    outs = np.zeros_like(img_pos_)
    for bs in range(point_coords.shape[0]):
        for i, p in enumerate(point_coords[bs]):
            if point_labels[bs, i] in pos_val: