Skip redundant full-image mask on StdShiftIntensity nonzero=False path

On the default nonzero=False path the boolean mask is all-True, so
img[slices] just gathers and scatters the whole image. Shift the image
directly instead, avoiding the mask allocation, .any(), gather and
scatter. Output is unchanged.

Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>

Author: aymuos15

monai/transforms/intensity/array.py

@@ -354,16 +354,15 @@ def __init__(
         self.dtype = dtype
 
     def _stdshift(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
-        ones: Callable
         std: Callable
         if isinstance(img, torch.Tensor):
-            ones = torch.ones
             std = partial(torch.std, unbiased=False)
         else:
-            ones = np.ones
             std = np.std
 
-        slices = (img != 0) if self.nonzero else ones(img.shape, dtype=bool)
+        if not self.nonzero:
+            return img + self.factor * std(img)
+        slices = img != 0
         if slices.any():
             offset = self.factor * std(img[slices])
             img[slices] = img[slices] + offset