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 """