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

Code before:
                l, r = max(x - margin, 0), min(x + margin + 1, features.shape[-3])
                t, d = max(y - margin, 0), min(y + margin + 1, features.shape[-2])
                f, b = max(z - margin, 0), min(z + margin + 1, features.shape[-1])
                if (features[bs, 0, l:r, t:d, f:b] > 0).any():
                    index = features[bs, 0, l:r, t:d, f:b].max()
                    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 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,

Review comment (suggestion):
```suggestion
        disc: If true, use regular disc, other use gaussian.
```

Code after:
                l, r = max(x - margin, 0), min(x + margin + 1, features.shape[-3])
                t, d = max(y - margin, 0), min(y + margin + 1, features.shape[-2])
                f, b = max(z - margin, 0), min(z + margin + 1, features.shape[-1])
                if (features[bs, 0, l:r, t:d, f:b] > 0).any():
                    index = features[bs, 0, l:r, t:d, f:b].max()
                    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,