Allow GeoDataset to list files in VSI path(s)

.gitignore

@@ -149,3 +149,6 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# Mac
+.DS_Store

pyproject.toml

@@ -39,7 +39,8 @@ dependencies = [
     # einops 0.3+ required for einops.repeat
     "einops>=0.3",
     # fiona 1.8.21+ required for Python 3.10 wheels
-    "fiona>=1.8.21",
+    # fiona 1.9+ required for fiona.listdir
+    "fiona>=1.9",
     # kornia 0.7.3+ required for instance segmentation support in AugmentationSequential
     "kornia>=0.7.3",
     # lightly 1.4.5+ required for LARS optimizer

requirements/min-reqs.old

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

tests/datasets/test_geo.py

@@ -1,11 +1,11 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
-
 import math
 import os
 import pickle
+import shutil
 import sys
-from collections.abc import Iterable
+from collections.abc import Generator, Iterable
 from pathlib import Path
 from typing import Any
 
@@ -30,6 +30,7 @@
     UnionDataset,
     VectorDataset,
 )
+from torchgeo.datasets.utils import list_directory_recursive
 
 
 class CustomGeoDataset(GeoDataset):
@@ -84,6 +85,19 @@ def __len__(self) -> int:
         return 2
 
 
+@pytest.fixture(scope='module')
+def temp_archive(request: SubRequest) -> Generator[tuple[str, str], None, None]:
+    # Runs before tests
+    dir_not_zipped = request.param
+    dir_zipped = f'{dir_not_zipped}.zip'
+    _ = shutil.make_archive(dir_not_zipped, 'zip', dir_not_zipped)
+    # make_archive returns absolute path, while input may be relative path.
+    # we opt to return the (relative) path as provided.
+    yield dir_not_zipped, dir_zipped
+    # Runs after tests
+    os.remove(dir_zipped)
+
+
 class TestGeoDataset:
     @pytest.fixture(scope='class')
     def dataset(self) -> GeoDataset:
@@ -178,29 +192,24 @@ def test_files_property_for_non_existing_file_or_dir(self, tmp_path: Path) -> No
         with pytest.warns(UserWarning, match='Path was ignored.'):
             assert len(CustomGeoDataset(paths=paths).files) == 0
 
-    def test_files_property_for_virtual_files(self) -> None:
-        # Tests only a subset of schemes and combinations.
-        paths = [
-            'file://directory/file.tif',
-            'zip://archive.zip!folder/file.tif',
-            'az://azure_bucket/prefix/file.tif',
-            '/vsiaz/azure_bucket/prefix/file.tif',
-            'zip+az://azure_bucket/prefix/archive.zip!folder_in_archive/file.tif',
-            '/vsizip//vsiaz/azure_bucket/prefix/archive.zip/folder_in_archive/file.tif',
-        ]
-        assert len(CustomGeoDataset(paths=paths).files) == len(paths)
-
-    def test_files_property_ordered(self) -> None:
+    def test_files_property_ordered(self, tmp_path: Path) -> None:
         """Ensure that the list of files is ordered."""
-        paths = ['file://file3.tif', 'file://file1.tif', 'file://file2.tif']
+
+        files = ['file3.tif', 'file1.tif', 'file2.tif']
+        paths = [tmp_path / fake_file for fake_file in files]
+        for fake_file in paths:
+            fake_file.touch()
         assert CustomGeoDataset(paths=paths).files == sorted(paths)
 
-    def test_files_property_deterministic(self) -> None:
+    def test_files_property_deterministic(self, tmp_path: Path) -> None:
         """Ensure that the list of files is consistent regardless of their original
         order.
         """
-        paths1 = ['file://file3.tif', 'file://file1.tif', 'file://file2.tif']
-        paths2 = ['file://file2.tif', 'file://file3.tif', 'file://file1.tif']
+        files = ['file3.tif', 'file1.tif', 'file2.tif']
+        paths1 = [tmp_path / fake_file for fake_file in files]
+        paths2 = paths1[::-1]  # reverse order
+        for fake_file in paths1:
+            fake_file.touch()
         assert (
             CustomGeoDataset(paths=paths1).files == CustomGeoDataset(paths=paths2).files
         )
