resolve valid_pos bug

Author: donglaiw

connectomics/config/defaults.py

@@ -278,6 +278,7 @@
 _C.DATASET.REJECT_SAMPLING.SIZE_THRES = -1
 _C.DATASET.REJECT_SAMPLING.DIVERSITY = -1
 _C.DATASET.REJECT_SAMPLING.P = 0.95
+_C.DATASET.REJECT_SAMPLING.NUM_TRIAL = 50
 
 # Normalize model inputs (the images are assumed to be gray-scale).
 _C.DATASET.MEAN = 0.5

connectomics/data/dataset/dataset_volume.py

@@ -1,6 +1,7 @@
 from typing import Optional, List
 import numpy as np
 import random
+import warnings
 
 import torch
 import torch.utils.data
@@ -68,6 +69,7 @@ def __init__(self,
                  do_relabel: bool = True,
                  # rejection sampling
                  reject_size_thres: int = 0,
+                 reject_num_trial: int = 50,
                  reject_diversity: int = 0,
                  reject_p: float = 0.95,
                  # normalization
@@ -98,6 +100,7 @@ def __init__(self,
         # rejection samping
         self.reject_size_thres = reject_size_thres
         self.reject_diversity = reject_diversity
+        self.reject_num_trial = reject_num_trial
         self.reject_p = reject_p
 
         # normalization
@@ -113,6 +116,17 @@ def __init__(self,
             assert len(set(x[0] for x in volume_size)) == 1, "All volumes should have the same number of channels"
         self.volume_size = [x[-3:] for x in volume_size]
 
+        volume_selection = [(sample_label_size <= x).all() for x in self.volume_size]
+        if not all(volume_selection):
+            print('remove volumes whose sizes are smaller than the model input', volume_selection)
+            self.volume = [x for i,x in enumerate(self.volume) if volume_selection[i]]
+            volume_size = [np.array(x.shape) for x in self.volume]
+            self.volume_size = [x[-3:] for x in volume_size]
+            if self.label is not None:
+                self.label = [x for i,x in enumerate(self.label) if volume_selection[i]]
+            if valid_mask is not None:
+                valid_mask = [x for i,x in enumerate(valid_mask) if volume_selection[i]]
+
         self.sample_volume_size = np.array(
             sample_volume_size).astype(int)  # model input size
         if self.label is not None:
@@ -122,7 +136,7 @@ def __init__(self,
             if self.augmentor is not None:
                 assert np.array_equal(
                     self.augmentor.sample_size, self.sample_label_size)
-        self._assert_valid_shape()
+        #self._assert_valid_shape()
 
         # compute number of samples for each dataset (multi-volume input)
         self.sample_stride = np.array(sample_stride).astype(int)
@@ -138,15 +152,19 @@ def __init__(self,
         self.valid_mask = valid_mask
         self.valid_ratio = valid_ratio
         # precompute valid region
+        # can be memory intensive
+        self.valid_pos = [None] * len(self.valid_mask) if self.valid_mask is not None else [None] * len(self.volume) 
+        """
         if self.valid_mask is not None:
-            self.valid_pos = [None] * len(self.valid_mask) 
             for i, x in enumerate(self.valid_mask):
                 if x is not None:
                     self.valid_pos[i] = get_valid_pos(x, sample_volume_size, valid_ratio)
                     self.sample_num[i] = self.valid_pos[i].shape[0]
+                    print(i, self.sample_num[i])
             self.sample_num_a = np.sum(self.sample_num)
             self.sample_num_c = np.cumsum([0] + list(self.sample_num))
-
+        """
+        
         if self.mode in ['val', 'test']:  # for validation and test
             self.sample_size_test = [
                 np.array([np.prod(x[1:3]), x[2]]) for x in self.sample_size]
@@ -240,17 +258,17 @@ def _get_pos_train(self, vol_size):
         # np.random: same seed
         pos = [0, 0, 0, 0]
         # pick a dataset
-        did = self._index_to_dataset(random.randint(0, self.sample_num_a))
+        did = self._index_to_dataset(random.randint(0, self.sample_num_a - 1))
         pos[0] = did
         # pick a position
         # all regions are valid
         if self.valid_pos[did] is None:
             tmp_size = count_volume(
                 self.volume_size[did], vol_size, self.sample_stride)
-            tmp_pos = [random.randint(0, tmp_size[x]) * self.sample_stride[x]
+            tmp_pos = [random.randint(0, tmp_size[x] - 1) * self.sample_stride[x]
                        for x in range(len(tmp_size))]
         else:
-            tmp_pos = self.valid_pos[did][random.randint(0, self.valid_pos[did].shape[0])]
+            tmp_pos = self.valid_pos[did][random.randint(0, self.valid_pos[did].shape[0]) - 1]
 
         pos[1:] = tmp_pos
         return pos
@@ -282,16 +300,21 @@ def _rejection_sampling(self, vol_size):
                 out_valid = augmented['valid_mask']
 
             if self._is_valid(out_valid) and self._is_fg(out_label):
+                #print('yes', sample_count)
                 return pos, out_volume, out_label, out_valid
 
             sample_count += 1
-            if sample_count > 100:
+            if sample_count > self.reject_num_trial:
                 err_msg = (
                     "Can not find any valid subvolume after sampling the "
-                    "dataset for more than 100 times. Please adjust the "
+                    f"dataset for more than {self.reject_num_trial} times. Please adjust the "
                     "valid mask or rejection sampling configurations."
                 )
-                raise RuntimeError(err_msg)
+                #raise RuntimeError(err_msg)
+                # return anyway with a useless sample
+                warnings.warn(err_msg)
+                #print('no..')
+                return pos, out_volume, out_label, out_valid
 
     def _random_sampling(self, vol_size):
         """Randomly sample a subvolume from all the volumes.

connectomics/data/utils/data_crop.py

@@ -1,5 +1,7 @@
 import numpy as np
 from scipy.ndimage import convolve
+import torch
+from torch.nn.functional import conv2d, conv3d
 
 ####################################################################
 ## Process image stacks.
@@ -18,27 +20,52 @@ def crop_volume(data, sz, st=(0, 0, 0)):
     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)
+def get_valid_pos_torch(mask, vol_sz, valid_ratio):
+    # torch version
+    # bug: out of memory
     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 
+        mask_sum = conv3d(torch.from_numpy(mask[None,None].astype(int)), torch.ones(tuple(vol_sz))[None,None], padding='valid')[0,0].numpy()>= 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
+        mask_sum = conv2d(torch.from_numpy(mask[None,None].astype(int)), torch.ones(tuple(vol_sz))[None,None], padding='valid')[0,0].numpy()>= 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
+
+def get_valid_pos(mask, vol_sz, valid_ratio):
+    # scipy version
+    valid_thres = valid_ratio * np.prod(vol_sz)
+    data_sz = mask.shape
+    mask_sum = convolve(mask.astype(int), np.ones(vol_sz), mode='constant', cval=0)
+    pad_sz_pre = (np.array(vol_sz) - 1) // 2
+    pad_sz_post = data_sz - (vol_sz - pad_sz_pre - 1) 
+    valid_pos = np.zeros([0,3])
+    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 
+        if mask_sum.max() > 0:
+            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.transpose([1,0,2])[mask_sum], \
+                                  yy.transpose([1,0,2])[mask_sum], \
+                                  xx.transpose([1,0,2])[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
+        if mask_sum.max() > 0:
+            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

docs/environment_docs.yml

@@ -100,7 +100,7 @@ dependencies:
       - connectomics==0.0.1.dev4
       - contourpy==1.1.1
       - cycler==0.12.1
-      - cython==0.29.21
+      - cython==0.29.22
       - debugpy==1.8.1
       - decorator==5.1.1
       - defusedxml==0.7.1

setup.py

@@ -10,12 +10,11 @@
     'scikit-image>=0.17.2',
     'opencv-python>=4.3.0',
     'matplotlib>=3.3.0',
-    'Cython==0.29.21',
+    'Cython>=0.29.22',
     'yacs>=0.1.8',
     'h5py>=2.10.0',
     'gputil>=1.4.0',
     'imageio>=2.9.0',
-    'tensorflow>=2.2.0',
     'tensorboard>=2.2.2',
     'einops>=0.3.0',
     'tqdm>=4.58.0',