Reduce number of for-loops during pixel manipulation

n2v/internals/N2V_DataWrapper.py

@@ -54,76 +54,75 @@ def __init__(self, X, Y, batch_size, perc_pix=0.198, shape=(64, 64),
 
         self.X_Batches = np.zeros((self.X.shape[0], *self.shape, self.n_chan), dtype=np.float32)
         self.Y_Batches = np.zeros((self.Y.shape[0], *self.shape, 2*self.n_chan), dtype=np.float32)
+
     def __len__(self):
         return int(np.ceil(len(self.X) / float(self.batch_size)))
 
     def on_epoch_end(self):
         self.perm = np.random.permutation(len(self.X))
+        self.X_Batches *= 0
+        self.Y_Batches *= 0
 
     def __getitem__(self, i):
         idx = slice(i * self.batch_size, (i + 1) * self.batch_size)
         idx = self.perm[idx]
-        self.patch_sampler(self.X, self.Y, self.X_Batches, self.Y_Batches, idx, self.range, self.shape)
+        self.patch_sampler(self.X, self.X_Batches, indices=idx, range=self.range, shape=self.shape)
 
-        for j in idx:
-            for c in range(self.n_chan):
+        for c in range(self.n_chan):
+            for j in idx:
                 coords = self.get_stratified_coords(self.rand_float, box_size=self.box_size,
-                                                    shape=np.array(self.X_Batches.shape)[1:-1])
-
-                y_val = []
-                x_val = []
-                for k in range(len(coords)):
-                    y_val.append(np.copy(self.Y_Batches[(j, *coords[k], ..., c)]))
-                    x_val.append(self.value_manipulation(self.X_Batches[j, ..., c][...,np.newaxis], coords[k], self.dims))
-
-                self.Y_Batches[j,...,c] *= 0
-                self.Y_Batches[j,...,self.n_chan+c] *= 0
+                                                    shape=self.shape)
 
-                for k in range(len(coords)):
-                    self.Y_Batches[(j, *coords[k], c)] = y_val[k]
-                    self.Y_Batches[(j, *coords[k], self.n_chan+c)] = 1
-                    self.X_Batches[(j, *coords[k], c)] = x_val[k]
+                indexing = (j,) + coords + (c,)
+                indexing_mask = (j,) + coords + (c + self.n_chan, )
+                y_val = self.X_Batches[indexing]
+                x_val = self.value_manipulation(self.X_Batches[j, ..., c], coords, self.dims)
 
+                self.Y_Batches[indexing] = y_val
+                self.Y_Batches[indexing_mask] = 1
+                self.X_Batches[indexing] = x_val
 
         return self.X_Batches[idx], self.Y_Batches[idx]
 
     @staticmethod
-    def __subpatch_sampling2D__(X, Y, X_Batches, Y_Batches, indices, range, shape):
+    def __subpatch_sampling2D__(X, X_Batches, indices, range, shape):
         for j in indices:
             y_start = np.random.randint(0, range[0] + 1)
             x_start = np.random.randint(0, range[1] + 1)
-            X_Batches[j] = X[j, y_start:y_start + shape[0], x_start:x_start + shape[1]]
-            Y_Batches[j] = Y[j, y_start:y_start + shape[0], x_start:x_start + shape[1]]
+            X_Batches[j] = np.copy(X[j, y_start:y_start + shape[0], x_start:x_start + shape[1]])
 
     @staticmethod
-    def __subpatch_sampling3D__(X, Y, X_Batches, Y_Batches, indices, range, shape):
+    def __subpatch_sampling3D__(X, X_Batches, indices, range, shape):
         for j in indices:
             z_start = np.random.randint(0, range[0] + 1)
             y_start = np.random.randint(0, range[1] + 1)
             x_start = np.random.randint(0, range[2] + 1)
-            X_Batches[j] = X[j, z_start:z_start + shape[0], y_start:y_start + shape[1], x_start:x_start + shape[2]]
-            Y_Batches[j] = Y[j, z_start:z_start + shape[0], y_start:y_start + shape[1], x_start:x_start + shape[2]]
+            X_Batches[j] = np.copy(X[j, z_start:z_start + shape[0], y_start:y_start + shape[1], x_start:x_start + shape[2]])
 
     @staticmethod
     def __get_stratified_coords2D__(coord_gen, box_size, shape):
