Remove additional dimension in single-patch extraction.

Closes #145.
juglab · Nov 15, 2023 · 3697518 · 3697518
1 parent 2ff781b
commit 3697518
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/n2v/internals/N2V_DataGenerator.py b/n2v/internals/N2V_DataGenerator.py
@@ -237,7 +237,7 @@ def __extract_patches__(self, data, num_patches=None, shape=(256, 256), n_dims=2
                 if len(patches) > 1:
                     return np.stack(patches)
                 else:
-                    return np.array(patches)[np.newaxis]
+                    return np.array(patches)
             elif n_dims == 3:
                 for i in range(num_patches):
                     z, y, x = np.random.randint(0, data.shape[1] - shape[0] + 1), np.random.randint(0,
@@ -250,7 +250,7 @@ def __extract_patches__(self, data, num_patches=None, shape=(256, 256), n_dims=2
                 if len(patches) > 1:
                     return np.stack(patches)
                 else:
-                    return np.array(patches)[np.newaxis]
+                    return np.array(patches)
             else:
                 print('Not implemented for more than 4 dimensional (ZYXC) data.')