Repository: Project-MONAI/MONAI

Issue:
Porting weighted sampling from NiftyNet
**Is your feature request related to a problem? Please describe.**
Niftynet's weighted sampler could now be implemented as a monai transform

context:
https://github.com/NifTK/NiftyNet/blob/dev/niftynet/engine/sampler_weighted_v2.py


Discussion:

Resolving patch:
diff --git a/monai/transforms/croppad/array.py b/monai/transforms/croppad/array.py
--- a/monai/transforms/croppad/array.py
+++ b/monai/transforms/croppad/array.py
@@ -24,6 +24,7 @@
     generate_pos_neg_label_crop_centers,
     generate_spatial_bounding_box,
     map_binary_to_indices,
+    weighted_patch_samples,
 )
 from monai.utils import Method, NumpyPadMode, ensure_tuple, fall_back_tuple
 
@@ -414,6 +415,58 @@ def __call__(self, img: np.ndarray):
         return cropped
 
 
+class RandWeightedCrop(Randomizable, Transform):
+    """
+    Samples a list of `num_samples` image patches according to the provided `weight_map`.
+
+    Args:
+        spatial_size: the spatial size of the image patch e.g. [224, 224, 128].
+            If its components have non-positive values, the corresponding size of `img` will be used.
+        num_samples: number of samples (image patches) to take in the returned list.
+        weight_map: weight map used to generate patch samples. The weights must be non-negative.
+            Each element denotes a sampling weight of the spatial location. 0 indicates no sampling.
+            It should be a single-channel array in shape, for example, `(1, spatial_dim_0, spatial_dim_1, ...)`.
+    """
+
+    def __init__(
+        self, spatial_size: Union[Sequence[int], int], num_samples: int = 1, weight_map: Optional[np.ndarray] = None
+    ):
+        self.spatial_size = ensure_tuple(spatial_size)
+        self.num_samples = int(num_samples)
+        self.weight_map = weight_map
+        self.centers: List[np.ndarray] = []
+
+    def randomize(self, weight_map: np.ndarray) -> None:
+        self.centers = weighted_patch_samples(
+            spatial_size=self.spatial_size, w=weight_map[0], n_samples=self.num_samples, r_state=self.R
+        )  # using only the first channel as weight map
+
+    def __call__(self, img: np.ndarray, weight_map: Optional[np.ndarray] = None) -> List[np.ndarray]:
+        """
+        Args:
+            img: input image to sample patches from. assuming `img` is a channel-first array.
+            weight_map: weight map used to generate patch samples. The weights must be non-negative.
+                Each element denotes a sampling weight of the spatial location. 0 indicates no sampling.
+                It should be a single-channel array in shape, for example, `(1, spatial_dim_0, spatial_dim_1, ...)`
+
+        Returns:
+            A list of image patches
+        """
+        if weight_map is None:
+            weight_map = self.weight_map
+        if weight_map is None:
+            raise ValueError("weight map must be provided for weighted patch sampling.")
+        if img.shape[1:] != weight_map.shape[1:]:
+            raise ValueError(f"image and weight map spatial shape mismatch: {img.shape[1:]} vs {weight_map.shape[1:]}.")
+        self.randomize(weight_map)
+        _spatial_size = fall_back_tuple(self.spatial_size, weight_map.shape[1:])
+        results = []
+        for center in self.centers:
+            cropper = SpatialCrop(roi_center=center, roi_size=_spatial_size)
+            results.append(cropper(img))
+        return results
+
+
 class RandCropByPosNegLabel(Randomizable, Transform):
     """
     Crop random fixed sized regions with the center being a foreground or background voxel
diff --git a/monai/transforms/croppad/dictionary.py b/monai/transforms/croppad/dictionary.py
--- a/monai/transforms/croppad/dictionary.py
+++ b/monai/transforms/croppad/dictionary.py
@@ -34,6 +34,7 @@
     generate_pos_neg_label_crop_centers,
     generate_spatial_bounding_box,
     map_binary_to_indices,
+    weighted_patch_samples,
 )
 from monai.utils import Method, NumpyPadMode, ensure_tuple, ensure_tuple_rep, fall_back_tuple
 
@@ -377,6 +378,70 @@ def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.nda
         return d
 
 
+class RandWeightedCropd(Randomizable, MapTransform):
+    """
+    Samples a list of `num_samples` image patches according to the provided `weight_map`.
+
+    Args:
+        keys: keys of the corresponding items to be transformed.
+            See also: :py:class:`monai.transforms.compose.MapTransform`
+        w_key: key for the weight map. The corresponding value will be used as the sampling weights,
+            it should be a single-channel array in size, for example, `(1, spatial_dim_0, spatial_dim_1, ...)`
+        spatial_size: the spatial size of the image patch e.g. [224, 224, 128].
+            If its components have non-positive values, the corresponding size of `img` will be used.
+        num_samples: number of samples (image patches) to take in the returned list.
+        center_coord_key: if specified, the actual sampling location will be stored with the corresponding key.
+
+    See Also:
+        :py:class:`monai.transforms.RandWeightedCrop`
+    """
+
+    def __init__(
+        self,
+        keys: KeysCollection,
+        w_key: str,
+        spatial_size: Union[Sequence[int], int],
+        num_samples: int = 1,
+        center_coord_key: Optional[str] = None,
+    ):
+        super().__init__(keys)
+        self.spatial_size = ensure_tuple(spatial_size)
+        self.w_key = w_key
+        self.num_samples = int(num_samples)
+        self.center_coord_key = center_coord_key
+        self.centers: List[np.ndarray] = []
+
+    def randomize(self, weight_map: np.ndarray) -> None:
+        self.centers = weighted_patch_samples(
+            spatial_size=self.spatial_size, w=weight_map[0], n_samples=self.num_samples, r_state=self.R
+        )
+
+    def __call__(self, data: Mapping[Hashable, np.ndarray]) -> List[Dict[Hashable, np.ndarray]]:
+        d = dict(data)
+        self.randomize(d[self.w_key])
+        _spatial_size = fall_back_tuple(self.spatial_size, d[self.w_key].shape[1:])
+
+        results: List[Dict[Hashable, np.ndarray]] = [dict() for _ in range(self.num_samples)]
+        for key in data.keys():
+            if key in self.keys:
+                img = d[key]
+                if img.shape[1:] != d[self.w_key].shape[1:]:
+                    raise ValueError(
+                        f"data {key} and weight map {self.w_key} spatial shape mismatch: "
+                        f"{img.shape[1:]} vs {d[self.w_key].shape[1:]}."
+                    )
+                for i, center in enumerate(self.centers):
+                    cropper = SpatialCrop(roi_center=center, roi_size=_spatial_size)
+                    results[i][key] = cropper(img)
+                    if self.center_coord_key:
+                        results[i][self.center_coord_key] = center
+            else:
+                for i in range(self.num_samples):
+                    results[i][key] = data[key]
+
+        return results
+
+
 class RandCropByPosNegLabeld(Randomizable, MapTransform):
     """
     Dictionary-based version :py:class:`monai.transforms.RandCropByPosNegLabel`.