@@ -368,6 +377,136 @@ def test_no_all_bands(self) -> None:
             CustomSentinelDataset(root, bands=bands, transforms=transforms, cache=cache)
 
 
+@pytest.mark.skipif(
+    sys.platform.startswith('win'),
+    reason='Have not been able to work this out on windows. '
+    'Try again when fiona version 1.10 arrives.',
+)
+class TestVirtualFilesystems:
+    @pytest.mark.parametrize(
+        'temp_archive', [os.path.join('tests', 'data', 'vector')], indirect=True
+    )
+    def test_zipped_file(self, temp_archive: tuple[str, str]) -> None:
+        _, dir_zipped = temp_archive
+        filename = 'vector_2024.geojson'
+
+        # Using '!' as separator between archive and file as per Apache Commons VFS
+        # https://rasterio.readthedocs.io/en/latest/topics/vsi.html
+        specific_file_zipped = f'{dir_zipped}!{filename}'
+
+        files_found = CustomGeoDataset(paths=f'zip://{specific_file_zipped}').files
+        assert len(files_found) == 1
+        file = str(files_found[0])
+        assert file.endswith(filename)
+
+    @pytest.mark.parametrize(
+        'temp_archive',
+        [
+            os.path.join(
+                'tests',
+                'data',
+                'sentinel2',
+                'S2A_MSIL2A_20220414T110751_N0400_R108_T26EMU_20220414T165533.SAFE',
+            )
+        ],
+        indirect=True,
+    )
+    def test_zipped_specific_file_dir(self, temp_archive: tuple[str, str]) -> None:
+        dir_not_zipped, dir_zipped = temp_archive
+
+        # Listing, instead of hardcoding filepath
+        filename_glob = Sentinel2.filename_glob.format('B02')
+        specific_file_not_zipped = list_directory_recursive(
+            dir_not_zipped, filename_glob
+        )[0]
+        filepath_within_dir = specific_file_not_zipped.replace(dir_not_zipped, '')
+
+        files_found = CustomGeoDataset(
+            paths=f'zip://{dir_zipped}!{filepath_within_dir}'
+        ).files
+        assert len(files_found) == 1
+        file = str(files_found[0])
+        assert file.endswith(filepath_within_dir)
+
+    @pytest.mark.parametrize(
+        'temp_archive',
+        [
+            os.path.join(
+                'tests',
+                'data',
+                'sentinel2',
+                'S2A_MSIL2A_20220414T110751_N0400_R108_T26EMU_20220414T165533.SAFE',
+            )
+        ],
+        indirect=True,
+    )
+    def test_zipped_path_non_existing_file(self, temp_archive: tuple[str, str]) -> None:
+        dir_not_zipped, dir_zipped = temp_archive
+        file_zipped_not_existing = os.path.join(dir_zipped, 'non_existing_file.tif')
+        with pytest.warns(UserWarning, match='Path was ignored.'):
+            assert (
+                len(CustomGeoDataset(paths=f'zip://{file_zipped_not_existing}').files)
+                == 0
+            )
+
+    @pytest.mark.parametrize(
+        'temp_archive',
+        [
+            os.path.join(
+                'tests',
+                'data',
+                'sentinel2',
+                'S2A_MSIL2A_20220414T110751_N0400_R108_T26EMU_20220414T165533.SAFE',
+            )
+        ],
+        indirect=True,
+    )
+    def test_zipped_directory(self, temp_archive: tuple[str, str]) -> None:
+        dir_not_zipped, dir_zipped = temp_archive
+        bands = Sentinel2.rgb_bands
+        transforms = nn.Identity()
+        cache = False
+
+        files_not_zipped = Sentinel2(
+            paths=dir_not_zipped, bands=bands, transforms=transforms, cache=cache
+        ).files
+
+        files_zipped = Sentinel2(
+            paths=f'zip://{dir_zipped}', bands=bands, transforms=transforms, cache=cache
+        ).files
+
+        basenames_not_zipped = [Path(path).stem for path in files_not_zipped]
+        basenames_zipped = [Path(path).stem for path in files_zipped]
+        assert basenames_zipped == basenames_not_zipped
+
+    def test_http_files(self) -> None:
+        url = 'https://github.com/microsoft/torchgeo/raw/main'
+        dir = os.path.join(
+            'tests',
+            'data',
+            'sentinel2',
+            'S2A_MSIL2A_20220414T110751_N0400_R108_T26EMU_20220414T165533.SAFE',
+            'GRANULE',
+            'L2A_T26EMU_A035569_20220414T110747',
+            'IMG_DATA',
+            'R10m',
+        )
+        files = [
+            'T26EMU_20220414T110751_B04_10m.jp2',
+            'T26EMU_20190414T110751_B03_10m.jp2',
+        ]
+        paths = [os.path.join(url, dir, file) for file in files]
+        bands = Sentinel2.rgb_bands
+        transforms = nn.Identity()
+        cache = False
+
+        files_found = Sentinel2(
+            paths=paths, bands=bands, transforms=transforms, cache=cache
+        ).files
+
+        assert len(files_found) == 1
+
+
 class TestVectorDataset:
     @pytest.fixture(scope='class')
     def dataset(self) -> CustomVectorDataset:

