add get_valid_pos

Author: donglaiw

connectomics/data/utils/data_crop.py

@@ -1,4 +1,5 @@
 import numpy as np
+from scipy.ndimage import convolve
 
 ####################################################################
 ## Process image stacks.
@@ -16,3 +17,28 @@ def crop_volume(data, sz, st=(0, 0, 0)):
         return data[st[0]:st[0]+sz[0], st[1]:st[1]+sz[1], st[2]:st[2]+sz[2]]
     else: # crop spatial dimensions
         return data[:, st[0]:st[0]+sz[0], st[1]:st[1]+sz[1], st[2]:st[2]+sz[2]]
+
+def get_valid_pos(mask, vol_sz, valid_ratio):
+    mask_sum = convolve(mask, np.ones(vol_sz), mode='constant', cval=0)
+    valid_thres = valid_ratio * np.prod(vol_sz)
+    data_sz = mask.shape
+    pad_sz_pre = (np.array(vol_sz) - 1) // 2
+    pad_sz_post = data_sz - (vol_sz - pad_sz_pre - 1) 
+    if len(vol_sz) == 3:
+        mask_sum = mask_sum[pad_sz_pre[0]:pad_sz_post[0], \
+                            pad_sz_pre[1]:pad_sz_post[1], \
+                            pad_sz_pre[2]:pad_sz_post[2]] >= valid_thres 
+        zz, yy, xx = np.meshgrid(np.arange(mask_sum.shape[0]), \
+                                 np.arange(mask_sum.shape[1]), \
+                                 np.arange(mask_sum.shape[2]))
+        valid_pos = np.stack([zz.T[mask_sum], \
+                              yy.T[mask_sum], \
+                              xx.T[mask_sum]], axis=1)
+    else:
+        mask_sum = mask_sum[pad_sz_pre[0]:pad_sz_post[0], \
+                            pad_sz_pre[1]:pad_sz_post[1]] >= valid_thres
+        yy, xx = np.meshgrid(np.arange(mask_sum.shape[0]), \
+                                 np.arange(mask_sum.shape[1]))
+        valid_pos = np.stack([yy.T[mask_sum], \
+                              xx.T[mask_sum]], axis=1)
+    return valid_pos