Update prediction.py

prediction.py

@@ -8,6 +8,8 @@ from tqdm import tqdm
 import warnings
 warnings.filterwarnings("ignore")
 from huggingface_hub import hf_hub_download
+from skimage.morphology import binary_erosion, binary_dilation
+from skimage import draw
 
 
 def predict_mask(image, dim_x, dim_y, dim_z, _resize=True, norm_=True, mode_='test', patch_size=(64,128,128,1), _step=64, _step_z=32, _patch_size_z=64):
@@ -144,4 +146,19 @@ def predict_mask(image, dim_x, dim_y, dim_z, _resize=True, norm_=True, mode_='te
     final_mask = final_mask*255.0
     final_mask = final_mask.astype('uint8')
 
-    return final_mask
+
+	#closing operation to fill small holes
+	mask = final_mask
+	mask[mask!=0] = 1
+	mask = mask.astype('uint8')
+	
+	ellipsoid = draw.ellipsoid(9,9,3, spacing=(1,1,1), levelset=False)
+	ellipsoid = ellipsoid.astype('uint8')
+	ellipsoid = ellipsoid[1:-1,1:-1,1:-1]
+	
+	#perform closing operation on the mask
+	dil = binary_dilation(mask, ellipsoid)
+	closed_mask = binary_erosion(dil, ellipsoid)
+	closed_mask = (closed_mask*255.0).astype('uint8')
+    
+    return closed_mask