Adding dataset from MOSAIKS paper

docs/api/datasets.rst

@@ -222,6 +222,11 @@ UC Merced
 
 .. autoclass:: UCMerced
 
+USAVars
+^^^^^^^
+
+.. autoclass:: USAVars
+
 Vaihingen
 ^^^^^^^^^
 

tests/data/usavars/data.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import glob
+import hashlib
+import os
+import shutil
+
+import numpy as np
+import pandas as pd
+import rasterio
+
+data_dir = "uar"
+labels = [
+    "elevation",
+    "population",
+    "treecover",
+    "income",
+    "nightlights",
+    "housing",
+    "roads",
+]
+SIZE = 3
+
+
+def create_file(path: str, dtype: str, num_channels: int) -> None:
+    profile = {}
+    profile["driver"] = "GTiff"
+    profile["dtype"] = dtype
+    profile["count"] = num_channels
+    profile["crs"] = "epsg:4326"
+    profile["transform"] = rasterio.transform.from_bounds(0, 0, 1, 1, 1, 1)
+    profile["height"] = SIZE
+    profile["width"] = SIZE
+    profile["compress"] = "lzw"
+    profile["predictor"] = 2
+
+    Z = np.random.randint(
+        np.iinfo(profile["dtype"]).max, size=(4, SIZE, SIZE), dtype=profile["dtype"]
+    )
+    src = rasterio.open(path, "w", **profile)
+    src.write(Z)
+
+
+# Remove old data
+filename = f"{data_dir}.zip"
+csvs = glob.glob("*.csv")
+
+for csv in csvs:
+    os.remove(csv)
+if os.path.exists(filename):
+    os.remove(filename)
+if os.path.exists(data_dir):
+    shutil.rmtree(data_dir)
+
+# Create tifs:
+os.makedirs(data_dir)
+create_file(os.path.join(data_dir, "tile_0,0.tif"), np.uint8, 4)
+create_file(os.path.join(data_dir, "tile_0,1.tif"), np.uint8, 4)
+
+# Create labels:
+columns = [["ID", "lon", "lat", lab] for lab in labels]
+fake_vals = [["0,0", 0.0, 0.0, 0.0], ["0,1", 0.1, 0.1, 1.0]]
+for lab, cols in zip(labels, columns):
+    df = pd.DataFrame(fake_vals, columns=cols)
+    df.to_csv(lab + ".csv")
+
+# Compress data
+shutil.make_archive(data_dir, "zip", ".", data_dir)
+
+# Compute checksums
+filename = f"{data_dir}.zip"
+with open(filename, "rb") as f:
+    md5 = hashlib.md5(f.read()).hexdigest()
+    print(repr(filename) + ": " + repr(md5) + ",")

tests/data/usavars/elevation.csv

@@ -0,0 +1,3 @@
+,ID,lon,lat,elevation
+0,"0,0",0.0,0.0,0.0
+1,"0,1",0.1,0.1,1.0

tests/data/usavars/housing.csv

@@ -0,0 +1,3 @@
+,ID,lon,lat,housing
+0,"0,0",0.0,0.0,0.0
+1,"0,1",0.1,0.1,1.0

tests/data/usavars/income.csv

@@ -0,0 +1,3 @@
+,ID,lon,lat,income
+0,"0,0",0.0,0.0,0.0
+1,"0,1",0.1,0.1,1.0

tests/data/usavars/nightlights.csv

@@ -0,0 +1,3 @@
+,ID,lon,lat,nightlights
+0,"0,0",0.0,0.0,0.0
+1,"0,1",0.1,0.1,1.0

tests/data/usavars/population.csv

@@ -0,0 +1,3 @@
+,ID,lon,lat,population
+0,"0,0",0.0,0.0,0.0
+1,"0,1",0.1,0.1,1.0

tests/data/usavars/roads.csv

@@ -0,0 +1,3 @@
+,ID,lon,lat,roads
+0,"0,0",0.0,0.0,0.0
+1,"0,1",0.1,0.1,1.0

tests/data/usavars/treecover.csv

@@ -0,0 +1,3 @@
+,ID,lon,lat,treecover
+0,"0,0",0.0,0.0,0.0
+1,"0,1",0.1,0.1,1.0

tests/datasets/test_usavars.py