@@ -523,5 +588,6 @@ def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.nda
 RandSpatialCropD = RandSpatialCropDict = RandSpatialCropd
 RandSpatialCropSamplesD = RandSpatialCropSamplesDict = RandSpatialCropSamplesd
 CropForegroundD = CropForegroundDict = CropForegroundd
+RandWeightedCropD = RandWeightedCropDict = RandWeightedCropd
 RandCropByPosNegLabelD = RandCropByPosNegLabelDict = RandCropByPosNegLabeld
 ResizeWithPadOrCropD = ResizeWithPadOrCropDict = ResizeWithPadOrCropd
diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py
--- a/monai/transforms/utils.py
+++ b/monai/transforms/utils.py
@@ -106,16 +106,18 @@ def rescale_array_int_max(arr: np.ndarray, dtype: np.dtype = np.uint16) -> np.nd
 
 
 def copypaste_arrays(
-    src: np.ndarray,
-    dest: np.ndarray,
+    src_shape,
+    dest_shape,
     srccenter: Sequence[int],
     destcenter: Sequence[int],
     dims: Sequence[Optional[int]],
 ) -> Tuple[Tuple[slice, ...], Tuple[slice, ...]]:
     """
-    Calculate the slices to copy a sliced area of array `src` into array `dest`. The area has dimensions `dims` (use 0
-    or None to copy everything in that dimension), the source area is centered at `srccenter` index in `src` and copied
-    into area centered at `destcenter` in `dest`. The dimensions of the copied area will be clipped to fit within the
+    Calculate the slices to copy a sliced area of array in `src_shape` into array in `dest_shape`.
+
+    The area has dimensions `dims` (use 0 or None to copy everything in that dimension),
+    the source area is centered at `srccenter` index in `src` and copied into area centered at `destcenter` in `dest`.
+    The dimensions of the copied area will be clipped to fit within the
     source and destination arrays so a smaller area may be copied than expected. Return value is the tuples of slice
     objects indexing the copied area in `src`, and those indexing the copy area in `dest`.
 
@@ -123,9 +125,10 @@ def copypaste_arrays(
 
     .. code-block:: python
 
-        src = np.random.randint(0,10,(6,6))
+        src_shape = (6,6)
+        src = np.random.randint(0,10,src_shape)
         dest = np.zeros_like(src)
-        srcslices, destslices = copypaste_arrays(src, dest, (3, 2),(2, 1),(3, 4))
+        srcslices, destslices = copypaste_arrays(src_shape, dest.shape, (3, 2),(2, 1),(3, 4))
         dest[destslices] = src[srcslices]
         print(src)
         print(dest)
@@ -144,10 +147,12 @@ def copypaste_arrays(
              [0 0 0 0 0 0]]
 
     """
