Fix axes order in tiff_dataset.SegmentationDataset samples

InnopolisUni · Jan 11, 2024 · 6cd28a7 · 6cd28a7
1 parent 088167f
commit 6cd28a7
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/innofw/core/datasets/semantic_segmentation/tiff_dataset.py b/innofw/core/datasets/semantic_segmentation/tiff_dataset.py
@@ -109,11 +109,11 @@ def __getitem__(self, index) -> dict:
                 # out = self.transform(image=image, mask=mask)
                 # image, mask = out["image"], out["mask"]
                 image, mask = self.transform(image, mask)
-
-        try:
-            image = np.moveaxis(image, 2, 0)  # todo: refactor
-        except:
-            pass
+        else:
+            try:
+                image = np.moveaxis(image, 2, 0) 
+            except:
+                pass
         try:
             image = image.astype(np.float32)  # todo: refactor
         except: