LandCover.ai: add GeoDataset version

docs/api/datasets.rst

@@ -93,6 +93,12 @@ iNaturalist
 
 .. autoclass:: INaturalist
 
+LandCover.ai Geo
+^^^^^^^^^^^^^^^^
+
+.. autoclass:: LandCoverAIBase
+.. autoclass:: LandCoverAIGeo
+
 Landsat
 ^^^^^^^
 

docs/api/geo_datasets.csv

@@ -12,6 +12,7 @@ Dataset,Type,Source,Size (px),Resolution (m)
 `GBIF`_,Points,Citizen Scientists,-,-
 `GlobBiomass`_,Masks,Landsat,"45,000x45,000",100
 `iNaturalist`_,Points,Citizen Scientists,-,-
+`LandCover.ai Geo`_,"Imagery, Masks",Aerial,"4,200--9,500",0.25--0.5
 `Landsat`_,Imagery,Landsat,"8,900x8,900",30
 `NAIP`_,Imagery,Aerial,"6,100x7,600",1
 `Open Buildings`_,Geometries,"Maxar, CNES/Airbus",-,-

tests/datasets/test_landcoverai.py

@@ -14,13 +14,62 @@
 from torch.utils.data import ConcatDataset
 
 import torchgeo.datasets.utils
-from torchgeo.datasets import LandCoverAI
+from torchgeo.datasets import BoundingBox, LandCoverAI, LandCoverAIGeo
 
 
 def download_url(url: str, root: str, *args: str, **kwargs: str) -> None:
     shutil.copy(url, root)
 
 
+class TestLandCoverAIGeo:
+    @pytest.fixture
+    def dataset(self, monkeypatch: MonkeyPatch, tmp_path: Path) -> LandCoverAIGeo:
+        monkeypatch.setattr(torchgeo.datasets.landcoverai, "download_url", download_url)
+        md5 = "ff8998857cc8511f644d3f7d0f3688d0"
+        monkeypatch.setattr(LandCoverAIGeo, "md5", md5)
+        url = os.path.join("tests", "data", "landcoverai", "landcover.ai.v1.zip")
+        monkeypatch.setattr(LandCoverAIGeo, "url", url)
+        root = str(tmp_path)
+        transforms = nn.Identity()
+        return LandCoverAIGeo(root, transforms=transforms, download=True, checksum=True)
+
+    def test_getitem(self, dataset: LandCoverAIGeo) -> None:
+        x = dataset[dataset.bounds]
+        assert isinstance(x, dict)
+        assert isinstance(x["image"], torch.Tensor)
+        assert isinstance(x["mask"], torch.Tensor)
+
+    def test_already_extracted(self, dataset: LandCoverAIGeo) -> None:
+        LandCoverAIGeo(root=dataset.root, download=True)
+
+    def test_already_downloaded(self, monkeypatch: MonkeyPatch, tmp_path: Path) -> None:
+        url = os.path.join("tests", "data", "landcoverai", "landcover.ai.v1.zip")
+        root = str(tmp_path)
+        shutil.copy(url, root)
+        LandCoverAIGeo(root)
+
+    def test_not_downloaded(self, tmp_path: Path) -> None:
+        with pytest.raises(RuntimeError, match="Dataset not found"):
+            LandCoverAIGeo(str(tmp_path))
+
+    def test_out_of_bounds_query(self, dataset: LandCoverAIGeo) -> None:
+        query = BoundingBox(0, 0, 0, 0, 0, 0)
+        with pytest.raises(
+            IndexError, match="query: .* not found in index with bounds:"
+        ):
+            dataset[query]
+
+    def test_plot(self, dataset: LandCoverAIGeo) -> None:
+        x = dataset[dataset.bounds].copy()
+        dataset.plot(x, suptitle="Test")
+        plt.close()
+        dataset.plot(x, show_titles=False)
+        plt.close()
+        x["prediction"] = x["mask"][:, :, 0].clone().unsqueeze(2)
+        dataset.plot(x)
+        plt.close()
+
+
 class TestLandCoverAI:
     @pytest.fixture(params=["train", "val", "test"])
     def dataset(

torchgeo/datasets/__init__.py

@@ -53,7 +53,7 @@
 from .idtrees import IDTReeS
 from .inaturalist import INaturalist
 from .inria import InriaAerialImageLabeling
-from .landcoverai import LandCoverAI
+from .landcoverai import LandCoverAI, LandCoverAIBase, LandCoverAIGeo
 from .landsat import (
     Landsat,
     Landsat1,
@@ -137,6 +137,8 @@
     "GBIF",
     "GlobBiomass",
     "INaturalist",
+    "LandCoverAIBase",
+    "LandCoverAIGeo",
     "Landsat",
     "Landsat1",
     "Landsat2",