-    srcslices = [slice(None)] * src.ndim
-    destslices = [slice(None)] * dest.ndim
+    s_ndim = len(src_shape)
+    d_ndim = len(dest_shape)
+    srcslices = [slice(None)] * s_ndim
+    destslices = [slice(None)] * d_ndim
 
-    for i, ss, ds, sc, dc, dim in zip(range(src.ndim), src.shape, dest.shape, srccenter, destcenter, dims):
+    for i, ss, ds, sc, dc, dim in zip(range(s_ndim), src_shape, dest_shape, srccenter, destcenter, dims):
         if dim:
             # dimension before midpoint, clip to size fitting in both arrays
             d1 = np.clip(dim // 2, 0, min(sc, dc))
@@ -160,23 +165,27 @@ def copypaste_arrays(
     return tuple(srcslices), tuple(destslices)
 
 
-def resize_center(img: np.ndarray, *resize_dims: Optional[int], fill_value: float = 0.0) -> np.ndarray:
+def resize_center(
+    img: np.ndarray, *resize_dims: Optional[int], fill_value: float = 0.0, inplace: bool = True
+) -> np.ndarray:
     """
     Resize `img` by cropping or expanding the image from the center. The `resize_dims` values are the output dimensions
     (or None to use original dimension of `img`). If a dimension is smaller than that of `img` then the result will be
     cropped and if larger padded with zeros, in both cases this is done relative to the center of `img`. The result is
     a new image with the specified dimensions and values from `img` copied into its center.
     """
-    resize_dims = tuple(resize_dims[i] or img.shape[i] for i in range(len(resize_dims)))
 
-    dest = np.full(resize_dims, fill_value, img.dtype)
-    half_img_shape = np.asarray(img.shape) // 2
-    half_dest_shape = np.asarray(dest.shape) // 2
+    resize_dims = fall_back_tuple(resize_dims, img.shape)
 
-    srcslices, destslices = copypaste_arrays(img, dest, half_img_shape, half_dest_shape, resize_dims)
-    dest[destslices] = img[srcslices]
+    half_img_shape = np.asarray(img.shape) // 2
+    half_dest_shape = np.asarray(resize_dims) // 2
+    srcslices, destslices = copypaste_arrays(img.shape, resize_dims, half_img_shape, half_dest_shape, resize_dims)
 
-    return dest
+    if not inplace:
+        dest = np.full(resize_dims, fill_value, img.dtype)
+        dest[destslices] = img[srcslices]
+        return dest
+    return img[srcslices]
 
 
 def map_binary_to_indices(
@@ -211,6 +220,51 @@ def map_binary_to_indices(
     return fg_indices, bg_indices
 
 
+def weighted_patch_samples(
+    spatial_size: Union[int, Sequence[int]],
+    w: np.ndarray,
+    n_samples: int = 1,
+    r_state: Optional[np.random.RandomState] = None,
+) -> List:
+    """
+    Computes `n_samples` of random patch sampling locations, given the sampling weight map `w` and patch `spatial_size`.
+
+    Args:
+        spatial_size: length of each spatial dimension of the patch.
+        w: weight map, the weights must be non-negative. each element denotes a sampling weight of the spatial location.
+            0 indicates no sampling.
+            The weight map shape is assumed ``(spatial_dim_0, spatial_dim_1, ..., spatial_dim_n)``.
+        n_samples: number of patch samples
+        r_state: a random state container
+
+    Returns:
+        a list of `n_samples` N-D integers representing the spatial sampling location of patches.
+
+    """
+    if w is None:
+        raise ValueError("w must be an ND array.")
+    if r_state is None:
+        r_state = np.random.RandomState()
+    img_size = np.asarray(w.shape, dtype=np.int)
+    win_size = np.asarray(fall_back_tuple(spatial_size, img_size), dtype=np.int)
+
+    s = tuple(slice(w // 2, m - w + w // 2) if m > w else slice(m // 2, m // 2 + 1) for w, m in zip(win_size, img_size))
+    v = w[s]  # weight map in the 'valid' mode
+    v_size = v.shape
+    v = v.ravel()
+    if np.any(v < 0):
+        v -= np.min(v)  # shifting to non-negative
+    v = v.cumsum()
+    if not v[-1] or not np.isfinite(v[-1]) or v[-1] < 0:  # uniform sampling
+        idx = r_state.randint(0, len(v), size=n_samples)
+    else:
+        idx = v.searchsorted(r_state.random(n_samples) * v[-1], side="right")
+    # compensate 'valid' mode
+    diff = np.minimum(win_size, img_size) // 2
+    centers = [np.unravel_index(i, v_size) + diff for i in np.asarray(idx, dtype=np.int)]
+    return centers
+
+
 def generate_pos_neg_label_crop_centers(
     spatial_size: Union[Sequence[int], int],
     num_samples: int,


Test patch:
diff --git a/tests/test_rand_weighted_crop.py b/tests/test_rand_weighted_crop.py
new file mode 100644
--- /dev/null
+++ b/tests/test_rand_weighted_crop.py
@@ -0,0 +1,139 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+from monai.transforms.croppad.array import RandWeightedCrop
+from tests.utils import NumpyImageTestCase2D, NumpyImageTestCase3D
+
+
+class TestRandWeightedCrop2D(NumpyImageTestCase2D):
+    def test_rand_weighted_crop_small_roi(self):
+        img = self.seg1[0]
+        n_samples = 3
+        crop = RandWeightedCrop((10, 12), n_samples)
+        weight = np.zeros_like(img)
+        weight[0, 30, 17] = 1.1
+        weight[0, 40, 31] = 1
+        weight[0, 80, 21] = 1
+        crop.set_random_state(10)
+        result = crop(img, weight)
+        self.assertTrue(len(result) == n_samples)
+        np.testing.assert_allclose(result[0].shape, (1, 10, 12))
+        np.testing.assert_allclose(np.asarray(crop.centers), [[80, 21], [30, 17], [40, 31]])
+
+    def test_rand_weighted_crop_default_roi(self):
+        img = self.imt[0]
+        n_samples = 3
+        crop = RandWeightedCrop((10, -1), n_samples)
+        weight = np.zeros_like(img)
+        weight[0, 30, 17] = 1.1
+        weight[0, 40, 31] = 1
+        weight[0, 80, 21] = 1
+        crop.set_random_state(10)
+        result = crop(img, weight)
+        self.assertTrue(len(result) == n_samples)
+        np.testing.assert_allclose(result[0].shape, (1, 10, 64))
+        np.testing.assert_allclose(np.asarray(crop.centers), [[14, 32], [105, 32], [20, 32]])
+
+    def test_rand_weighted_crop_large_roi(self):
+        img = self.segn[0]
+        n_samples = 3
+        crop = RandWeightedCrop((10000, 400), n_samples)
+        weight = np.zeros_like(img)
+        weight[0, 30, 17] = 1.1
+        weight[0, 10, 1] = 1
+        crop.set_random_state(10)
+        result = crop(img, weight)
+        self.assertTrue(len(result) == n_samples)
+        np.testing.assert_allclose(result[0].shape, (1, 128, 64))
+        np.testing.assert_allclose(np.asarray(crop.centers), [[64, 32], [64, 32], [64, 32]])
+        for res in result:
+            np.testing.assert_allclose(res, self.segn[0])
+
+    def test_rand_weighted_crop_bad_w(self):
+        img = self.imt[0]
+        n_samples = 3
+        crop = RandWeightedCrop((20, 40), n_samples)
+        weight = np.zeros_like(img)
+        weight[0, 30, 17] = np.inf
+        weight[0, 10, 1] = -np.inf
+        weight[0, 10, 20] = -np.nan
+        crop.set_random_state(10)
+        result = crop(img, weight)
+        self.assertTrue(len(result) == n_samples)
+        np.testing.assert_allclose(result[0].shape, (1, 20, 40))
+        np.testing.assert_allclose(np.asarray(crop.centers), [[63, 37], [31, 43], [66, 20]])
+
+
+class TestRandWeightedCrop(NumpyImageTestCase3D):
+    def test_rand_weighted_crop_small_roi(self):
+        img = self.seg1[0]
+        n_samples = 3
+        crop = RandWeightedCrop((8, 10, 12), n_samples)
+        weight = np.zeros_like(img)
+        weight[0, 5, 30, 17] = 1.1
+        weight[0, 8, 40, 31] = 1
+        weight[0, 11, 23, 21] = 1
+        crop.set_random_state(10)
+        result = crop(img, weight)
+        self.assertTrue(len(result) == n_samples)
+        np.testing.assert_allclose(result[0].shape, (1, 8, 10, 12))
+        np.testing.assert_allclose(np.asarray(crop.centers), [[11, 23, 21], [5, 30, 17], [8, 40, 31]])
+
+    def test_rand_weighted_crop_default_roi(self):
+        img = self.imt[0]
+        n_samples = 3
+        crop = RandWeightedCrop((10, -1, -1), n_samples)
+        weight = np.zeros_like(img)
+        weight[0, 7, 17] = 1.1
+        weight[0, 13, 31] = 1.1
+        weight[0, 24, 21] = 1
+        crop.set_random_state(10)
+        result = crop(img, weight)
+        self.assertTrue(len(result) == n_samples)
+        np.testing.assert_allclose(result[0].shape, (1, 10, 64, 80))
+        np.testing.assert_allclose(np.asarray(crop.centers), [[14, 32, 40], [41, 32, 40], [20, 32, 40]])
+
+    def test_rand_weighted_crop_large_roi(self):
+        img = self.segn[0]
+        n_samples = 3
+        crop = RandWeightedCrop((10000, 400, 80), n_samples)
+        weight = np.zeros_like(img)
+        weight[0, 30, 17, 20] = 1.1
+        weight[0, 10, 1, 17] = 1
+        crop.set_random_state(10)
+        result = crop(img, weight)
+        self.assertTrue(len(result) == n_samples)
+        np.testing.assert_allclose(result[0].shape, (1, 48, 64, 80))
+        np.testing.assert_allclose(np.asarray(crop.centers), [[24, 32, 40], [24, 32, 40], [24, 32, 40]])
+        for res in result:
+            np.testing.assert_allclose(res, self.segn[0])
+
+    def test_rand_weighted_crop_bad_w(self):
+        img = self.imt[0]
+        n_samples = 3
+        crop = RandWeightedCrop((48, 64, 80), n_samples)
+        weight = np.zeros_like(img)
+        weight[0, 30, 17] = np.inf
+        weight[0, 10, 1] = -np.inf
+        weight[0, 10, 20] = -np.nan
+        crop.set_random_state(10)
+        result = crop(img, weight)
+        self.assertTrue(len(result) == n_samples)
+        np.testing.assert_allclose(result[0].shape, (1, 48, 64, 80))
+        np.testing.assert_allclose(np.asarray(crop.centers), [[24, 32, 40], [24, 32, 40], [24, 32, 40]])
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_rand_weighted_cropd.py b/tests/test_rand_weighted_cropd.py
new file mode 100644
--- /dev/null
+++ b/tests/test_rand_weighted_cropd.py
@@ -0,0 +1,144 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+from monai.transforms.croppad.dictionary import RandWeightedCropd
+from tests.utils import NumpyImageTestCase2D, NumpyImageTestCase3D
+
+
+class TestRandWeightedCrop(NumpyImageTestCase2D):
+    def test_rand_weighted_crop_small_roi(self):
+        img = self.seg1[0]
+        n_samples = 3
+        crop = RandWeightedCropd("img", "w", (10, 12), n_samples)
+        weight = np.zeros_like(img)
+        weight[0, 30, 17] = 1.1
+        weight[0, 40, 31] = 1
+        weight[0, 80, 21] = 1
+        crop.set_random_state(10)
+        d = {"img": img, "w": weight}
+        result = crop(d)
+        self.assertTrue(len(result) == n_samples)
+        np.testing.assert_allclose(result[0]["img"].shape, (1, 10, 12))
+        np.testing.assert_allclose(np.asarray(crop.centers), [[80, 21], [30, 17], [40, 31]])
+
+    def test_rand_weighted_crop_default_roi(self):
+        img = self.imt[0]
+        n_samples = 3
+        crop = RandWeightedCropd("im", "weight", (10, -1), n_samples, "coords")
+        weight = np.zeros_like(img)
+        weight[0, 30, 17] = 1.1
+        weight[0, 40, 31] = 1
+        weight[0, 80, 21] = 1
+        crop.set_random_state(10)
+        data = {"im": img, "weight": weight, "others": np.nan}
+        result = crop(data)
+        self.assertTrue(len(result) == n_samples)
+        np.testing.assert_allclose(result[0]["im"].shape, (1, 10, 64))
+        np.testing.assert_allclose(np.asarray(crop.centers), [[14, 32], [105, 32], [20, 32]])
+        np.testing.assert_allclose(result[1]["coords"], [105, 32])
+
+    def test_rand_weighted_crop_large_roi(self):
+        img = self.segn[0]
+        n_samples = 3
+        crop = RandWeightedCropd(("img", "seg"), "weight", (10000, 400), n_samples, "location")
+        weight = np.zeros_like(img)
+        weight[0, 30, 17] = 1.1
+        weight[0, 10, 1] = 1
+        crop.set_random_state(10)
+        data = {"img": img, "seg": self.imt[0], "weight": weight}
+        result = crop(data)
+        self.assertTrue(len(result) == n_samples)
+        np.testing.assert_allclose(result[0]["img"].shape, (1, 128, 64))
+        np.testing.assert_allclose(result[0]["seg"].shape, (1, 128, 64))
+        np.testing.assert_allclose(np.asarray(crop.centers), [[64, 32], [64, 32], [64, 32]])
+        np.testing.assert_allclose(result[1]["location"], [64, 32])
+
+    def test_rand_weighted_crop_bad_w(self):
+        img = self.imt[0]
+        n_samples = 3
+        crop = RandWeightedCropd(("img", "seg"), "w", (20, 40), n_samples)
+        weight = np.zeros_like(img)
+        weight[0, 30, 17] = np.inf
+        weight[0, 10, 1] = -np.inf
+        weight[0, 10, 20] = -np.nan
+        crop.set_random_state(10)
+        result = crop({"img": img, "seg": self.segn[0], "w": weight})
+        self.assertTrue(len(result) == n_samples)
+        np.testing.assert_allclose(result[0]["img"].shape, (1, 20, 40))
+        np.testing.assert_allclose(result[0]["seg"].shape, (1, 20, 40))
+        np.testing.assert_allclose(np.asarray(crop.centers), [[63, 37], [31, 43], [66, 20]])
+
+
+class TestRandWeightedCrop3D(NumpyImageTestCase3D):
+    def test_rand_weighted_crop_small_roi(self):
+        img = self.seg1[0]
+        n_samples = 3
+        crop = RandWeightedCropd("img", "w", (8, 10, 12), n_samples)
+        weight = np.zeros_like(img)
+        weight[0, 5, 30, 17] = 1.1
+        weight[0, 8, 40, 31] = 1
+        weight[0, 11, 23, 21] = 1
+        crop.set_random_state(10)
+        result = crop({"img": img, "w": weight})
+        self.assertTrue(len(result) == n_samples)
+        np.testing.assert_allclose(result[0]["img"].shape, (1, 8, 10, 12))
+        np.testing.assert_allclose(np.asarray(crop.centers), [[11, 23, 21], [5, 30, 17], [8, 40, 31]])
+
+    def test_rand_weighted_crop_default_roi(self):
+        img = self.imt[0]
+        n_samples = 3
+        crop = RandWeightedCropd(("img", "seg"), "w", (10, -1, -1), n_samples)
+        weight = np.zeros_like(img)
+        weight[0, 7, 17] = 1.1
+        weight[0, 13, 31] = 1.1
+        weight[0, 24, 21] = 1
+        crop.set_random_state(10)
+        result = crop({"img": img, "seg": self.segn[0], "w": weight})
+        self.assertTrue(len(result) == n_samples)
+        np.testing.assert_allclose(result[0]["img"].shape, (1, 10, 64, 80))
+        np.testing.assert_allclose(result[0]["seg"].shape, (1, 10, 64, 80))
+        np.testing.assert_allclose(np.asarray(crop.centers), [[14, 32, 40], [41, 32, 40], [20, 32, 40]])
+
+    def test_rand_weighted_crop_large_roi(self):
+        img = self.segn[0]
+        n_samples = 3
+        crop = RandWeightedCropd("img", "w", (10000, 400, 80), n_samples)
+        weight = np.zeros_like(img)
+        weight[0, 30, 17, 20] = 1.1
+        weight[0, 10, 1, 17] = 1
+        crop.set_random_state(10)
+        result = crop({"img": img, "w": weight})
+        self.assertTrue(len(result) == n_samples)
+        np.testing.assert_allclose(result[0]["img"].shape, (1, 48, 64, 80))
+        np.testing.assert_allclose(np.asarray(crop.centers), [[24, 32, 40], [24, 32, 40], [24, 32, 40]])
+
+    def test_rand_weighted_crop_bad_w(self):
+        img = self.imt[0]
+        n_samples = 3
+        crop = RandWeightedCropd(("img", "seg"), "w", (48, 64, 80), n_samples)
+        weight = np.zeros_like(img)
+        weight[0, 30, 17] = np.inf
+        weight[0, 10, 1] = -np.inf
+        weight[0, 10, 20] = -np.nan
+        crop.set_random_state(10)
+        result = crop({"img": img, "seg": self.segn[0], "w": weight})
+        self.assertTrue(len(result) == n_samples)
+        np.testing.assert_allclose(result[0]["img"].shape, (1, 48, 64, 80))
+        np.testing.assert_allclose(result[0]["seg"].shape, (1, 48, 64, 80))
+        np.testing.assert_allclose(np.asarray(crop.centers), [[24, 32, 40], [24, 32, 40], [24, 32, 40]])
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_rotated.py b/tests/test_rotated.py
--- a/tests/test_rotated.py
+++ b/tests/test_rotated.py
@@ -84,7 +84,7 @@ def test_correct_results(self, angle, keep_size, mode, padding_mode, align_corne
             self.segn[0, 0], np.rad2deg(angle), (0, 2), not keep_size, order=0, mode=_mode, prefilter=False
         )
         expected = np.stack(expected).astype(int)
-        self.assertLessEqual(np.count_nonzero(expected != rotated["seg"][0]), 100)
+        self.assertLessEqual(np.count_nonzero(expected != rotated["seg"][0]), 105)
 
 
 class TestRotated3DXY(NumpyImageTestCase3D):