-        coords = []
         box_count_y = int(np.ceil(shape[0] / box_size))
         box_count_x = int(np.ceil(shape[1] / box_size))
+        x_coords = []
+        y_coords = []
         for i in range(box_count_y):
             for j in range(box_count_x):
                 y, x = next(coord_gen)
                 y = int(i * box_size + y)
                 x = int(j * box_size + x)
                 if (y < shape[0] and x < shape[1]):
-                    coords.append((y, x))
-        return coords
+                    y_coords.append(y)
+                    x_coords.append(x)
+        return (y_coords, x_coords)
 
     @staticmethod
     def __get_stratified_coords3D__(coord_gen, box_size, shape):
-        coords = []
         box_count_z = int(np.ceil(shape[0] / box_size))
         box_count_y = int(np.ceil(shape[1] / box_size))
         box_count_x = int(np.ceil(shape[2] / box_size))
+        x_coords = []
+        y_coords = []
+        z_coords = []
         for i in range(box_count_z):
             for j in range(box_count_y):
                 for k in range(box_count_x):
@@ -132,8 +131,10 @@ def __get_stratified_coords3D__(coord_gen, box_size, shape):
                     y = int(j * box_size + y)
                     x = int(k * box_size + x)
                     if (z < shape[0] and y < shape[1] and x < shape[2]):
-                        coords.append((z, y, x))
-        return coords
+                        z_coords.append(z)
+                        y_coords.append(y)
+                        x_coords.append(x)
+        return (z_coords, y_coords, x_coords)
 
     @staticmethod
     def __rand_float_coords2D__(boxsize):

n2v/utils/n2v_utils.py

@@ -7,9 +7,6 @@ def get_subpatch(patch, coord, local_sub_patch_radius):
     start = np.maximum(0, np.array(coord) - local_sub_patch_radius)
     end = start + local_sub_patch_radius*2 + 1
 
-    start = np.append(start, 0)
-    end = np.append(end, patch.shape[-1])
-
     shift = np.minimum(0, patch.shape - end)
 
     start += shift
@@ -37,37 +34,52 @@ def normal_int(mean, sigma, w):
 
 
 def pm_normal_withoutCP(local_sub_patch_radius):
-    def normal_withoutCP(patch, coord, dims):
-        rand_coords = random_neighbor(patch.shape, coord)
-        return patch[tuple(rand_coords)]
+    def normal_withoutCP(patch, coords, dims):
+        vals = []
+        for coord in zip(*coords):
+            rand_coords = random_neighbor(patch.shape, coord)
+            vals.append(patch[tuple(rand_coords)])
+        return vals
     return normal_withoutCP
 
 
 def pm_uniform_withCP(local_sub_patch_radius):
-    def random_neighbor_withCP_uniform(patch, coord, dims):
-        sub_patch = get_subpatch(patch, coord,local_sub_patch_radius)
-        rand_coords = [np.random.randint(0, s) for s in sub_patch.shape[0:dims]]
-        return sub_patch[tuple(rand_coords)]
+    def random_neighbor_withCP_uniform(patch, coords, dims):
+        vals = []
+        for coord in zip(*coords):
+            sub_patch = get_subpatch(patch, coord,local_sub_patch_radius)
+            rand_coords = [np.random.randint(0, s) for s in sub_patch.shape[0:dims]]
+            vals.append(sub_patch[tuple(rand_coords)])
+        return vals
     return random_neighbor_withCP_uniform
 
 
 def pm_normal_additive(pixel_gauss_sigma):
-    def pixel_gauss(patch, coord, dims):
-        return np.random.normal(patch[tuple(coord)], pixel_gauss_sigma)
+    def pixel_gauss(patch, coords, dims):
+        vals = []
+        for coord in zip(*coords):
+            vals.append(np.random.normal(patch[tuple(coord)], pixel_gauss_sigma))
+        return vals
     return pixel_gauss
 
 
 def pm_normal_fitted(local_sub_patch_radius):
