Allow GeoDataset to list files in VSI path(s)

requirements/min-reqs.old

@@ -3,7 +3,7 @@ setuptools==61.0.0
 
 # install
 einops==0.3.0
-fiona==1.8.21
+fiona==1.9.6
 kornia==0.7.3
 lightly==1.4.5
 lightning[pytorch-extra]==2.0.0

tests/datasets/test_geo.py

@@ -1,8 +1,8 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
-
 import math
 import os
+import pathlib
 import pickle
 import sys
 from collections.abc import Iterable
@@ -205,6 +205,25 @@ def test_files_property_deterministic(self) -> None:
             CustomGeoDataset(paths=paths1).files == CustomGeoDataset(paths=paths2).files
         )
 
+    def test_zipped_sentinel2_dataset(self) -> None:
+        bands = ['B04', 'B03', 'B02']
+        transforms = nn.Identity()
+        cache = False
+
+        dir_not_zipped = 'tests/data/sentinel2/S2A_MSIL2A_20220414T110751_N0400_R108_T26EMU_20220414T165533.SAFE'
+        dir_zipped = f'zip://{dir_not_zipped}.zip'
+
+        files_not_zipped = Sentinel2(
+            paths=dir_not_zipped, bands=bands, transforms=transforms, cache=cache
+        ).files
+        files_zipped = Sentinel2(
+            paths=[dir_zipped], bands=bands, transforms=transforms, cache=cache
+        ).files
+
+        basenames_not_zipped = [pathlib.Path(path).stem for path in files_not_zipped]
+        basenames_zipped = [pathlib.Path(path).stem for path in files_zipped]
+        assert basenames_zipped == basenames_not_zipped
+
 
 class TestRasterDataset:
     @pytest.fixture(params=zip([['R', 'G', 'B'], None], [True, False]))

torchgeo/datasets/geo.py

@@ -5,7 +5,6 @@
 
 import abc
 import functools
-import glob
 import os
 import pathlib
 import re
@@ -39,8 +38,8 @@
     array_to_tensor,
     concat_samples,
     disambiguate_timestamp,
+    list_directory_recursive,
     merge_samples,
-    path_is_vsi,
 )
 
 
@@ -307,12 +306,11 @@ def files(self) -> list[Path]:
         # Using set to remove any duplicates if directories are overlapping
         files: set[Path] = set()
         for path in paths:
-            if os.path.isdir(path):
-                pathname = os.path.join(path, '**', self.filename_glob)
-                files |= set(glob.iglob(pathname, recursive=True))
-            elif os.path.isfile(path) or path_is_vsi(path):
+            if os.path.exists(path) and os.path.isfile(path):
                 files.add(path)
-            elif not hasattr(self, 'download'):
+            else:
+                files |= set(list_directory_recursive(path, self.filename_glob))
+            if not files and not hasattr(self, 'download'):
                 warnings.warn(
                     f"Could not find any relevant files for provided path '{path}'. "
                     f'Path was ignored.',

torchgeo/datasets/utils.py

@@ -9,6 +9,8 @@
 import bz2
 import collections
 import contextlib
+import fnmatch
+import glob
 import gzip
 import importlib
 import lzma
@@ -23,9 +25,11 @@
 from datetime import datetime, timedelta
 from typing import Any, TypeAlias, cast, overload
 
+import fiona
 import numpy as np
 import rasterio
 import torch
+from fiona.errors import FionaValueError
 from torch import Tensor
 from torchvision.datasets.utils import check_integrity, download_url
 from torchvision.utils import draw_segmentation_masks
@@ -772,6 +776,47 @@ def path_is_vsi(path: Path) -> bool:
     return '://' in str(path) or str(path).startswith('/vsi')
 
 
+def listdir_vsi_recursive(root: Path) -> list[Path]:
+    dirs = [root]
+    files = []
+    while dirs:
+        dir = dirs.pop()
+        try:
+            subdirs = fiona.listdir(dir)
+            dirs.extend([os.path.join(dir, subdir) for subdir in subdirs])
+        except FionaValueError:
+            # Assuming dir is a file as it is not a directory
+            # fiona.listdir can throw FionaValueError for only two reasons
+            # 1. 'is not a directory'
+            # 2. 'does not exist'
+            # We currently don't have tests for existing vsi, and will thus
+            # allow vsi to not exist.
+            files.append(dir)
+    return files
+
+
+def list_directory_recursive(root: Path, filename_glob: str) -> list[Path]:
+    """Lists files in directory recursively matching the given glob expression.
+
+    Also supports gdal virtual file systems (vsi).
+
+    Args:
+        root: directory to list. For vsi these will have prefix
+            e.g. /vsiaz or az:// for azure blob storage
+        filename_glob: filename pattern to filter filenames
+    """
+    files: list[Path]
+    if path_is_vsi(root):
+        # Change type to match expected input to filter
+        files_as_str: list[str] = [str(file) for file in listdir_vsi_recursive(root)]
+        # Prefix glob with wildcard to ignore directories
+        files = cast(list[Path], fnmatch.filter(files_as_str, '*' + filename_glob))
+    else:
+        pathname = os.path.join(root, '**', filename_glob)
+        files = list(glob.iglob(pathname, recursive=True))
+    return files
+
+
 def array_to_tensor(array: np.typing.NDArray[Any]) -> Tensor:
     """Converts a :class:`numpy.ndarray` to :class:`torch.Tensor`.