ETCI2021: fix file list when 'vv' in directory name (#2532)

microsoft · Jan 24, 2025 · 171b498 · 171b498
1 parent ada5122
commit 171b498
Showing 1 changed file with 7 additions and 9 deletions.
diff --git a/torchgeo/datasets/etci2021.py b/torchgeo/datasets/etci2021.py
@@ -170,17 +170,15 @@ def _load_files(self, root: Path, split: str) -> list[dict[str, str]]:
         folders = [os.path.join(folder, 'tiles') for folder in folders]
         for folder in folders:
             vvs = sorted(glob.glob(os.path.join(folder, 'vv', '*.png')))
-            vhs = [vv.replace('vv', 'vh') for vv in vvs]
-            water_masks = [
-                vv.replace('_vv.png', '.png').replace('vv', 'water_body_label')
-                for vv in vvs
-            ]
+            vhs = sorted(glob.glob(os.path.join(folder, 'vh', '*.png')))
+            water_masks = sorted(
+                glob.glob(os.path.join(folder, 'water_body_label', '*.png'))
+            )
 
             if split != 'test':
-                flood_masks = [
-                    vv.replace('_vv.png', '.png').replace('vv', 'flood_label')
-                    for vv in vvs
-                ]
+                flood_masks = sorted(
+                    glob.glob(os.path.join(folder, 'flood_label', '*.png'))
+                )
 
                 for vv, vh, flood_mask, water_mask in zip(
                     vvs, vhs, flood_masks, water_masks