-    def local_gaussian(patch, coord, dims):
-        sub_patch = get_subpatch(patch, coord, local_sub_patch_radius)
-        axis = tuple(range(dims))
-        return np.random.normal(np.mean(sub_patch, axis=axis), np.std(sub_patch, axis=axis))
+    def local_gaussian(patch, coords, dims):
+        vals = []
+        for coord in zip(*coords):
+            sub_patch = get_subpatch(patch, coord, local_sub_patch_radius)
+            axis = tuple(range(dims))
+            vals.append(np.random.normal(np.mean(sub_patch, axis=axis), np.std(sub_patch, axis=axis)))
+        return vals
     return local_gaussian
 
 
 def pm_identity(local_sub_patch_radius):
-    def identity(patch, coord, dims):
-        return patch[tuple(coord)]
+    def identity(patch, coords, dims):
+        vals = []
+        for coord in zip(*coords):
+            vals.append(patch[coord])
+        return vals
     return identity
 
 

n2v/version.py

@@ -1 +1 @@
-__version__ = '0.1.9'
+__version__ = '0.1.10'

tests/test_Noise2VoidDataWrapper.py

@@ -15,7 +15,7 @@ def create_data(in_shape, out_shape):
     def _sample2D(in_shape, out_shape, seed):
         X, Y, X_Batches, Y_Batches, indices = create_data(in_shape, out_shape)
         np.random.seed(seed)
-        N2V_DataWrapper.__subpatch_sampling2D__(X, Y, X_Batches, Y_Batches, indices,
+        N2V_DataWrapper.__subpatch_sampling2D__(X, X_Batches, indices,
                                                       range=in_shape[1:3]-out_shape[1:3], shape=out_shape[1:3])
 
         assert ([*X_Batches.shape] == out_shape).all()
@@ -24,20 +24,18 @@ def _sample2D(in_shape, out_shape, seed):
         range_x = in_shape[2] - out_shape[2]
         for j in indices:
             assert np.sum(X_Batches[j]) != 0
-            assert np.sum(Y_Batches[j]) != 0
             y_start = np.random.randint(0, range_y + 1)
             x_start = np.random.randint(0, range_x + 1)
             assert np.sum(X_Batches[j] - X[j, y_start:y_start+out_shape[1], x_start:x_start+out_shape[2]]) == 0
 
         for j in range(in_shape[0]):
             if j not in indices:
                 assert np.sum(X_Batches[j]) == 0
-                assert np.sum(Y_Batches[j]) == 0
 
     def _sample3D(in_shape, out_shape, seed):
         X, Y, X_Batches, Y_Batches, indices = create_data(in_shape, out_shape)
         np.random.seed(seed)
-        N2V_DataWrapper.__subpatch_sampling3D__(X, Y, X_Batches, Y_Batches, indices,
+        N2V_DataWrapper.__subpatch_sampling3D__(X, X_Batches, indices,
                                                       range=in_shape[1:4]-out_shape[1:4], shape=out_shape[1:4])
 
         assert ([*X_Batches.shape] == out_shape).all()
@@ -47,7 +45,6 @@ def _sample3D(in_shape, out_shape, seed):
         range_x = in_shape[3] - out_shape[3]
         for j in indices:
             assert np.sum(X_Batches[j]) != 0
-            assert np.sum(Y_Batches[j]) != 0
             z_start = np.random.randint(0, range_z + 1)
             y_start = np.random.randint(0, range_y + 1)
             x_start = np.random.randint(0, range_x + 1)
@@ -56,7 +53,6 @@ def _sample3D(in_shape, out_shape, seed):
         for j in range(in_shape[0]):
             if j not in indices:
                 assert np.sum(X_Batches[j]) == 0
-                assert np.sum(Y_Batches[j]) == 0
 
     _sample2D(np.array([20, 64, 64, 2]), np.array([20, 32, 32, 2]), 1)
     _sample2D(np.array([10, 25, 25, 1]), np.array([10, 12, 12, 1]), 2)