Skip to content

Commit 02a8ef3

Browse files
authored
perf(RandCropBoxByPosNegLabeld): avoid N-fold mask allocation in randomize (#8993)
### Description `RandCropBoxByPosNegLabeld.randomize` builds a foreground/background mask when precomputed indices are not available. The previous implementation called `convert_box_to_mask` for each box (creating an N-channel stacked mask) then reduced with `np.amax` to a single channel, allocating an (N_boxes, H, W, D) int16 intermediate that grows linearly with the number of boxes. The fix paints directly into a single (1, H, W, D) union mask by iterating over each extended box and setting `mask[0, box_slices] = 1`. ### Changes **`monai/apps/detection/transforms/dictionary.py`** 1. Remove the now-unused import of `convert_box_to_mask` from `monai.apps.detection.transforms.box_ops`. 2. Replace the `convert_box_to_mask` + `np.amax(...)[0:1]` block in `randomize()` with direct slicing into a pre-allocated zero mask. ### Performance Benchmarked old (`convert_box_to_mask` + `np.amax`) vs new (direct slicing) implementations on synthetic 3D volumes with random boxes (`spatial_size=(64,64,64)`, `num_samples=4`, `pos=1`, `neg=1`, `whole_box=True`). **System specs** - CPU: 12th Gen Intel Core i7-12800H (14 cores / 20 threads, up to 4.8 GHz) - RAM: 30 GiB - OS: Ubuntu 22.04.5 LTS, kernel 6.8.0-124-generic - Python 3.10.12, NumPy 2.2.6, PyTorch 2.13.0+cu130 **`randomize()` only** | image_size | n_boxes | old (ms) | new (ms) | speedup | |---|---|---|---|---| | (96,96,96) | 5 | 2.471 | 1.331 | 1.86x | | (96,96,96) | 20 | 17.998 | 1.921 | 9.37x | | (160,160,160) | 5 | 22.931 | 15.241 | 1.50x | | (160,160,160) | 20 | 78.651 | 6.563 | 11.98x | | (256,256,128) | 5 | 62.477 | 26.386 | 2.37x | | (256,256,128) | 20 | 174.831 | 24.427 | 7.16x | **Full `__call__()`** | image_size | n_boxes | old (ms) | new (ms) | speedup | |---|---|---|---|---| | (96,96,96) | 5 | 6.416 | 4.823 | 1.33x | | (96,96,96) | 20 | 21.745 | 5.078 | 4.28x | | (160,160,160) | 5 | 22.820 | 13.778 | 1.66x | | (160,160,160) | 20 | 85.259 | 14.000 | 6.09x | | (256,256,128) | 5 | 76.623 | 38.455 | 1.99x | | (256,256,128) | 20 | 188.849 | 36.955 | 5.11x | ### Types of changes - [x] Non-breaking change (fix or new feature that would not break existing functionality). - [ ] Breaking change (fix or new feature that would cause existing functionality to change). - [ ] New tests added to cover the changes. - [ ] Integration tests passed locally by running `./runtests.sh -f -u --net --coverage`. - [x] Quick tests passed locally by running `./runtests.sh --quick --unittests --disttests`. - [ ] In-line docstrings updated. - [ ] Documentation updated, tested `make html` command in the `docs/` folder. --------- Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>
1 parent 8c5bede commit 02a8ef3

1 file changed

Lines changed: 7 additions & 5 deletions

File tree

‎monai/apps/detection/transforms/dictionary.py‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,6 @@
3737
StandardizeEmptyBox,
3838
ZoomBox,
3939
)
40-
from monai.apps.detection.transforms.box_ops import convert_box_to_mask
4140
from monai.config import KeysCollection, SequenceStr
4241
from monai.config.type_definitions import DtypeLike, NdarrayOrTensor
4342
from monai.data.box_utils import COMPUTE_DTYPE, BoxMode, clip_boxes_to_image
@@ -1154,10 +1153,13 @@ def randomize( # type: ignore
11541153
# As along as the cropped patch contains a box, it is considered as a foreground patch.
11551154
# Positions within extended_boxes are crop centers for foreground patches
11561155
extended_boxes_np = self.generate_fg_center_boxes_np(boxes, image_size)
1157-
mask_img = convert_box_to_mask(
1158-
extended_boxes_np, np.ones(extended_boxes_np.shape[0]), image_size, bg_label=0, ellipse_mask=False
1159-
)
1160-
mask_img = np.amax(mask_img, axis=0, keepdims=True)[0:1, ...]
1156+
spatial_dims = len(image_size)
1157+
mask_img = np.zeros((1,) + tuple(image_size), dtype=np.int16)
1158+
for b in range(extended_boxes_np.shape[0]):
1159+
slicing = (0,) + tuple(
1160+
slice(extended_boxes_np[b, d], extended_boxes_np[b, d + spatial_dims]) for d in range(spatial_dims)
1161+
)
1162+
mask_img[slicing] = 1
11611163
fg_indices_, bg_indices_ = map_binary_to_indices(mask_img, thresh_image, self.image_threshold)
11621164
else:
11631165
fg_indices_ = fg_indices

0 commit comments

Comments
 (0)