SSL4EO-L: add download support

tests/data/ssl4eo/l/data.py

@@ -14,13 +14,14 @@
 from rasterio.crs import CRS
 
 SIZE = 36
+CHUNK_SIZE = 2**12
 
 np.random.seed(0)
 
 FILENAME_HIERARCHY = Union[dict[str, "FILENAME_HIERARCHY"], list[str]]
 
 filenames: FILENAME_HIERARCHY = {
-    "tm_toa": {
+    "ssl4eo_l_tm_toa": {
         "0000002": {
             "LT05_172034_20010526": ["all_bands.tif"],
             "LT05_172034_20020310": ["all_bands.tif"],
@@ -34,7 +35,7 @@
             "LT5_223084_20020923": ["all_bands.tif"],
         },
     },
-    "etm_sr": {
+    "ssl4eo_l_etm_toa": {
         "0000002": {
             "LE07_172034_20010526": ["all_bands.tif"],
             "LE07_172034_20020310": ["all_bands.tif"],
@@ -48,7 +49,7 @@
             "LE07_223084_20020923": ["all_bands.tif"],
         },
     },
-    "etm_toa": {
+    "ssl4eo_l_etm_sr": {
         "0000002": {
             "LE07_172034_20010526": ["all_bands.tif"],
             "LE07_172034_20020310": ["all_bands.tif"],
@@ -62,7 +63,7 @@
             "LE07_223084_20020923": ["all_bands.tif"],
         },
     },
-    "oli_tirs_toa": {
+    "ssl4eo_l_oli_tirs_toa": {
         "0000002": {
             "LC08_172034_20210306": ["all_bands.tif"],
             "LC08_172034_20210829": ["all_bands.tif"],
@@ -76,7 +77,7 @@
             "LC08_223084_20221211": ["all_bands.tif"],
         },
     },
-    "oli_sr": {
+    "ssl4eo_l_oli_sr": {
         "0000002": {
             "LC08_172034_20210306": ["all_bands.tif"],
             "LC08_172034_20210829": ["all_bands.tif"],
@@ -92,7 +93,13 @@
     },
 }
 
-num_bands = {"tm_toa": 7, "etm_sr": 6, "etm_toa": 9, "oli_tirs_toa": 11, "oli_sr": 7}
+num_bands = {
+    "ssl4eo_l_tm_toa": 7,
+    "ssl4eo_l_etm_toa": 9,
+    "ssl4eo_l_etm_sr": 6,
+    "ssl4eo_l_oli_tirs_toa": 11,
+    "ssl4eo_l_oli_sr": 7,
+}
 
 
 def create_file(path: str) -> None:
@@ -141,10 +148,27 @@ def create_directory(directory: str, hierarchy: FILENAME_HIERARCHY) -> None:
 
     directories = filenames.keys()
     for directory in directories:
-        # Create tarballs
+        # Create tarball
         shutil.make_archive(directory, "gztar", ".", directory)
 
+        # Split tarball
+        path = f"{directory}.tar.gz"
+        paths = []
+        with open(path, "rb") as f:
+            chunk = f.read(CHUNK_SIZE)
+            suffix = "a"
+            while chunk:
+                split = f"{path}a{suffix}"
+                with open(split, "wb") as g:
+                    g.write(chunk)
+                chunk = f.read(CHUNK_SIZE)
+                suffix = chr(ord(suffix) + 1)
+                paths.append(split)
+
+        os.remove(path)
+
         # Compute checksums
-        with open(f"{directory}.tar.gz", "rb") as f:
-            md5 = hashlib.md5(f.read()).hexdigest()
-            print(directory, md5)
+        for path in paths:
+            with open(path, "rb") as f:
+                md5 = hashlib.md5(f.read()).hexdigest()
+                print(path, md5)

tests/datasets/test_ssl4eo.py

@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 
+import glob
 import os
 import shutil
 from pathlib import Path
@@ -13,16 +14,57 @@
 from pytest import MonkeyPatch
 from torch.utils.data import ConcatDataset
 
+import torchgeo
 from torchgeo.datasets import SSL4EOL, SSL4EOS12
 
 
+def download_url(url: str, root: str, *args: str, **kwargs: str) -> None:
+    shutil.copy(url, root)
+
+
 class TestSSL4EOL:
     @pytest.fixture(params=zip(SSL4EOL.metadata.keys(), [1, 1, 2, 2, 4]))
-    def dataset(self, request: SubRequest) -> SSL4EOL:
+    def dataset(
+        self, monkeypatch: MonkeyPatch, tmp_path: Path, request: SubRequest
+    ) -> SSL4EOL:
+        monkeypatch.setattr(torchgeo.datasets.ssl4eo, "download_url", download_url)
+
+        url = os.path.join("tests", "data", "ssl4eo", "l", "ssl4eo_l_{0}.tar.gz{1}")
+        monkeypatch.setattr(SSL4EOL, "url", url)
+
+        checksums = {
+            "tm_toa": {
+                "aa": "010b9d72b476e0e30741c17725f84e5c",
+                "ab": "39171bd7bca8a56a8cb339a0f88da9d3",
+                "ac": "3cfc407ce3f4f4d6e3c5fdb457bb87da",
+            },
+            "etm_toa": {
+                "aa": "87e47278f5a30acd3b696b6daaa4713b",
+                "ab": "59295e1816e08a5acd3a18ae56b6f32e",
+                "ac": "f3ff76eb6987501000228ce15684e09f",
+            },
+            "etm_sr": {
+                "aa": "fd61a4154eafaeb350dbb01a2551a818",
+                "ab": "0c3117bc7682ba9ffdc6871e6c364b36",
+                "ac": "93d3385e47de4578878ca5c4fa6a628d",
+            },
+            "oli_tirs_toa": {
+                "aa": "defb9e91a73b145b2dbe347649bded06",
+                "ab": "97f7edaa4e288fc14ec7581dccea766f",
+                "ac": "7472fad9929a0dc96ccf4dc6c804b92f",
+            },
+            "oli_sr": {
+                "aa": "8fd3aa6b581d024299f44457956faa05",
+                "ab": "7eb4d761ce1afd89cae9c6142ca17882",
+                "ac": "a3210da9fcc71e3a4efde71c30d78c59",
+            },
+        }
+        monkeypatch.setattr(SSL4EOL, "checksums", checksums)
+
+        root = str(tmp_path)
         split, seasons = request.param
-        root = os.path.join("tests", "data", "ssl4eo", "l", split)
         transforms = nn.Identity()
-        return SSL4EOL(root, split, seasons, transforms)
+        return SSL4EOL(root, split, seasons, transforms, download=True, checksum=True)
 
     def test_getitem(self, dataset: SSL4EOL) -> None:
         x = dataset[0]
@@ -41,6 +83,20 @@ def test_add(self, dataset: SSL4EOL) -> None:
         assert isinstance(ds, ConcatDataset)
         assert len(ds) == 2 * 2
 
+    def test_already_extracted(self, dataset: SSL4EOL) -> None:
+        SSL4EOL(dataset.root, dataset.split, dataset.seasons)
+
+    def test_already_downloaded(self, dataset: SSL4EOL, tmp_path: Path) -> None:
+        pathname = os.path.join("tests", "data", "ssl4eo", "l", "*.tar.gz*")
+        root = str(tmp_path)
+        for tarfile in glob.iglob(pathname):
+            shutil.copy(tarfile, root)
+        SSL4EOL(root)
+
+    def test_not_downloaded(self, tmp_path: Path) -> None:
+        with pytest.raises(RuntimeError, match="Dataset not found"):
+            SSL4EOL(str(tmp_path))
+
     def test_invalid_split(self) -> None:
         with pytest.raises(AssertionError):
             SSL4EOL(split="foo")