monai-weekly 1.4.dev2435__py3-none-any.whl → 1.4.dev2436__py3-none-any.whl

monai/networks/utils.py

@@ -36,6 +36,8 @@ from monai.utils.type_conversion import convert_to_dst_type, convert_to_tensor
 onnx, _ = optional_import("onnx")
 onnxreference, _ = optional_import("onnx.reference")
 onnxruntime, _ = optional_import("onnxruntime")
+polygraphy, polygraphy_imported = optional_import("polygraphy")
+torch_tensorrt, _ = optional_import("torch_tensorrt", "1.4.0")
 
 __all__ = [
     "one_hot",
@@ -61,6 +63,7 @@ __all__ = [
     "look_up_named_module",
     "set_named_module",
     "has_nvfuser_instance_norm",
+    "get_profile_shapes",
 ]
 
 logger = get_logger(module_name=__name__)
@@ -68,6 +71,26 @@ logger = get_logger(module_name=__name__)
 _has_nvfuser = None
 
 
+def get_profile_shapes(input_shape: Sequence[int], dynamic_batchsize: Sequence[int] | None):
+    """
+    Given a sample input shape, calculate min/opt/max shapes according to dynamic_batchsize.
+    """
+
+    def scale_batch_size(input_shape: Sequence[int], scale_num: int):
+        scale_shape = [*input_shape]
+        scale_shape[0] = scale_num
+        return scale_shape
+
+    # Use the dynamic batchsize range to generate the min, opt and max model input shape
+    if dynamic_batchsize:
+        min_input_shape = scale_batch_size(input_shape, dynamic_batchsize[0])
+        opt_input_shape = scale_batch_size(input_shape, dynamic_batchsize[1])
+        max_input_shape = scale_batch_size(input_shape, dynamic_batchsize[2])
+    else:
+        min_input_shape = opt_input_shape = max_input_shape = input_shape
+    return min_input_shape, opt_input_shape, max_input_shape
+
+
 def has_nvfuser_instance_norm():
     """whether the current environment has InstanceNorm3dNVFuser
     https://github.com/NVIDIA/apex/blob/23.05-devel/apex/normalization/instance_norm.py#L15-L16
@@ -606,6 +629,9 @@ def convert_to_onnx(
     rtol: float = 1e-4,
     atol: float = 0.0,
     use_trace: bool = True,
+    do_constant_folding: bool = True,
+    constant_size_threshold: int = 16 * 1024 * 1024 * 1024,
+    dynamo=False,
     **kwargs,
 ):
     """
@@ -632,7 +658,10 @@ def convert_to_onnx(
         rtol: the relative tolerance when comparing the outputs of PyTorch model and TorchScript model.
         atol: the absolute tolerance when comparing the outputs of PyTorch model and TorchScript model.
         use_trace: whether to use `torch.jit.trace` to export the torchscript model.
-        kwargs: other arguments except `obj` for `torch.jit.script()` to convert model, for more details:
+        do_constant_folding: passed to onnx.export(). If True, extra polygraphy folding pass is done.
+        constant_size_threshold: passed to polygrapy conatant forling, default = 16M
+        kwargs: if use_trace=True: additional arguments to pass to torch.onnx.export()
+            else: other arguments except `obj` for `torch.jit.script()` to convert model, for more details:
             https://pytorch.org/docs/master/generated/torch.jit.script.html.
 
     """
@@ -642,6 +671,7 @@ def convert_to_onnx(
         if use_trace:
             # let torch.onnx.export to trace the model.
             mode_to_export = model
+            torch_versioned_kwargs = kwargs
         else:
             if not pytorch_after(1, 10):
                 if "example_outputs" not in kwargs:
@@ -654,32 +684,37 @@ def convert_to_onnx(
                 del kwargs["example_outputs"]
             mode_to_export = torch.jit.script(model, **kwargs)
 
+        if torch.is_tensor(inputs) or isinstance(inputs, dict):
+            onnx_inputs = (inputs,)
+        else:
+            onnx_inputs = tuple(inputs)
+
         if filename is None:
             f = io.BytesIO()
-            torch.onnx.export(
-                mode_to_export,
-                tuple(inputs),
-                f=f,
-                input_names=input_names,
-                output_names=output_names,
-                dynamic_axes=dynamic_axes,
-                opset_version=opset_version,
-                **torch_versioned_kwargs,
-            )
+        else:
+            f = filename
+
+        torch.onnx.export(
+            mode_to_export,
+            onnx_inputs,
+            f=f,
+            input_names=input_names,
+            output_names=output_names,
+            dynamic_axes=dynamic_axes,
+            opset_version=opset_version,
+            do_constant_folding=do_constant_folding,
+            **torch_versioned_kwargs,
+        )
+        if filename is None:
             onnx_model = onnx.load_model_from_string(f.getvalue())
         else:
-            torch.onnx.export(
-                mode_to_export,
-                tuple(inputs),
-                f=filename,
-                input_names=input_names,
-                output_names=output_names,
-                dynamic_axes=dynamic_axes,
-                opset_version=opset_version,
-                **torch_versioned_kwargs,
-            )
             onnx_model = onnx.load(filename)
 
+    if do_constant_folding and polygraphy_imported:
+        from polygraphy.backend.onnx.loader import fold_constants
+
+        fold_constants(onnx_model, size_threshold=constant_size_threshold)
+
     if verify:
         if device is None:
             device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -814,7 +849,6 @@ def _onnx_trt_compile(
 
     """
     trt, _ = optional_import("tensorrt", "8.5.3")
-    torch_tensorrt, _ = optional_import("torch_tensorrt", "1.4.0")
 
     input_shapes = (min_shape, opt_shape, max_shape)
     # default to an empty list to fit the `torch_tensorrt.ts.embed_engine_in_new_module` function.
@@ -916,8 +950,6 @@ def convert_to_trt(
             to compile model, for more details: https://pytorch.org/TensorRT/py_api/torch_tensorrt.html#torch-tensorrt-py.
     """
 
-    torch_tensorrt, _ = optional_import("torch_tensorrt", version="1.4.0")
-
     if not torch.cuda.is_available():
         raise Exception("Cannot find any GPU devices.")
 
@@ -935,23 +967,9 @@ def convert_to_trt(
     convert_precision = torch.float32 if precision == "fp32" else torch.half
     inputs = [torch.rand(ensure_tuple(input_shape)).to(target_device)]
 
-    def scale_batch_size(input_shape: Sequence[int], scale_num: int):
-        scale_shape = [*input_shape]
-        scale_shape[0] *= scale_num
-        return scale_shape
-
-    # Use the dynamic batchsize range to generate the min, opt and max model input shape
-    if dynamic_batchsize:
-        min_input_shape = scale_batch_size(input_shape, dynamic_batchsize[0])
-        opt_input_shape = scale_batch_size(input_shape, dynamic_batchsize[1])
-        max_input_shape = scale_batch_size(input_shape, dynamic_batchsize[2])
-    else:
-        min_input_shape = opt_input_shape = max_input_shape = input_shape
-
     # convert the torch model to a TorchScript model on target device
     model = model.eval().to(target_device)
-    ir_model = convert_to_torchscript(model, device=target_device, inputs=inputs, use_trace=use_trace)
-    ir_model.eval()
+    min_input_shape, opt_input_shape, max_input_shape = get_profile_shapes(input_shape, dynamic_batchsize)
 
     if use_onnx:
         # set the batch dim as dynamic
@@ -960,7 +978,6 @@ def convert_to_trt(
         ir_model = convert_to_onnx(
             model, inputs, onnx_input_names, onnx_output_names, use_trace=use_trace, dynamic_axes=dynamic_axes
         )
-
         # convert the model through the ONNX-TensorRT way
         trt_model = _onnx_trt_compile(
             ir_model,
@@ -973,6 +990,8 @@ def convert_to_trt(
             output_names=onnx_output_names,
         )
     else:
+        ir_model = convert_to_torchscript(model, device=target_device, inputs=inputs, use_trace=use_trace)
+        ir_model.eval()
         # convert the model through the Torch-TensorRT way
         ir_model.to(target_device)
         with torch.no_grad():
@@ -1189,3 +1208,168 @@ class CastTempType(nn.Module):
         if dtype == self.initial_type:
             x = x.to(self.initial_type)
         return x
+
+
+def cast_tensor(x, from_dtype=torch.float16, to_dtype=torch.float32):
+    """
+    Utility function to cast a single tensor from from_dtype to to_dtype
+    """
+    return x.to(dtype=to_dtype) if x.dtype == from_dtype else x
+
+
+def cast_all(x, from_dtype=torch.float16, to_dtype=torch.float32):
+    """
+    Utility function to cast all tensors in a tuple from from_dtype to to_dtype
+    """
+    if isinstance(x, torch.Tensor):
+        return cast_tensor(x, from_dtype=from_dtype, to_dtype=to_dtype)
+    else:
+        if isinstance(x, dict):
+            new_dict = {}
+            for k in x.keys():
+                new_dict[k] = cast_all(x[k], from_dtype=from_dtype, to_dtype=to_dtype)
+            return new_dict
+        elif isinstance(x, tuple):
+            return tuple(cast_all(y, from_dtype=from_dtype, to_dtype=to_dtype) for y in x)
+
+
+class CastToFloat(torch.nn.Module):
+    """
+    Class used to add autocast protection for ONNX export
+    for forward methods with single return vaue
+    """
+
+    def __init__(self, mod):
+        super().__init__()
+        self.mod = mod
+
+    def forward(self, x):
+        dtype = x.dtype
+        with torch.amp.autocast("cuda", enabled=False):
+            ret = self.mod.forward(x.to(torch.float32)).to(dtype)
+        return ret
+
+
+class CastToFloatAll(torch.nn.Module):
+    """
+    Class used to add autocast protection for ONNX export
+    for forward methods with multiple return values
+    """
+
+    def __init__(self, mod):
+        super().__init__()
+        self.mod = mod
+
+    def forward(self, *args):
+        from_dtype = args[0].dtype
+        with torch.amp.autocast("cuda", enabled=False):
+            ret = self.mod.forward(*cast_all(args, from_dtype=from_dtype, to_dtype=torch.float32))
+        return cast_all(ret, from_dtype=torch.float32, to_dtype=from_dtype)
+
+
+def wrap_module(base_t: type[nn.Module], dest_t: type[nn.Module]) -> Callable[[nn.Module], nn.Module | None]:
+    """
+    Generic function generator to replace base_t module with dest_t wrapper.
+    Args:
+        base_t : module type to replace
+        dest_t : destination module type
+    Returns:
+        swap function to replace base_t module with dest_t
+    """
+
+    def expansion_fn(mod: nn.Module) -> nn.Module | None:
+        out = dest_t(mod)
+        return out
+
+    return expansion_fn
+
+
+def simple_replace(base_t: type[nn.Module], dest_t: type[nn.Module]) -> Callable[[nn.Module], nn.Module | None]:
+    """
+    Generic function generator to replace base_t module with dest_t.
+    base_t and dest_t should have same atrributes. No weights are copied.
+    Args:
+        base_t : module type to replace
+        dest_t : destination module type
+    Returns:
+        swap function to replace base_t module with dest_t
+    """
+
+    def expansion_fn(mod: nn.Module) -> nn.Module | None:
+        if not isinstance(mod, base_t):
+            return None
+        args = [getattr(mod, name, None) for name in mod.__constants__]
+        out = dest_t(*args)
+        return out
+
+    return expansion_fn
+
+
+def _swap_modules(model: nn.Module, mapping: dict[str, nn.Module]) -> nn.Module:
+    """
+    This function swaps nested modules as specified by "dot paths" in mod with a desired replacement. This allows
+    for swapping nested modules through arbitrary levels if children
+
+    NOTE: This occurs in place, if you want to preserve model then make sure to copy it first.
+
+    """
+    for path, new_mod in mapping.items():
+        expanded_path = path.split(".")
+        parent_mod = model
+        for sub_path in expanded_path[:-1]:
+            submod = parent_mod._modules[sub_path]
+            if submod is None:
+                break
+            else:
+                parent_mod = submod
+        parent_mod._modules[expanded_path[-1]] = new_mod
+
+    return model
+
+
+def replace_modules_by_type(
+    model: nn.Module, expansions: dict[str, Callable[[nn.Module], nn.Module | None]]
+) -> nn.Module:
+    """
+    Top-level function to replace modules in model, specified by class name with a desired replacement.
+    NOTE: This occurs in place, if you want to preserve model then make sure to copy it first.
+    Args:
+        model : top level module
+        expansions : replacement dictionary: module class name -> replacement function generator
+    Returns:
+        model, possibly modified in-place
+    """
+    mapping: dict[str, nn.Module] = {}
+    for name, m in model.named_modules():
+        m_type = type(m).__name__
+        if m_type in expansions:
+            # print (f"Found {m_type} in expansions ...")
+            swapped = expansions[m_type](m)
+            if swapped:
+                mapping[name] = swapped
+
+    print(f"Swapped {len(mapping)} modules")
+    _swap_modules(model, mapping)
+    return model
+
+
+def add_casts_around_norms(model: nn.Module) -> nn.Module:
+    """
+    Top-level function to add cast wrappers around modules known to cause issues for FP16/autocast ONNX export
+    NOTE: This occurs in place, if you want to preserve model then make sure to copy it first.
+    Args:
+        model : top level module
+    Returns:
+        model, possibly modified in-place
+    """
+    print("Adding casts around norms...")
+    cast_replacements = {
+        "BatchNorm1d": wrap_module(nn.BatchNorm1d, CastToFloat),
+        "BatchNorm2d": wrap_module(nn.BatchNorm2d, CastToFloat),
+        "BatchNorm3d": wrap_module(nn.BatchNorm2d, CastToFloat),
+        "LayerNorm": wrap_module(nn.LayerNorm, CastToFloat),
+        "InstanceNorm1d": wrap_module(nn.InstanceNorm1d, CastToFloat),
+        "InstanceNorm3d": wrap_module(nn.InstanceNorm3d, CastToFloat),
+    }
+    replace_modules_by_type(model, cast_replacements)
+    return model

monai/transforms/__init__.py

@@ -396,6 +396,8 @@ from .smooth_field.dictionary import (
 from .spatial.array import (
     Affine,
     AffineGrid,
+    ConvertBoxToPoints,
+    ConvertPointsToBoxes,
     Flip,
     GridDistortion,
     GridPatch,
@@ -427,6 +429,12 @@ from .spatial.dictionary import (
     Affined,
     AffineD,
     AffineDict,
+    ConvertBoxToPointsd,
+    ConvertBoxToPointsD,
+    ConvertBoxToPointsDict,
+    ConvertPointsToBoxesd,
+    ConvertPointsToBoxesD,
+    ConvertPointsToBoxesDict,
     Flipd,
     FlipD,
     FlipDict,
@@ -503,6 +511,7 @@ from .transform import LazyTransform, MapTransform, Randomizable, RandomizableTr
 from .utility.array import (
     AddCoordinateChannels,
     AddExtremePointsChannel,
+    ApplyTransformToPoints,
     AsChannelLast,
     CastToType,
     ClassesToIndices,
@@ -542,6 +551,9 @@ from .utility.dictionary import (
     AddExtremePointsChanneld,
     AddExtremePointsChannelD,
     AddExtremePointsChannelDict,
+    ApplyTransformToPointsd,
+    ApplyTransformToPointsD,
+    ApplyTransformToPointsDict,
     AsChannelLastd,
     AsChannelLastD,
     AsChannelLastDict,

monai/transforms/spatial/array.py

@@ -25,6 +25,7 @@ import torch
 
 from monai.config import USE_COMPILED, DtypeLike
 from monai.config.type_definitions import NdarrayOrTensor
+from monai.data.box_utils import BoxMode, StandardMode
 from monai.data.meta_obj import get_track_meta, set_track_meta
 from monai.data.meta_tensor import MetaTensor
 from monai.data.utils import AFFINE_TOL, affine_to_spacing, compute_shape_offset, iter_patch, to_affine_nd, zoom_affine
@@ -34,6 +35,8 @@ from monai.transforms.croppad.array import CenterSpatialCrop, ResizeWithPadOrCro
 from monai.transforms.inverse import InvertibleTransform
 from monai.transforms.spatial.functional import (
     affine_func,
+    convert_box_to_points,
+    convert_points_to_box,
     flip,
     orientation,
     resize,
@@ -3544,3 +3547,44 @@ class RandSimulateLowResolution(RandomizableTransform):
 
         else:
             return img
+
+
+class ConvertBoxToPoints(Transform):
+    """
+    Converts an axis-aligned bounding box to points. It can automatically convert the boxes to the points based on the box mode.
+    Bounding boxes of the shape (N, C) for N boxes. C is [x1, y1, x2, y2] for 2D or [x1, y1, z1, x2, y2, z2] for 3D for each box.
+    Return shape will be (N, 4, 2) for 2D or (N, 8, 3) for 3D.
+    """
+
+    backend = [TransformBackends.TORCH, TransformBackends.NUMPY]
+
+    def __init__(self, mode: str | BoxMode | type[BoxMode] | None = None) -> None:
+        """
+        Args:
+            mode: the mode of the box, can be a string, a BoxMode instance or a BoxMode class. Defaults to StandardMode.
+        """
+        super().__init__()
+        self.mode = StandardMode if mode is None else mode
+
+    def __call__(self, data: Any):
+        data = convert_to_tensor(data, track_meta=get_track_meta())
+        points = convert_box_to_points(data, mode=self.mode)
+        return convert_to_dst_type(points, data)[0]
+
+
+class ConvertPointsToBoxes(Transform):
+    """
+    Converts points to an axis-aligned bounding box.
+    Points representing the corners of the bounding box. Shape (N, 8, 3) for the 8 corners of a 3D cuboid or
+    (N, 4, 2) for the 4 corners of a 2D rectangle.
+    """
+
+    backend = [TransformBackends.TORCH, TransformBackends.NUMPY]
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def __call__(self, data: Any):
+        data = convert_to_tensor(data, track_meta=get_track_meta())
+        box = convert_points_to_box(data)
+        return convert_to_dst_type(box, data)[0]

monai/transforms/spatial/dictionary.py

@@ -26,6 +26,7 @@ import torch
 
 from monai.config import DtypeLike, KeysCollection, SequenceStr
 from monai.config.type_definitions import NdarrayOrTensor
+from monai.data.box_utils import BoxMode, StandardMode
 from monai.data.meta_obj import get_track_meta
 from monai.data.meta_tensor import MetaTensor
 from monai.networks.layers.simplelayers import GaussianFilter
@@ -33,6 +34,8 @@ from monai.transforms.croppad.array import CenterSpatialCrop
 from monai.transforms.inverse import InvertibleTransform
 from monai.transforms.spatial.array import (
     Affine,
+    ConvertBoxToPoints,
+    ConvertPointsToBoxes,
     Flip,
     GridDistortion,
     GridPatch,
@@ -2585,6 +2588,7 @@ class RandSimulateLowResolutiond(RandomizableTransform, MapTransform):
         self, seed: int | None = None, state: np.random.RandomState | None = None
     ) -> RandSimulateLowResolutiond:
         super().set_random_state(seed, state)
+        self.sim_lowres_tfm.set_random_state(seed, state)
         return self
 
     def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]:
@@ -2611,6 +2615,61 @@ class RandSimulateLowResolutiond(RandomizableTransform, MapTransform):
         return d
 
 
+class ConvertBoxToPointsd(MapTransform):
+    """
+    Dictionary-based wrapper of :py:class:`monai.transforms.ConvertBoxToPoints`.
+    """
+
+    backend = ConvertBoxToPoints.backend
+
+    def __init__(
+        self,
+        keys: KeysCollection,
+        point_key="points",
+        mode: str | BoxMode | type[BoxMode] | None = StandardMode,
+        allow_missing_keys: bool = False,
+    ):
+        """
+        Args:
+            keys: keys of the corresponding items to be transformed.
+            point_key: key to store the point data.
+            mode: the mode of the input boxes. Defaults to StandardMode.
+            allow_missing_keys: don't raise exception if key is missing.
+        """
+        super().__init__(keys, allow_missing_keys)
+        self.point_key = point_key
+        self.converter = ConvertBoxToPoints(mode=mode)
+
+    def __call__(self, data):
+        d = dict(data)
+        for key in self.key_iterator(d):
+            data[self.point_key] = self.converter(d[key])
+        return data
+
+
+class ConvertPointsToBoxesd(MapTransform):
+    """
+    Dictionary-based wrapper of :py:class:`monai.transforms.ConvertPointsToBoxes`.
+    """
+
+    def __init__(self, keys: KeysCollection, box_key="box", allow_missing_keys: bool = False):
+        """
+        Args:
+            keys: keys of the corresponding items to be transformed.
+            box_key: key to store the box data.
+            allow_missing_keys: don't raise exception if key is missing.
+        """
+        super().__init__(keys, allow_missing_keys)
+        self.box_key = box_key
+        self.converter = ConvertPointsToBoxes()
+
+    def __call__(self, data):
+        d = dict(data)
+        for key in self.key_iterator(d):
+            data[self.box_key] = self.converter(d[key])
+        return data
+
+
 SpatialResampleD = SpatialResampleDict = SpatialResampled
 ResampleToMatchD = ResampleToMatchDict = ResampleToMatchd
 SpacingD = SpacingDict = Spacingd
@@ -2635,3 +2694,5 @@ GridSplitD = GridSplitDict = GridSplitd
 GridPatchD = GridPatchDict = GridPatchd
 RandGridPatchD = RandGridPatchDict = RandGridPatchd
 RandSimulateLowResolutionD = RandSimulateLowResolutionDict = RandSimulateLowResolutiond
+ConvertBoxToPointsD = ConvertBoxToPointsDict = ConvertBoxToPointsd
+ConvertPointsToBoxesD = ConvertPointsToBoxesDict = ConvertPointsToBoxesd

monai/transforms/spatial/functional.py

@@ -24,6 +24,7 @@ import torch
 import monai
 from monai.config import USE_COMPILED
 from monai.config.type_definitions import NdarrayOrTensor
+from monai.data.box_utils import get_boxmode
 from monai.data.meta_obj import get_track_meta
 from monai.data.meta_tensor import MetaTensor
 from monai.data.utils import AFFINE_TOL, compute_shape_offset, to_affine_nd
@@ -32,7 +33,7 @@ from monai.transforms.croppad.array import ResizeWithPadOrCrop
 from monai.transforms.intensity.array import GaussianSmooth
 from monai.transforms.inverse import TraceableTransform
 from monai.transforms.utils import create_rotate, create_translate, resolves_modes, scale_affine
-from monai.transforms.utils_pytorch_numpy_unification import allclose
+from monai.transforms.utils_pytorch_numpy_unification import allclose, concatenate, stack
 from monai.utils import (
     LazyAttr,
     TraceKeys,
@@ -610,3 +611,71 @@ def affine_func(
         out = _maybe_new_metatensor(img, dtype=torch.float32, device=resampler.device)
     out = out.copy_meta_from(meta_info) if isinstance(out, MetaTensor) else out
     return out if image_only else (out, affine)
+
+
+def convert_box_to_points(bbox, mode):
+    """
+    Converts an axis-aligned bounding box to points.
+
+    Args:
+        mode: The mode specifying how to interpret the bounding box.
+        bbox: Bounding boxes of the shape (N, C) for N boxes. C is [x1, y1, x2, y2] for 2D or [x1, y1, z1, x2, y2, z2]
+            for 3D for each box. Return shape will be (N, 4, 2) for 2D or (N, 8, 3) for 3D.
+
+    Returns:
+        sequence of points representing the corners of the bounding box.
+    """
+
+    mode = get_boxmode(mode)
+
+    points_list = []
+    for _num in range(bbox.shape[0]):
+        corners = mode.boxes_to_corners(bbox[_num : _num + 1])
+        if len(corners) == 4:
+            points_list.append(
+                concatenate(
+                    [
+                        concatenate([corners[0], corners[1]], axis=1),
+                        concatenate([corners[2], corners[1]], axis=1),
+                        concatenate([corners[2], corners[3]], axis=1),
+                        concatenate([corners[0], corners[3]], axis=1),
+                    ],
+                    axis=0,
+                )
+            )
+        else:
+            points_list.append(
+                concatenate(
+                    [
+                        concatenate([corners[0], corners[1], corners[2]], axis=1),
+                        concatenate([corners[3], corners[1], corners[2]], axis=1),
+                        concatenate([corners[3], corners[4], corners[2]], axis=1),
+                        concatenate([corners[0], corners[4], corners[2]], axis=1),
+                        concatenate([corners[0], corners[1], corners[5]], axis=1),
+                        concatenate([corners[3], corners[1], corners[5]], axis=1),
+                        concatenate([corners[3], corners[4], corners[5]], axis=1),
+                        concatenate([corners[0], corners[4], corners[5]], axis=1),
+                    ],
+                    axis=0,
+                )
+            )
+
+    return stack(points_list, dim=0)
+
+
+def convert_points_to_box(points):
+    """
+    Converts points to an axis-aligned bounding box.
+
+    Args:
+        points: Points representing the corners of the bounding box. Shape (N, 8, 3) for the 8 corners of
+            a 3D cuboid or (N, 4, 2) for the 4 corners of a 2D rectangle.
+    """
+    from monai.transforms.utils_pytorch_numpy_unification import max, min
+
+    mins = min(points, dim=1)
+    maxs = max(points, dim=1)
+    # Concatenate the min and max values to get the bounding boxes
+    bboxes = concatenate([mins, maxs], axis=1)
+
+    return bboxes