Fix rotation.

juglab · Oct 5, 2022 · 92d2c5f · 92d2c5f
1 parent 847d9bf
commit 92d2c5f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/cryocare/internals/CryoCAREDataModule.py b/cryocare/internals/CryoCAREDataModule.py
@@ -134,8 +134,8 @@ def augment(self, x, y):
             rot_k = np.random.randint(0, 4, x.shape[0])
 
             for i in range(x.shape[0]):
-                x[i] = np.rot90(x[i], k=rot_k[i], axes=self.rot_axes)
-                y[i] = np.rot90(y[i], k=rot_k[i], axes=self.rot_axes)
+                x[i,...,0] = np.rot90(x[i,...,0], k=rot_k[i], axes=self.rot_axes)
+                y[i,...,0] = np.rot90(y[i,...,0], k=rot_k[i], axes=self.rot_axes)
 
 
         if np.random.rand() > 0.5: