diff --git a/docs/api/datasets.rst b/docs/api/datasets.rst
index ea3615dc404..10d49dde8f5 100644
--- a/docs/api/datasets.rst
+++ b/docs/api/datasets.rst
@@ -288,6 +288,7 @@ SpaceNet
.. autoclass:: SpaceNet3
.. autoclass:: SpaceNet4
.. autoclass:: SpaceNet5
+.. autoclass:: SpaceNet6
.. autoclass:: SpaceNet7
Tropical Cyclone
diff --git a/tests/data/spacenet/data.py b/tests/data/spacenet/data.py
new file mode 100644
index 00000000000..37d67e98518
--- /dev/null
+++ b/tests/data/spacenet/data.py
@@ -0,0 +1,283 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+import shutil
+from collections import OrderedDict
+from typing import List, cast
+
+import fiona
+import numpy as np
+import rasterio
+from rasterio.crs import CRS
+from rasterio.transform import Affine
+from torchvision.datasets.utils import calculate_md5
+
+from torchgeo.datasets import (
+ SpaceNet,
+ SpaceNet1,
+ SpaceNet2,
+ SpaceNet3,
+ SpaceNet4,
+ SpaceNet5,
+ SpaceNet6,
+ SpaceNet7,
+)
+
+transform = Affine(0.3, 0.0, 616500.0, 0.0, -0.3, 3345000.0)
+crs = CRS.from_epsg(4326)
+
+img_count = {
+ "MS.tif": 8,
+ "PAN.tif": 1,
+ "PS-MS.tif": 8,
+ "PS-RGB.tif": 3,
+ "PS-RGBNIR.tif": 4,
+ "RGB.tif": 3,
+ "RGBNIR.tif": 4,
+ "SAR-Intensity.tif": 1,
+ "mosaic.tif": 3,
+ "8Band.tif": 8,
+}
+
+
+sn4_catalog = [
+ "10300100023BC100",
+ "10300100036D5200",
+ "1030010003BDDC00",
+ "1030010003CD4300",
+]
+sn4_angles = [8, 30, 52, 53]
+
+sn4_imgdirname = "sn4_SN4_buildings_train_AOI_6_Atlanta_732701_3730989-nadir{}_catid_{}"
+sn4_lbldirname = "sn4_SN4_buildings_train_AOI_6_Atlanta_732701_3730989-labels"
+sn4_emptyimgdirname = (
+ "sn4_SN4_buildings_train_AOI_6_Atlanta_732701_3720639-nadir53_"
+ + "catid_1030010003CD4300"
+)
+sn4_emptylbldirname = "sn4_SN4_buildings_train_AOI_6_Atlanta_732701_3720639-labels"
+
+
+datasets = [SpaceNet1, SpaceNet2, SpaceNet3, SpaceNet4, SpaceNet5, SpaceNet6, SpaceNet7]
+
+
+def create_test_image(img_dir: str, imgs: List[str]) -> List[List[float]]:
+ """Create test image
+
+ Args:
+ img_dir (str): Name of image directory
+ imgs (List[str]): List of images to be created
+
+ Returns:
+ List[List[float]]: Boundary coordinates
+ """
+ for img in imgs:
+ imgpath = os.path.join(img_dir, img)
+ Z = np.arange(4, dtype="uint16").reshape(2, 2)
+ count = img_count[img]
+ with rasterio.open(
+ imgpath,
+ "w",
+ driver="GTiff",
+ height=Z.shape[0],
+ width=Z.shape[1],
+ count=count,
+ dtype=Z.dtype,
+ crs=crs,
+ transform=transform,
+ ) as dst:
+ for i in range(1, dst.count + 1):
+ dst.write(Z, i)
+
+ tim = rasterio.open(imgpath)
+ slice_index = [[1, 1], [1, 2], [2, 2], [2, 1], [1, 1]]
+ return [list(tim.transform * p) for p in slice_index]
+
+
+def create_test_label(
+ lbldir: str,
+ lblname: str,
+ coords: List[List[float]],
+ det_type: str,
+ empty: bool = False,
+ diff_crs: bool = False,
+) -> None:
+ """Create test label
+
+ Args:
+ lbldir (str): Name of label directory
+ lblname (str): Name of label file
+ coords (List[Tuple[float, float]]): Boundary coordinates
+ det_type (str): Type of dataset. Must be either buildings or roads.
+ empty (bool, optional): Creates empty label file if True. Defaults to False.
+ diff_crs (bool, optional): Assigns EPSG:3857 as CRS instead of
+ default EPSG:4326. Defaults to False.
+ """
+ if empty:
+ # Creates a new file
+ with open(os.path.join(lbldir, lblname), "w"):
+ pass
+ return
+
+ if det_type == "buildings":
+ meta_properties = OrderedDict()
+ geom = "Polygon"
+ rec = {
+ "type": "Feature",
+ "id": "0",
+ "properties": OrderedDict(),
+ "geometry": {"type": "Polygon", "coordinates": [coords]},
+ }
+ else:
+ meta_properties = OrderedDict(
+ [
+ ("heading", "str"),
+ ("lane_number", "str"),
+ ("one_way_ty", "str"),
+ ("paved", "str"),
+ ("road_id", "int"),
+ ("road_type", "str"),
+ ("origarea", "int"),
+ ("origlen", "float"),
+ ("partialDec", "int"),
+ ("truncated", "int"),
+ ("bridge_type", "str"),
+ ("inferred_speed_mph", "float"),
+ ("inferred_speed_mps", "float"),
+ ]
+ )
+ geom = "LineString"
+
+ dummy_vals = {"str": "a", "float": 45.0, "int": 0}
+ ROAD_DICT = [(k, dummy_vals[v]) for k, v in meta_properties.items()]
+ rec = {
+ "type": "Feature",
+ "id": "0",
+ "properties": OrderedDict(ROAD_DICT),
+ "geometry": {"type": "LineString", "coordinates": [coords[0], coords[2]]},
+ }
+
+ meta = {
+ "driver": "GeoJSON",
+ "schema": {"properties": meta_properties, "geometry": geom},
+ "crs": {"init": "epsg:4326"},
+ }
+ if diff_crs:
+ meta["crs"] = {"init": "epsg:3857"}
+ out_file = os.path.join(lbldir, lblname)
+ with fiona.open(out_file, "w", **meta) as dst:
+ dst.write(rec)
+
+
+def main() -> None:
+ ROOT_DIR = os.path.dirname(os.path.realpath(__file__))
+
+ num_samples = 2
+
+ for dataset in datasets:
+
+ collections = list(dataset.collection_md5_dict.keys())
+ for collection in collections:
+ dataset = cast(SpaceNet, dataset)
+ if dataset.dataset_id == "spacenet4":
+ num_samples = 4
+ elif collection == "sn5_AOI_7_Moscow" or collection not in [
+ "sn5_AOI_8_Mumbai",
+ "sn7_test_source",
+ ]:
+ num_samples = 2
+ elif collection == "sn5_AOI_8_Mumbai":
+ num_samples = 3
+ else:
+ num_samples = 1
+ for sample in range(num_samples):
+ out_dir = os.path.join(ROOT_DIR, collection)
+ if collection == "sn6_AOI_11_Rotterdam":
+ out_dir = os.path.join(ROOT_DIR, "spacenet6", collection)
+
+ # Create img dir
+ if dataset.dataset_id == "spacenet4":
+ assert num_samples == 4
+ if sample != 3:
+ imgdirname = sn4_imgdirname.format(
+ sn4_angles[sample], sn4_catalog[sample]
+ )
+ lbldirname = sn4_lbldirname
+ else:
+ imgdirname = sn4_emptyimgdirname.format(
+ sn4_angles[sample], sn4_catalog[sample]
+ )
+ lbldirname = sn4_emptylbldirname
+ else:
+ imgdirname = f"{collection}_img{sample + 1}"
+ lbldirname = f"{collection}_img{sample + 1}-labels"
+
+ imgdir = os.path.join(out_dir, imgdirname)
+ os.makedirs(imgdir, exist_ok=True)
+ bounds = create_test_image(imgdir, list(dataset.imagery.values()))
+
+ # Create lbl dir
+ lbldir = os.path.join(out_dir, lbldirname)
+ os.makedirs(lbldir, exist_ok=True)
+ det_type = "roads" if dataset in [SpaceNet3, SpaceNet5] else "buildings"
+ if dataset.dataset_id == "spacenet4" and sample == 3:
+ # Creates an empty file
+ create_test_label(
+ lbldir, dataset.label_glob, bounds, det_type, empty=True
+ )
+ else:
+ create_test_label(lbldir, dataset.label_glob, bounds, det_type)
+
+ if collection == "sn5_AOI_8_Mumbai":
+ if sample == 1:
+ create_test_label(
+ lbldir, dataset.label_glob, bounds, det_type, empty=True
+ )
+ if sample == 2:
+ create_test_label(
+ lbldir, dataset.label_glob, bounds, det_type, diff_crs=True
+ )
+
+ if collection == "sn1_AOI_1_RIO" and sample == 1:
+ create_test_label(
+ lbldir, dataset.label_glob, bounds, det_type, diff_crs=True
+ )
+
+ if collection not in [
+ "sn2_AOI_2_Vegas",
+ "sn3_AOI_5_Khartoum",
+ "sn4_AOI_6_Atlanta",
+ "sn5_AOI_8_Mumbai",
+ "sn6_AOI_11_Rotterdam",
+ "sn7_train_source",
+ ]:
+ # Create collection.json
+ with open(
+ os.path.join(ROOT_DIR, collection, "collection.json"), "w"
+ ):
+ pass
+ if collection == "sn6_AOI_11_Rotterdam":
+ # Create collection.json
+ with open(
+ os.path.join(
+ ROOT_DIR, "spacenet6", collection, "collection.json"
+ ),
+ "w",
+ ):
+ pass
+
+ # Create archive
+ if collection == "sn6_AOI_11_Rotterdam":
+ break
+ archive_path = os.path.join(ROOT_DIR, collection)
+ shutil.make_archive(
+ archive_path, "gztar", root_dir=ROOT_DIR, base_dir=collection
+ )
+ shutil.rmtree(out_dir)
+ print(f'{collection}: {calculate_md5(f"{archive_path}.tar.gz")}')
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tests/data/spacenet/sn1_AOI_1_RIO.tar.gz b/tests/data/spacenet/sn1_AOI_1_RIO.tar.gz
index 231c81e1c56..bf7ad82733b 100644
Binary files a/tests/data/spacenet/sn1_AOI_1_RIO.tar.gz and b/tests/data/spacenet/sn1_AOI_1_RIO.tar.gz differ
diff --git a/tests/data/spacenet/sn2_AOI_2_Vegas.tar.gz b/tests/data/spacenet/sn2_AOI_2_Vegas.tar.gz
index fe01659fe2e..c7cbcd5b4fc 100644
Binary files a/tests/data/spacenet/sn2_AOI_2_Vegas.tar.gz and b/tests/data/spacenet/sn2_AOI_2_Vegas.tar.gz differ
diff --git a/tests/data/spacenet/sn2_AOI_3_Paris.tar.gz b/tests/data/spacenet/sn2_AOI_3_Paris.tar.gz
index 3eaec04c45d..6c26bde44d7 100644
Binary files a/tests/data/spacenet/sn2_AOI_3_Paris.tar.gz and b/tests/data/spacenet/sn2_AOI_3_Paris.tar.gz differ
diff --git a/tests/data/spacenet/sn2_AOI_4_Shanghai.tar.gz b/tests/data/spacenet/sn2_AOI_4_Shanghai.tar.gz
index 138a9079e43..b7d8ba655e1 100644
Binary files a/tests/data/spacenet/sn2_AOI_4_Shanghai.tar.gz and b/tests/data/spacenet/sn2_AOI_4_Shanghai.tar.gz differ
diff --git a/tests/data/spacenet/sn2_AOI_5_Khartoum.tar.gz b/tests/data/spacenet/sn2_AOI_5_Khartoum.tar.gz
index f0131ad3b8c..b0a4b29ca34 100644
Binary files a/tests/data/spacenet/sn2_AOI_5_Khartoum.tar.gz and b/tests/data/spacenet/sn2_AOI_5_Khartoum.tar.gz differ
diff --git a/tests/data/spacenet/sn3_AOI_2_Vegas.tar.gz b/tests/data/spacenet/sn3_AOI_2_Vegas.tar.gz
new file mode 100644
index 00000000000..0e17e78befc
Binary files /dev/null and b/tests/data/spacenet/sn3_AOI_2_Vegas.tar.gz differ
diff --git a/tests/data/spacenet/sn3_AOI_3_Paris.tar.gz b/tests/data/spacenet/sn3_AOI_3_Paris.tar.gz
index d8cb305dc35..3a960eca9af 100644
Binary files a/tests/data/spacenet/sn3_AOI_3_Paris.tar.gz and b/tests/data/spacenet/sn3_AOI_3_Paris.tar.gz differ
diff --git a/tests/data/spacenet/sn3_AOI_4_Shanghai.tar.gz b/tests/data/spacenet/sn3_AOI_4_Shanghai.tar.gz
new file mode 100644
index 00000000000..f3b479e43b1
Binary files /dev/null and b/tests/data/spacenet/sn3_AOI_4_Shanghai.tar.gz differ
diff --git a/tests/data/spacenet/sn3_AOI_5_Khartoum.tar.gz b/tests/data/spacenet/sn3_AOI_5_Khartoum.tar.gz
index 0daea2f5392..f3a1b809291 100644
Binary files a/tests/data/spacenet/sn3_AOI_5_Khartoum.tar.gz and b/tests/data/spacenet/sn3_AOI_5_Khartoum.tar.gz differ
diff --git a/tests/data/spacenet/sn4_AOI_6_Atlanta.tar.gz b/tests/data/spacenet/sn4_AOI_6_Atlanta.tar.gz
index 1c382ff25dd..a1e0c8a6910 100644
Binary files a/tests/data/spacenet/sn4_AOI_6_Atlanta.tar.gz and b/tests/data/spacenet/sn4_AOI_6_Atlanta.tar.gz differ
diff --git a/tests/data/spacenet/sn5_AOI_7_Moscow.tar.gz b/tests/data/spacenet/sn5_AOI_7_Moscow.tar.gz
index acb289ebacd..c666f1b837c 100644
Binary files a/tests/data/spacenet/sn5_AOI_7_Moscow.tar.gz and b/tests/data/spacenet/sn5_AOI_7_Moscow.tar.gz differ
diff --git a/tests/data/spacenet/sn5_AOI_8_Mumbai.tar.gz b/tests/data/spacenet/sn5_AOI_8_Mumbai.tar.gz
index a6c6e353c27..4f6b9cd7ad4 100644
Binary files a/tests/data/spacenet/sn5_AOI_8_Mumbai.tar.gz and b/tests/data/spacenet/sn5_AOI_8_Mumbai.tar.gz differ
diff --git a/tests/data/spacenet/sn7_test_source.tar.gz b/tests/data/spacenet/sn7_test_source.tar.gz
index f1ca4f1c87e..e411894fb3a 100644
Binary files a/tests/data/spacenet/sn7_test_source.tar.gz and b/tests/data/spacenet/sn7_test_source.tar.gz differ
diff --git a/tests/data/spacenet/sn7_train_labels.tar.gz b/tests/data/spacenet/sn7_train_labels.tar.gz
index ae3bb5f4004..b3f583771c0 100644
Binary files a/tests/data/spacenet/sn7_train_labels.tar.gz and b/tests/data/spacenet/sn7_train_labels.tar.gz differ
diff --git a/tests/data/spacenet/sn7_train_source.tar.gz b/tests/data/spacenet/sn7_train_source.tar.gz
index c2f8b2c8e5b..e847fd070a4 100644
Binary files a/tests/data/spacenet/sn7_train_source.tar.gz and b/tests/data/spacenet/sn7_train_source.tar.gz differ
diff --git a/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/collection.json b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/collection.json
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1-labels/labels.geojson b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1-labels/labels.geojson
new file mode 100644
index 00000000000..0a418938820
--- /dev/null
+++ b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1-labels/labels.geojson
@@ -0,0 +1,7 @@
+{
+"type": "FeatureCollection",
+"crs": { "type": "name", "properties": { "name": "urn:ogc:def:crs:OGC:1.3:CRS84" } },
+"features": [
+{ "type": "Feature", "properties": { }, "geometry": { "type": "Polygon", "coordinates": [ [ [ 616500.300000000046566, 3344999.700000000186265 ], [ 616500.300000000046566, 3344999.4 ], [ 616500.599999999976717, 3344999.4 ], [ 616500.599999999976717, 3344999.700000000186265 ], [ 616500.300000000046566, 3344999.700000000186265 ] ] ] } }
+]
+}
diff --git a/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1/PAN.tif b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1/PAN.tif
new file mode 100644
index 00000000000..a9aef1da576
Binary files /dev/null and b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1/PAN.tif differ
diff --git a/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1/PS-RGB.tif b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1/PS-RGB.tif
new file mode 100644
index 00000000000..022510c2df5
Binary files /dev/null and b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1/PS-RGB.tif differ
diff --git a/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1/PS-RGBNIR.tif b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1/PS-RGBNIR.tif
new file mode 100644
index 00000000000..daadc4a2e39
Binary files /dev/null and b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1/PS-RGBNIR.tif differ
diff --git a/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1/RGBNIR.tif b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1/RGBNIR.tif
new file mode 100644
index 00000000000..daadc4a2e39
Binary files /dev/null and b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1/RGBNIR.tif differ
diff --git a/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1/SAR-Intensity.tif b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1/SAR-Intensity.tif
new file mode 100644
index 00000000000..a9aef1da576
Binary files /dev/null and b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img1/SAR-Intensity.tif differ
diff --git a/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2-labels/labels.geojson b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2-labels/labels.geojson
new file mode 100644
index 00000000000..0a418938820
--- /dev/null
+++ b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2-labels/labels.geojson
@@ -0,0 +1,7 @@
+{
+"type": "FeatureCollection",
+"crs": { "type": "name", "properties": { "name": "urn:ogc:def:crs:OGC:1.3:CRS84" } },
+"features": [
+{ "type": "Feature", "properties": { }, "geometry": { "type": "Polygon", "coordinates": [ [ [ 616500.300000000046566, 3344999.700000000186265 ], [ 616500.300000000046566, 3344999.4 ], [ 616500.599999999976717, 3344999.4 ], [ 616500.599999999976717, 3344999.700000000186265 ], [ 616500.300000000046566, 3344999.700000000186265 ] ] ] } }
+]
+}
diff --git a/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2/PAN.tif b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2/PAN.tif
new file mode 100644
index 00000000000..a9aef1da576
Binary files /dev/null and b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2/PAN.tif differ
diff --git a/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2/PS-RGB.tif b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2/PS-RGB.tif
new file mode 100644
index 00000000000..022510c2df5
Binary files /dev/null and b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2/PS-RGB.tif differ
diff --git a/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2/PS-RGBNIR.tif b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2/PS-RGBNIR.tif
new file mode 100644
index 00000000000..daadc4a2e39
Binary files /dev/null and b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2/PS-RGBNIR.tif differ
diff --git a/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2/RGBNIR.tif b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2/RGBNIR.tif
new file mode 100644
index 00000000000..daadc4a2e39
Binary files /dev/null and b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2/RGBNIR.tif differ
diff --git a/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2/SAR-Intensity.tif b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2/SAR-Intensity.tif
new file mode 100644
index 00000000000..a9aef1da576
Binary files /dev/null and b/tests/data/spacenet/spacenet6/sn6_AOI_11_Rotterdam/sn6_AOI_11_Rotterdam_img2/SAR-Intensity.tif differ
diff --git a/tests/datasets/test_spacenet.py b/tests/datasets/test_spacenet.py
index bf84b60633d..a20d6e9b2d8 100644
--- a/tests/datasets/test_spacenet.py
+++ b/tests/datasets/test_spacenet.py
@@ -19,10 +19,12 @@
SpaceNet3,
SpaceNet4,
SpaceNet5,
+ SpaceNet6,
SpaceNet7,
)
TEST_DATA_DIR = "tests/data/spacenet"
+radiant_mlhub = pytest.importorskip("radiant_mlhub", minversion="0.2.1")
class Collection:
@@ -35,18 +37,33 @@ def download(self, output_dir: str, **kwargs: str) -> None:
shutil.copy(tarball, output_dir)
+class Dataset:
+ def __init__(self, dataset_id: str) -> None:
+ self.dataset_id = dataset_id
+
+ def download(self, output_dir: str, **kwargs: str) -> None:
+ glob_path = os.path.join(TEST_DATA_DIR, "spacenet*")
+ for directory in glob.iglob(glob_path):
+ dataset_name = os.path.basename(directory)
+ output_dir = os.path.join(output_dir, dataset_name)
+ shutil.copytree(directory, output_dir)
+
+
def fetch_collection(collection_id: str, **kwargs: str) -> Collection:
return Collection(collection_id)
+def fetch_dataset(dataset_id: str, **kwargs: str) -> Dataset:
+ return Dataset(dataset_id)
+
+
class TestSpaceNet1:
@pytest.fixture(params=["rgb", "8band"])
def dataset(
self, request: SubRequest, monkeypatch: MonkeyPatch, tmp_path: Path
) -> SpaceNet1:
- radiant_mlhub = pytest.importorskip("radiant_mlhub", minversion="0.2.1")
monkeypatch.setattr(radiant_mlhub.Collection, "fetch", fetch_collection)
- test_md5 = {"sn1_AOI_1_RIO": "829652022c2df4511ee4ae05bc290250"}
+ test_md5 = {"sn1_AOI_1_RIO": "246e27fcd7ae73496212a7f585a43dbb"}
# Refer https://github.com/python/mypy/issues/1032
monkeypatch.setattr(SpaceNet1, "collection_md5_dict", test_md5)
@@ -58,6 +75,7 @@ def dataset(
def test_getitem(self, dataset: SpaceNet1) -> None:
x = dataset[0]
+ dataset[1]
assert isinstance(x, dict)
assert isinstance(x["image"], torch.Tensor)
assert isinstance(x["mask"], torch.Tensor)
@@ -90,13 +108,12 @@ class TestSpaceNet2:
def dataset(
self, request: SubRequest, monkeypatch: MonkeyPatch, tmp_path: Path
) -> SpaceNet2:
- radiant_mlhub = pytest.importorskip("radiant_mlhub", minversion="0.2.1")
monkeypatch.setattr(radiant_mlhub.Collection, "fetch", fetch_collection)
test_md5 = {
- "sn2_AOI_2_Vegas": "6ceae7ff8c557346e8a4c8b6c61cc1b9",
- "sn2_AOI_3_Paris": "811e6a26fdeb8be445fed99769fa52c5",
- "sn2_AOI_4_Shanghai": "139d1627d184c74426a85ad0222f7355",
- "sn2_AOI_5_Khartoum": "435535120414b74165aa87f051c3a2b3",
+ "sn2_AOI_2_Vegas": "131048686ba21a45853c05f227f40b7f",
+ "sn2_AOI_3_Paris": "62242fd198ee32b59f0178cf656e1513",
+ "sn2_AOI_4_Shanghai": "563b0817ecedd8ff3b3e4cb2991bf3fb",
+ "sn2_AOI_5_Khartoum": "e4185a2e9a12cf7b3d0cd1db6b3e0f06",
}
monkeypatch.setattr(SpaceNet2, "collection_md5_dict", test_md5)
@@ -123,9 +140,8 @@ def test_getitem(self, dataset: SpaceNet2) -> None:
else:
assert x["image"].shape[0] == 1
- # TODO: Change len to 4 when radiantearth/radiant-mlhub#65 is fixed
def test_len(self, dataset: SpaceNet2) -> None:
- assert len(dataset) == 5
+ assert len(dataset) == 4
def test_already_downloaded(self, dataset: SpaceNet2) -> None:
SpaceNet2(root=dataset.root, download=True)
@@ -153,11 +169,10 @@ class TestSpaceNet3:
def dataset(
self, request: SubRequest, monkeypatch: MonkeyPatch, tmp_path: Path
) -> SpaceNet3:
- radiant_mlhub = pytest.importorskip("radiant_mlhub", minversion="0.2.1")
monkeypatch.setattr(radiant_mlhub.Collection, "fetch", fetch_collection)
test_md5 = {
- "sn3_AOI_3_Paris": "197440e0ade970169a801a173a492c27",
- "sn3_AOI_5_Khartoum": "b21ff7dd33a15ec32bd380c083263cdf",
+ "sn3_AOI_3_Paris": "93452c68da11dd6b57dc83dba43c2c9d",
+ "sn3_AOI_5_Khartoum": "7c9d96810198bf101cbaf54f7a5e8b3b",
}
monkeypatch.setattr(SpaceNet3, "collection_md5_dict", test_md5)
@@ -218,9 +233,8 @@ class TestSpaceNet4:
def dataset(
self, request: SubRequest, monkeypatch: MonkeyPatch, tmp_path: Path
) -> SpaceNet4:
- radiant_mlhub = pytest.importorskip("radiant_mlhub", minversion="0.2.1")
monkeypatch.setattr(radiant_mlhub.Collection, "fetch", fetch_collection)
- test_md5 = {"sn4_AOI_6_Atlanta": "ea37c2d87e2c3a1d8b2a7c2230080d46"}
+ test_md5 = {"sn4_AOI_6_Atlanta": "097a76a2319b7ba34dac1722862fc93b"}
test_angles = ["nadir", "off-nadir", "very-off-nadir"]
@@ -281,11 +295,10 @@ class TestSpaceNet5:
def dataset(
self, request: SubRequest, monkeypatch: MonkeyPatch, tmp_path: Path
) -> SpaceNet5:
- radiant_mlhub = pytest.importorskip("radiant_mlhub", minversion="0.2.1")
monkeypatch.setattr(radiant_mlhub.Collection, "fetch", fetch_collection)
test_md5 = {
- "sn5_AOI_7_Moscow": "e0d5f41f1b6b0ee7696c15e5ff3141f5",
- "sn5_AOI_8_Mumbai": "ab898700ee586a137af492b84a08e662",
+ "sn5_AOI_7_Moscow": "5c511dd31eea739cc1f81ef5962f3d56",
+ "sn5_AOI_8_Mumbai": "e00452b87bbe87feaef65f373be3978e",
}
monkeypatch.setattr(SpaceNet5, "collection_md5_dict", test_md5)
@@ -339,17 +352,55 @@ def test_plot(self, dataset: SpaceNet5) -> None:
plt.close()
+class TestSpaceNet6:
+ @pytest.fixture(params=["PAN", "RGBNIR", "PS-RGB", "PS-RGBNIR", "SAR-Intensity"])
+ def dataset(
+ self, request: SubRequest, monkeypatch: MonkeyPatch, tmp_path: Path
+ ) -> SpaceNet6:
+ monkeypatch.setattr(radiant_mlhub.Dataset, "fetch", fetch_dataset)
+ root = str(tmp_path)
+ transforms = nn.Identity()
+ return SpaceNet6(
+ root, image=request.param, transforms=transforms, download=True, api_key=""
+ )
+
+ def test_getitem(self, dataset: SpaceNet6) -> None:
+ x = dataset[0]
+ assert isinstance(x, dict)
+ assert isinstance(x["image"], torch.Tensor)
+ assert isinstance(x["mask"], torch.Tensor)
+ if dataset.image == "PS-RGB":
+ assert x["image"].shape[0] == 3
+ elif dataset.image in ["RGBNIR", "PS-RGBNIR"]:
+ assert x["image"].shape[0] == 4
+ else:
+ assert x["image"].shape[0] == 1
+
+ def test_len(self, dataset: SpaceNet6) -> None:
+ assert len(dataset) == 2
+
+ def test_already_downloaded(self, dataset: SpaceNet6) -> None:
+ SpaceNet6(root=dataset.root, download=True)
+
+ def test_plot(self, dataset: SpaceNet6) -> None:
+ x = dataset[0].copy()
+ x["prediction"] = x["mask"]
+ dataset.plot(x, suptitle="Test")
+ plt.close()
+ dataset.plot(x, show_titles=False)
+ plt.close()
+
+
class TestSpaceNet7:
@pytest.fixture(params=["train", "test"])
def dataset(
self, request: SubRequest, monkeypatch: MonkeyPatch, tmp_path: Path
) -> SpaceNet7:
- radiant_mlhub = pytest.importorskip("radiant_mlhub", minversion="0.2.1")
monkeypatch.setattr(radiant_mlhub.Collection, "fetch", fetch_collection)
test_md5 = {
- "sn7_train_source": "254fd6b16e350b071137b2658332091f",
- "sn7_train_labels": "05befe86b037a3af75c7143553033664",
- "sn7_test_source": "37d98d44a9da39657ed4b7beee22a21e",
+ "sn7_train_source": "197bfa8842a40b09b6837b824a6370e0",
+ "sn7_train_labels": "625ad8a989a5105bc766a53e53df4d0e",
+ "sn7_test_source": "461f59eb21bb4f416c867f5037dfceeb",
}
monkeypatch.setattr(SpaceNet7, "collection_md5_dict", test_md5)
diff --git a/torchgeo/datasets/__init__.py b/torchgeo/datasets/__init__.py
index 6b6aa6fbc6e..f85ff83510d 100644
--- a/torchgeo/datasets/__init__.py
+++ b/torchgeo/datasets/__init__.py
@@ -92,6 +92,7 @@
SpaceNet3,
SpaceNet4,
SpaceNet5,
+ SpaceNet6,
SpaceNet7,
)
from .ucmerced import UCMerced
@@ -184,6 +185,7 @@
"SpaceNet3",
"SpaceNet4",
"SpaceNet5",
+ "SpaceNet6",
"SpaceNet7",
"TropicalCyclone",
"UCMerced",
diff --git a/torchgeo/datasets/spacenet.py b/torchgeo/datasets/spacenet.py
index 3bbf24daa11..3e5baa7ad75 100644
--- a/torchgeo/datasets/spacenet.py
+++ b/torchgeo/datasets/spacenet.py
@@ -28,6 +28,7 @@
from .utils import (
check_integrity,
download_radiant_mlhub_collection,
+ download_radiant_mlhub_dataset,
extract_archive,
percentile_normalization,
)
@@ -40,6 +41,13 @@ class SpaceNet(NonGeoDataset, abc.ABC):
datasets that all together contain >11M building footprints and ~20,000 km
of road labels mapped over high-resolution satellite imagery obtained from
a variety of sensors such as Worldview-2, Worldview-3 and Dove.
+
+ .. note::
+
+ The SpaceNet datasets require the following additional library to be installed:
+
+ * `radiant-mlhub `_ to download the
+ imagery and labels from the Radiant Earth MLHub
"""
@property
@@ -384,13 +392,6 @@ class SpaceNet1(SpaceNet):
* https://arxiv.org/abs/1807.01232
- .. note::
-
- This dataset requires the following additional library to be installed:
-
- * `radiant-mlhub `_ to download the
- imagery and labels from the Radiant Earth MLHub
-
"""
dataset_id = "spacenet1"
@@ -490,13 +491,6 @@ class SpaceNet2(SpaceNet):
* https://arxiv.org/abs/1807.01232
- .. note::
-
- This dataset requires the following additional library to be installed:
-
- * `radiant-mlhub `_ to download the
- imagery and labels from the Radiant Earth MLHub
-
"""
dataset_id = "spacenet2"
@@ -616,13 +610,6 @@ class SpaceNet3(SpaceNet):
* https://arxiv.org/abs/1807.01232
- .. note::
-
- This dataset requires the following additional library to be installed:
-
- * `radiant-mlhub `_ to download the
- imagery and labels from the Radiant Earth MLHub
-
.. versionadded:: 0.3
"""
@@ -853,13 +840,6 @@ class SpaceNet4(SpaceNet):
* https://arxiv.org/abs/1903.12239
- .. note::
-
- This dataset requires the following additional library to be installed:
-
- * `radiant-mlhub `_ to download the
- imagery and labels from the Radiant Earth MLHub
-
"""
dataset_id = "spacenet4"
@@ -1050,13 +1030,6 @@ class SpaceNet5(SpaceNet3):
Route Travel Time Estimation from Satellite Imageryā€¯,
https://spacenet.ai/sn5-challenge/
- .. note::
-
- This dataset requires the following additional library to be installed:
-
- * `radiant-mlhub `_ to download the
- imagery and labels from the Radiant Earth MLHub
-
.. versionadded:: 0.2
"""
@@ -1122,6 +1095,149 @@ def __init__(
)
+class SpaceNet6(SpaceNet):
+ r"""SpaceNet 6: Multi-Sensor All-Weather Mapping.
+
+ `SpaceNet 6 `_ is a dataset
+ of optical and SAR imagery over the city of Rotterdam.
+
+ Collection features:
+
+ +------------+---------------------+------------+-----------------------------+
+ | AOI | Area (km\ :sup:`2`\)| # Images | # Building Footprint Labels |
+ +============+=====================+============+=============================+
+ | Rotterdam | 120 | 3401 | 48000 |
+ +------------+---------------------+------------+-----------------------------+
+
+
+ Imagery features:
+
+ .. list-table::
+ :widths: 10 10 10 10 10 10
+ :header-rows: 1
+ :stub-columns: 1
+
+ * -
+ - PAN
+ - RGBNIR
+ - PS-RGB
+ - PS-RGBNIR
+ - SAR-Intensity
+ * - GSD (m)
+ - 0.5
+ - 2.0
+ - 0.5
+ - 0.5
+ - 0.5
+ * - Chip size (px)
+ - 900 x 900
+ - 450 x 450
+ - 900 x 900
+ - 900 x 900
+ - 900 x 900
+
+
+ Dataset format:
+
+ * Imagery - GeoTIFFs from Worldview-2 (optical) and Capella Space (SAR)
+
+ * PAN.tif (Panchromatic)
+ * RGBNIR.tif (Multispectral)
+ * PS-RGB (Pansharpened RGB)
+ * PS-RGBNIR (Pansharpened RGBNIR)
+ * SAR-Intensity (SAR Intensity)
+
+ * Labels - GeoJSON
+
+ * labels.geojson
+
+ If you use this dataset in your research, please cite the following paper:
+
+ * https://arxiv.org/abs/2004.06500
+
+ .. note::
+
+ This dataset requires the following additional library to be installed:
+
+ * `radiant-mlhub `_ to download the
+ imagery and labels from the Radiant Earth MLHub
+
+ .. versionadded:: 0.4
+ """
+
+ dataset_id = "spacenet6"
+ collections = ["sn6_AOI_11_Rotterdam"]
+ # This is actually the metadata hash
+ collection_md5_dict = {"sn6_AOI_11_Rotterdam": "66f7312218fec67a1e0b3b02b22c95cc"}
+ imagery = {
+ "PAN": "PAN.tif",
+ "RGBNIR": "RGBNIR.tif",
+ "PS-RGB": "PS-RGB.tif",
+ "PS-RGBNIR": "PS-RGBNIR.tif",
+ "SAR-Intensity": "SAR-Intensity.tif",
+ }
+ chip_size = {
+ "PAN": (900, 900),
+ "RGBNIR": (450, 450),
+ "PS-RGB": (900, 900),
+ "PS-RGBNIR": (900, 900),
+ "SAR-Intensity": (900, 900),
+ }
+ label_glob = "labels.geojson"
+
+ def __init__(
+ self,
+ root: str = "data",
+ image: str = "PS-RGB",
+ transforms: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
+ download: bool = False,
+ api_key: Optional[str] = None,
+ ) -> None:
+ """Initialize a new SpaceNet 6 Dataset instance.
+
+ Args:
+ root: root directory where dataset can be found
+ image: image selection which must be in ["PAN", "RGBNIR",
+ "PS-RGB", "PS-RGBNIR", "SAR-Intensity"]
+ transforms: a function/transform that takes input sample and its target as
+ entry and returns a transformed version
+ download: if True, download dataset and store it in the root directory.
+ api_key: a RadiantEarth MLHub API key to use for downloading the dataset
+
+ Raises:
+ RuntimeError: if ``download=False`` but dataset is missing
+ """
+ self.root = root
+ self.image = image # For testing
+
+ self.filename = self.imagery[image]
+ self.transforms = transforms
+
+ if download:
+ self.__download(api_key)
+
+ self.files = self._load_files(os.path.join(root, self.dataset_id))
+
+ def __download(self, api_key: Optional[str] = None) -> None:
+ """Download the dataset and extract it.
+
+ Args:
+ api_key: a RadiantEarth MLHub API key to use for downloading the dataset
+
+ Raises:
+ RuntimeError: if download doesn't work correctly or checksums don't match
+ """
+ if os.path.exists(
+ os.path.join(
+ self.root, self.dataset_id, self.collections[0], "collection.json"
+ )
+ ):
+ print("Files already downloaded and verified")
+ return
+
+ download_radiant_mlhub_dataset(self.dataset_id, self.root, api_key)
+
+
class SpaceNet7(SpaceNet):
"""SpaceNet 7: Multi-Temporal Urban Development Challenge.
@@ -1155,13 +1271,6 @@ class SpaceNet7(SpaceNet):
* https://arxiv.org/abs/2102.04420
- .. note::
-
- This dataset requires the following additional library to be installed:
-
- * `radiant-mlhub `_ to download the
- imagery and labels from the Radiant Earth MLHub
-
.. versionadded:: 0.2
"""