@@ -0,0 +1,135 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import builtins
+import os
+import shutil
+from pathlib import Path
+from typing import Any, Generator
+
+import pytest
+import torch
+import torch.nn as nn
+from _pytest.fixtures import SubRequest
+from _pytest.monkeypatch import MonkeyPatch
+from matplotlib import pyplot as plt
+from torch.utils.data import ConcatDataset
+
+import torchgeo.datasets.utils
+from torchgeo.datasets import USAVars
+
+pytest.importorskip("pandas", minversion="0.19.1")
+
+
+def download_url(url: str, root: str, *args: str, **kwargs: str) -> None:
+    shutil.copy(url, root)
+
+
+class TestUSAVars:
+    @pytest.fixture()
+    def dataset(
+        self,
+        monkeypatch: Generator[MonkeyPatch, None, None],
+        tmp_path: Path,
+        request: SubRequest,
+    ) -> USAVars:
+
+        monkeypatch.setattr(  # type: ignore[attr-defined]
+            torchgeo.datasets.usavars, "download_url", download_url
+        )
+
+        md5 = "b504580a00bdc27097d5421dec50481b"
+        monkeypatch.setattr(USAVars, "md5", md5)  # type: ignore[attr-defined]
+
+        data_url = os.path.join("tests", "data", "usavars", "uar.zip")
+        monkeypatch.setattr(USAVars, "data_url", data_url)  # type: ignore[attr-defined]
+
+        label_urls = {
+            "elevation": os.path.join("tests", "data", "usavars", "elevation.csv"),
+            "population": os.path.join("tests", "data", "usavars", "population.csv"),
+            "treecover": os.path.join("tests", "data", "usavars", "treecover.csv"),
+            "income": os.path.join("tests", "data", "usavars", "income.csv"),
+            "nightlights": os.path.join("tests", "data", "usavars", "nightlights.csv"),
+            "roads": os.path.join("tests", "data", "usavars", "roads.csv"),
+            "housing": os.path.join("tests", "data", "usavars", "housing.csv"),
+        }
+        monkeypatch.setattr(  # type: ignore[attr-defined]
+            USAVars, "label_urls", label_urls
+        )
+
+        root = str(tmp_path)
+        transforms = nn.Identity()  # type: ignore[attr-defined]
+
+        return USAVars(root, transforms=transforms, download=True, checksum=True)
+
+    def test_getitem(self, dataset: USAVars) -> None:
+        x = dataset[0]
+        assert isinstance(x, dict)
+        assert isinstance(x["image"], torch.Tensor)
+        assert x["image"].ndim == 3
+        assert len(x.keys()) == 2  # image, elevation, population, treecover
+        assert x["image"].shape[0] == 4  # R, G, B, Inf
+
+    def test_len(self, dataset: USAVars) -> None:
+        assert len(dataset) == 2
+
+    def test_add(self, dataset: USAVars) -> None:
+        ds = dataset + dataset
+        assert isinstance(ds, ConcatDataset)
+
+    def test_already_extracted(self, dataset: USAVars) -> None:
+        USAVars(root=dataset.root, download=True)
+
+    def test_already_downloaded(self, tmp_path: Path) -> None:
+        pathname = os.path.join("tests", "data", "usavars", "uar.zip")
+        root = str(tmp_path)
+        shutil.copy(pathname, root)
+        csvs = [
+            "elevation.csv",
+            "population.csv",
+            "treecover.csv",
+            "income.csv",
+            "nightlights.csv",
+            "roads.csv",
+            "housing.csv",
+        ]
+        for csv in csvs:
+            shutil.copy(os.path.join("tests", "data", "usavars", csv), root)
+
+        USAVars(root)
+
+    def test_not_downloaded(self, tmp_path: Path) -> None:
+        with pytest.raises(RuntimeError, match="Dataset not found"):
+            USAVars(str(tmp_path))
+
+    @pytest.fixture(params=["pandas"])
+    def mock_missing_module(
+        self, monkeypatch: Generator[MonkeyPatch, None, None], request: SubRequest
+    ) -> str:
+        import_orig = builtins.__import__
+        package = str(request.param)
+
+        def mocked_import(name: str, *args: Any, **kwargs: Any) -> Any:
+            if name == package:
+                raise ImportError()
+            return import_orig(name, *args, **kwargs)
+
+        monkeypatch.setattr(  # type: ignore[attr-defined]
+            builtins, "__import__", mocked_import
+        )
+        return package
+
+    def test_mock_missing_module(
+        self, dataset: USAVars, mock_missing_module: str
+    ) -> None:
+        package = mock_missing_module
+        if package == "pandas":
+            with pytest.raises(
+                ImportError,
+                match=f"{package} is not installed and is required to use this dataset",
+            ):
+                USAVars(dataset.root)
+
+    def test_plot(self, dataset: USAVars) -> None:
+        dataset.plot(dataset[0], suptitle="Test")
+        plt.close()

torchgeo/datasets/__init__.py

@@ -72,6 +72,7 @@
 from .so2sat import So2Sat
 from .spacenet import SpaceNet, SpaceNet1, SpaceNet2, SpaceNet4, SpaceNet5, SpaceNet7
 from .ucmerced import UCMerced
+from .usavars import USAVars
 from .utils import (
     BoundingBox,
     concat_samples,
@@ -149,6 +150,7 @@
     "SpaceNet7",
     "TropicalCycloneWindEstimation",
     "UCMerced",
+    "USAVars",
     "Vaihingen2D",
     "VHR10",
     "XView2",