torchgeo/datasets/geo.py

@@ -6,7 +6,6 @@
 import abc
 import fnmatch
 import functools
-import glob
 import os
 import pathlib
 import re
@@ -40,8 +39,8 @@
     array_to_tensor,
     concat_samples,
     disambiguate_timestamp,
+    list_directory_recursive,
     merge_samples,
-    path_is_vsi,
 )
 
 
@@ -308,14 +307,16 @@ 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)) and fnmatch.fnmatch(
-                str(path), f'*{self.filename_glob}'
-            ):
-                files.add(path)
-            elif not hasattr(self, 'download'):
+            file_not_found = True
+            if os.path.isfile(path):
+                if fnmatch.fnmatch(str(path), os.path.join('*', self.filename_glob)):
+                    files.add(path)
+                    file_not_found = False
+            elif files_found := set(list_directory_recursive(path, self.filename_glob)):
+                files |= files_found
+                file_not_found = False
+
+            if file_not_found 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

@@ -8,6 +8,8 @@
 
 import collections
 import contextlib
+import fnmatch
+import glob
 import importlib
 import os
 import pathlib
@@ -19,9 +21,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,
@@ -631,6 +635,74 @@ def path_is_vsi(path: Path) -> bool:
     return '://' in str(path) or str(path).startswith('/vsi')
 
 
+def listdir_vsi_recursive(root: Path) -> list[str]:
+    """Lists all files in Virtual File Systems (VSI) recursively.
+
+    Args:
+        root: directory to list. These must contain the prefix for the VSI
+            e.g. '/vsiaz/' or 'az://' for azure blob storage
+                 '/vsizip/' or 'zip://' for zipped archives
+
+    Returns:
+        A list of all file paths matching filename_glob in the root VSI directory or its
+        subdirectories.
+
+    .. versionadded:: #0.6
+    """
+    dirs = [str(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 as e:
+            if 'is not a directory' in str(e):
+                # Assuming dir is a file as it is not a directory
+                # fiona.listdir can throw FionaValueError for only two reasons
+                files.append(dir)
+            else:
+                # fiona.listdir can throw FionaValueError for only two reasons
+                # 1. 'is not a directory'
+                # 2. 'does not exist'
+                raise
+    return files
+
+
+def list_directory_recursive(root: Path, filename_glob: str) -> list[str]:
+    """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
+
+    Returns:
+        A list of all file paths matching filename_glob in the root directory or its
+        subdirectories.
+
+    .. versionadded:: #0.6
+    """
+    files: list[str]
+    if path_is_vsi(root):
+        # Change type to match expected input to filter
+        all_files: list[str] = []
+        try:
+            all_files = listdir_vsi_recursive(root)
+        except FionaValueError:
+            # To match the behaviour of glob.iglob we silently return empty list
+            # for non-existing root.
+            pass
+        # Prefix glob with wildcard to ignore directories
+        files = fnmatch.filter(all_files, os.path.join('*', 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`.