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

Code before:
    if len(unique(img)) == 1:
        return img

    if not has_morphology:
        raise RuntimeError("Skimage required.")

    if by_measure:
        sr = len(img.shape[1:])
        if isinstance(img, monai.data.MetaTensor):
            _pixdim = img.pixdim
        elif pixdim is not None:
            _pixdim = ensure_tuple_rep(pixdim, sr)
        else:
            warnings.warn("`img` is not of type MetaTensor and `pixdim` is None, assuming affine to be identity.")
            _pixdim = (1.0,) * sr
        voxel_volume = np.prod(np.array(_pixdim))
        if voxel_volume == 0:
            warnings.warn("Invalid `pixdim` value detected, set it to 1. Please verify the pixdim settings.")
            voxel_volume = 1
        min_size = np.ceil(min_size / voxel_volume)
    elif pixdim is not None:
        warnings.warn("`pixdim` is specified but not in use when computing the volume.")

    img_np: np.ndarray
    img_np, *_ = convert_data_type(img, np.ndarray)

    # morphology.remove_small_objects assumes them to be independent by default
    # else, convert to foreground vs background, remove small objects, then convert
    # back by multiplying the output by the input
    if not independent_channels:
        img_np = img_np > 0
    else:
        # if binary, convert to boolean, else int
        img_np = img_np.astype(bool if img_np.max() <= 1 else np.int32)

    out_np = morphology.remove_small_objects(img_np, min_size, connectivity)
    out, *_ = convert_to_dst_type(out_np, img)

    # convert back by multiplying
    if not independent_channels:
        out = img * out  # type: ignore
    return out


def get_unique_labels(img: NdarrayOrTensor, is_onehot: bool, discard: int | Iterable[int] | None = None) -> set[int]:
    """Get list of non-background labels in an image.

    Args:
        img: Image to be processed. Shape should be [C, W, H, [D]] with C=1 if not onehot else `num_classes`.
        is_onehot: Boolean as to whether input image is one-hotted. If one-hotted, only return channels with
        discard: Can be used to remove labels (e.g., background). Can be any value, sequence of values, or

Review comment (suggestion):
```suggestion
    _array = [torch.arange(end=image_size[i], dtype=torch.float32, device=point.device) for i in range(3)]
```
You can omit these arguments since start=0 and step=1 is the default. 

Code after:
                    outs[[bs]] += lib.isin(features[[bs]], index)
                    break
    outs[outs > 1] = 1
    return convert_to_dst_type(outs, dst=img_pos, dtype=outs.dtype)[0]


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 3D tuple.
        point: [B, N, 3], 3D point coordinates.
        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 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, n in np.ndindex(*point.shape[:2]):
        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 sample_points_from_label(
    labels: Tensor,
    label_set: Sequence[int],
    max_ppoint: int = 1,
    max_npoint: int = 0,
    device: torch.device | str | None = "cpu",
    use_center: bool = False,
):
    """Sample points from labels.