Allow GeoDataset to list files in VSI path(s)

torchgeo/datasets/geo.py

@@ -10,7 +10,7 @@
 import re
 import sys
 from collections.abc import Sequence
-from typing import Any, Callable, Optional, cast
+from typing import Any, Callable, Optional, Union, cast
 
 import fiona
 import fiona.transform
@@ -29,7 +29,13 @@
 from torchvision.datasets import ImageFolder
 from torchvision.datasets.folder import default_loader as pil_loader
 
-from .utils import BoundingBox, concat_samples, disambiguate_timestamp, merge_samples
+from .utils import (
+    BoundingBox,
+    concat_samples,
+    disambiguate_timestamp,
+    list_directory_recursive,
+    merge_samples,
+)
 
 
 class GeoDataset(Dataset[dict[str, Any]], abc.ABC):
@@ -329,7 +335,7 @@ def dtype(self) -> torch.dtype:
 
     def __init__(
         self,
-        root: str = "data",
+        root: Union[str, list[str]] = "data",
         crs: Optional[CRS] = None,
         res: Optional[float] = None,
         bands: Optional[Sequence[str]] = None,
@@ -339,7 +345,8 @@ def __init__(
         """Initialize a new Dataset instance.
 
         Args:
-            root: root directory where dataset can be found
+            root: root directory or list of absolute filepaths where
+                dataset can be found
             crs: :term:`coordinate reference system (CRS)` to warp to
                 (defaults to the CRS of the first file found)
             res: resolution of the dataset in units of CRS
@@ -358,11 +365,22 @@ def __init__(
         self.bands = bands or self.all_bands
         self.cache = cache
 
+        if isinstance(root, str):
+            root = [root]
+
+        filespaths: list[str] = []
+        for dir_or_file in root:
+            if os.path.exists(dir_or_file) and os.path.isfile(dir_or_file):
+                filespaths.append(dir_or_file)
+            else:
+                filespaths.extend(
+                    list_directory_recursive(dir_or_file, self.filename_glob)
+                )
+
         # Populate the dataset index
         i = 0
-        pathname = os.path.join(root, "**", self.filename_glob)
         filename_regex = re.compile(self.filename_regex, re.VERBOSE)
-        for filepath in glob.iglob(pathname, recursive=True):
+        for filepath in filespaths:
             match = re.match(filename_regex, os.path.basename(filepath))
             if match is not None:
                 try:

torchgeo/datasets/utils.py

@@ -9,6 +9,8 @@
 import bz2
 import collections
 import contextlib
+import fnmatch
+import glob
 import gzip
 import lzma
 import os
@@ -19,9 +21,11 @@
 from datetime import datetime, timedelta
 from typing import Any, 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
@@ -43,6 +47,7 @@
     "draw_semantic_segmentation_masks",
     "rgb_to_mask",
     "percentile_normalization",
+    "list_directory_recursive",
 )
 
 
@@ -737,3 +742,48 @@ def percentile_normalization(
         (img - lower_percentile) / (upper_percentile - lower_percentile + 1e-5), 0, 1
     )
     return img_normalized
+
+
+def _path_is_vsi(path: str) -> bool:
+    from rasterio._path import SCHEMES
+
+    prefix = path.split("://")[0]
+    schemes = prefix.split("+")
+    is_apache_vfs_scheme = set(schemes).issubset(set(SCHEMES))
+    is_gdal_vsi = path.startswith("/vsi")
+    return is_gdal_vsi or is_apache_vfs_scheme
+
+
+def _listdir_vsi_recursive(root: str) -> list[str]:
+    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 as e:
+            if "is not a directory" in str(e):
+                files.append(dir)
+            else:
+                raise e
+    return files
+
+
+def list_directory_recursive(root: str, filename_glob: str) -> list[str]:
+    """Lists files in directory recursively.
+
+    Also supports gdal virtual file systems (vsi).
+
+    Args:
+        root: directory to list. For vsi these can start with
+            e.g. /vsiaz or az:// for azure blob storage
+        filename_glob: filename pattern to filter filenames
+    """
+    if not _path_is_vsi(root):
+        filepaths = _listdir_vsi_recursive(root)
+        filepaths = fnmatch.filter(filepaths, filename_glob)
+    else:
+        pathname = os.path.join(root, "**", filename_glob)
+        filepaths = list(glob.iglob(pathname, recursive=True))
+    return filepaths