Add utils for vista3d

monai/transforms/utils.py

@@ -22,6 +22,7 @@
 
 import numpy as np
 import torch
+from torch import Tensor
 
 import monai
 from monai.config import DtypeLike, IndexSelection
@@ -65,6 +66,8 @@
     min_version,
     optional_import,
     pytorch_after,
+    unsqueeze_left,
+    unsqueeze_right,
 )
 from monai.utils.enums import TransformBackends
 from monai.utils.type_conversion import (
@@ -103,6 +106,7 @@
     "generate_spatial_bounding_box",
     "get_extreme_points",
     "get_largest_connected_component_mask",
+    "convert_points_to_disc",
     "remove_small_objects",
     "img_bounds",
     "in_bounds",
@@ -1172,6 +1176,43 @@ def get_largest_connected_component_mask(
     return convert_to_dst_type(out, dst=img, dtype=out.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
+        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,