From fbd92688c960ba0543eecb86a741ec172e821267 Mon Sep 17 00:00:00 2001
From: aditiiyer <iyera@mskcc.org>
Date: Wed, 11 Sep 2024 16:20:41 -0400
Subject: [PATCH] Corrections to box filt radius, padding

---
 cerr/utils/image_proc.py | 11 ++++-------
 cerr/utils/mask.py       | 31 ++++++++++++++++++++++---------
 2 files changed, 26 insertions(+), 16 deletions(-)

diff --git a/cerr/utils/image_proc.py b/cerr/utils/image_proc.py
index 21aaa3f..27df923 100644
--- a/cerr/utils/image_proc.py
+++ b/cerr/utils/image_proc.py
@@ -128,7 +128,7 @@ def resizeScanAndMask(scan3M, mask4M, gridS, outputImgSizeV, method, \
 
         xOutV = np.arange(xV[0] - (xPad/2)*voxSizeV[0],
                           xV[-1]+(xPad/2)*voxSizeV[0], voxSizeV[0])
-        yOutV = np.arange(yV[0] - (yPad/2)*voxSizeV[1], \
+        yOutV = np.arange(yV[0] - (yPad/2)*voxSizeV[1],
                           yV[-1]+(yPad/2)*voxSizeV[1], voxSizeV[1])
         zOutV = zV
 
@@ -289,12 +289,9 @@ def resizeScanAndMask(scan3M, mask4M, gridS, outputImgSizeV, method, \
                 maskOut4M[rMin:rMax+1, cMin:cMax+1, slcNum, :] = mask4M[outRmin:outRmax+1,\
                                                                  outCmin:outCmax+1, slcNum, :]
             # Return co-ordinates of cropped region
-            xOutM[:, slcNum] = np.arange(xV[outCmin, slcNum],
-                                         xV[outCmax, slcNum] + voxSizeV[0], voxSizeV[0] + EPS)
-            yOutM[:, slcNum] = np.arange(yV[outRmin, slcNum],
-                                         yV[outRmax, slcNum] + voxSizeV[1], voxSizeV[1] - EPS)
-
-        gridOutS = (xOutM, yOutM, zV)                           
+            xOutM[:, slcNum] = np.arange(xV[outCmin], xV[outCmax] + voxSizeV[0], voxSizeV[0] + EPS)
+            yOutM[:, slcNum] = np.arange(yV[outRmin], yV[outRmax] + voxSizeV[1], voxSizeV[1] - EPS)
+            gridOutS = (xOutM, yOutM, zV)
 
     elif methodLower in ['bilinear','bicubic','nearest']:
         #3D resizing. TBD: 2D
diff --git a/cerr/utils/mask.py b/cerr/utils/mask.py
index 6830d52..e451394 100644
--- a/cerr/utils/mask.py
+++ b/cerr/utils/mask.py
@@ -191,29 +191,42 @@ def blurring(binaryMask, sigmaVox, filtType='gaussian'):
 
     Args:
         binaryMask (numpy.array): Binary mask to blur.
-        sigmaVox (float): Sigma for Gaussian in units of voxels.
+        sigmaVox (float): Sigma for Gaussian OR filter radius for box filter (in units of voxels).
         filtType (string): [optional, default:'gaussian'] 'gaussian' or 'box' smoothing filter.
 
     Returns:
         numpy.ndarray(dtype=bool): Blurred mask using Gaussian blur with input sigma.
     """
 
+    origSize = binaryMask.shape
+    binaryMaskImg = sitk.GetImageFromArray(binaryMask.astype(int))
+    padded = False
+
     if filtType == 'gaussian':
         filter = sitk.SmoothingRecursiveGaussianImageFilter()
         filter.SetSigma(sigmaVox)
     elif filtType == 'box':
+        padded = True
+        padMaskImg = sitk.MirrorPad(binaryMaskImg, [sigmaVox] * 3, [sigmaVox] * 3)
+        padMaskArr = sitk.GetArrayFromImage(padMaskImg)
         filter = sitk.BoxMeanImageFilter()
         filter.SetRadius(sigmaVox)
 
-    dim = binaryMask.shape
-    blurredMask3M = np.empty_like(binaryMask, dtype=float)
-    for slc in range(dim[2]):
-        if not np.any(binaryMask[:,:,slc]):
-            blurredMask3M[:,:,slc] = binaryMask[:,:,slc]
+    fullSize = padMaskArr.shape
+    blurredMask3M = np.empty_like(padMaskArr, dtype=float)
+    for slc in range(fullSize[2]):
+        if not np.any(binaryMask[:, :, slc]):
+            blurredMask3M[:, :, slc] = padMaskArr[:, :, slc]
             continue
-        img = sitk.GetImageFromArray(binaryMask[:,:,slc].astype(float))
-        blurImage = filter.Execute(img)
-        blurredMask3M[:,:,slc] = sitk.GetArrayFromImage(blurImage)
+        img = sitk.GetImageFromArray(padMaskArr[:, :, slc].astype(float))
+        blurImageFull = filter.Execute(img)
+        if padded:
+            startIdx = [(f - o) // 2 for f, o in zip(fullSize, origSize)]
+            blurImage = sitk.Crop(blurImageFull, startIdx)
+        else:
+            blurImage = blurImageFull
+
+        blurredMask3M[:, :, slc] = sitk.GetArrayFromImage(blurImage)
     return blurredMask3M