remove negative part for saving

Author: donglaiw

connectomics/data/utils/data_segmentation.py

@@ -291,7 +291,7 @@ def seg2inst_edt(label, topt):
     _, mode, padding, quant, z_res, erosion = topt.split('-')
     resolution = (float(z_res), 1.0, 1.0)
     return edt_instance(label.copy(), mode, resolution=resolution,
-                        quantize=bool(int(quant)), padding=bool(int(padding)), erosion=erosion)
+                        quantize=bool(int(quant)), padding=bool(int(padding)), erosion=int(erosion))
 
 
 def seg_to_targets(

connectomics/data/utils/data_transform.py

@@ -5,7 +5,7 @@
 import scipy
 import numpy as np
 from scipy.ndimage import distance_transform_edt
-from skimage.morphology import remove_small_holes, skeletonize, binary_erosion, disk
+from skimage.morphology import remove_small_holes, skeletonize, binary_erosion, disk, ball
 from skimage.measure import label as label_cc  # avoid namespace conflict
 from skimage.filters import gaussian
 
@@ -61,7 +61,7 @@ def edt_instance(label: np.ndarray,
                  mode: str = '2d',
                  quantize: bool = True,
                  resolution: Tuple[float] = (1.0, 1.0, 1.0),
-                 padding: bool = False),
+                 padding: bool = False,
                  erosion: int = 0):
     assert mode in ['2d', '3d']
     if mode == '3d':
@@ -151,11 +151,14 @@ def distance_transform(label: np.ndarray,
 
     if not all_bg_sample:
         if erosion > 0:
-            erosion_disk = disk(erosion)
+            if label.ndim == 2:
+                footprint = disk(erosion)
+            elif label.ndim == 3:
+                footprint = ball(erosion)
         for idx in indices:
             temp2 = remove_small_holes(label == idx, 16, connectivity=1)
             if erosion > 0:
-                temp2 = binary_erosion(temp2, erosion_disk)
+                temp2 = binary_erosion(temp2, footprint)
 
             semantic += temp2.astype(np.uint8)
             boundary_edt = distance_transform_edt(temp2, resolution)

connectomics/engine/trainer.py

@@ -250,6 +250,8 @@ def test(self):
             if result[vol_id].ndim > weight[vol_id].ndim:
                 weight[vol_id] = np.expand_dims(weight[vol_id], axis=0)
             result[vol_id] /= weight[vol_id]  # in-place to save memory
+            # remoe negative part for quantization
+            result[vol_id][result[vol_id]<0] = 0
             result[vol_id] *= 255
             result[vol_id] = result[vol_id].astype(np.uint8)