add FireRisk dataset

docs/api/datamodules.rst

@@ -60,6 +60,11 @@ FAIR1M
 
 .. autoclass:: FAIR1MDataModule
 
+FireRisk
+^^^^^^^^
+
+.. autoclass:: FireRiskDataModule
+
 GID-15
 ^^^^^^
 

docs/api/datasets.rst

@@ -217,6 +217,11 @@ FAIR1M
 
 .. autoclass:: FAIR1M
 
+FireRisk
+^^^^^^^^
+
+.. autoclass:: FireRisk
+
 Forest Damage
 ^^^^^^^^^^^^^
 

docs/api/non_geo_datasets.csv

@@ -10,6 +10,7 @@ Dataset,Task,Source,# Samples,# Classes,Size (px),Resolution (m),Bands
 `ETCI2021 Flood Detection`_,S,Sentinel-1,"66,810",2,256x256,5--20,SAR
 `EuroSAT`_,C,Sentinel-2,"27,000",10,64x64,10,MSI
 `FAIR1M`_,OD,Gaofen/Google Earth,"15,000",37,"1,024x1,024",0.3--0.8,RGB
+`FireRisk`_,C,NAIP Aerial,"91,872",7,"320x320",1,RGB
 `Forest Damage`_,OD,Drone imagery,"1,543",4,"1,500x1,500",,RGB
 `GID-15`_,S,Gaofen-2,150,15,"6,800x7,200",3,RGB
 `IDTReeS`_,"OD,C",Aerial,591,33,200x200,0.1--1,RGB

tests/conf/fire_risk.yaml

@@ -0,0 +1,15 @@
+experiment:
+  task: "fire_risk"
+  module:
+    loss: "ce"
+    model: "resnet18"
+    learning_rate: 1e-3
+    learning_rate_schedule_patience: 6
+    weights: null
+    in_channels: 3
+    num_classes: 5
+  datamodule:
+    root: "tests/data/fire_risk"
+    download: false
+    batch_size: 1
+    num_workers: 0

tests/data/fire_risk/data.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import hashlib
+import os
+import shutil
+
+import numpy as np
+from PIL import Image
+
+SIZE = 32
+
+np.random.seed(0)
+
+PATHS = [
+    os.path.join(
+        "FireRisk", "train", "High", "27032281_4_-103.430441201095_44.2804260315038.png"
+    ),
+    os.path.join(
+        "FireRisk", "train", "Low", "27032391_2_-103.058289903541_44.3007203324261.png"
+    ),
+    os.path.join(
+        "FireRisk",
+        "train",
+        "Moderate",
+        "27033601_3_-98.95279624632_44.455109470962.png",
+    ),
+    os.path.join(
+        "FireRisk",
+        "train",
+        "Non-burnable",
+        "27033161_6_-100.447787439271_44.4136022778593.png",
+    ),
+    os.path.join(
+        "FireRisk",
+        "train",
+        "Very_High",
+        "27041631_5_-123.547051830273_41.5463004986268.png",
+    ),
+    os.path.join(
+        "FireRisk", "val", "High", "35501951_4_-73.9911660056379_41.2755665931274.png"
+    ),
+    os.path.join(
+        "FireRisk", "val", "Low", "35501621_2_-75.0371666057303_41.4540009148918.png"
+    ),
+    os.path.join(
+        "FireRisk",
+        "val",
+        "Moderate",
+        "35501731_3_-74.6879125510064_41.3954685534897.png",
+    ),
+    os.path.join(
+        "FireRisk",
+        "val",
+        "Non-burnable",
+        "35502061_6_-73.6436892181052_41.2142019946826.png",
+    ),
+    os.path.join(
+        "FireRisk",
+        "val",
+        "Very_High",
+        "35502941_5_-122.968467383602_40.2960022654498.png",
+    ),
+]
+
+
+def create_file(path: str) -> None:
+    Z = np.random.randint(255, size=(SIZE, SIZE, 3), dtype=np.uint8)
+    img = Image.fromarray(Z).convert("RGB")
+    img.save(path)
+
+
+if __name__ == "__main__":
+    directory = "FireRisk"
+    filename = "FireRisk.zip"
+
+    # remove old data
+    if os.path.isdir(directory):
+        shutil.rmtree(directory)
+
+    for path in PATHS:
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        create_file(path)
+
+    # compress data
+    shutil.make_archive(directory, "zip", ".", directory)
+
+    # compute checksum
+    with open(filename, "rb") as f:
+        md5 = hashlib.md5(f.read()).hexdigest()
+        print(f"{filename}: {md5}")

tests/datasets/test_fire_risk.py

@@ -0,0 +1,70 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+import shutil
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import pytest
+import torch
+import torch.nn as nn
+from _pytest.fixtures import SubRequest
+from _pytest.monkeypatch import MonkeyPatch
+
+import torchgeo.datasets.utils
+from torchgeo.datasets import FireRisk
+
+
+def download_url(url: str, root: str, *args: str, **kwargs: str) -> None:
+    shutil.copy(url, root)
+
+
+class TestFireRisk:
+    @pytest.fixture(params=["train", "val"])
+    def dataset(
+        self, monkeypatch: MonkeyPatch, tmp_path: Path, request: SubRequest
+    ) -> FireRisk:
+        monkeypatch.setattr(torchgeo.datasets.fire_risk, "download_url", download_url)
+        url = os.path.join("tests", "data", "fire_risk", "FireRisk.zip")
+        md5 = "db22106d61b10d855234b4a74db921ac"
+        monkeypatch.setattr(FireRisk, "md5", md5)
+        monkeypatch.setattr(FireRisk, "url", url)
+        root = str(tmp_path)
+        split = request.param
+        transforms = nn.Identity()
+        return FireRisk(root, split, transforms, download=True, checksum=True)
+
+    def test_getitem(self, dataset: FireRisk) -> None:
+        x = dataset[0]
+        assert isinstance(x, dict)
+        assert isinstance(x["image"], torch.Tensor)
+        assert isinstance(x["label"], torch.Tensor)
+        assert x["image"].shape[0] == 3
+
+    def test_len(self, dataset: FireRisk) -> None:
+        assert len(dataset) == 5
+
+    def test_already_downloaded(self, dataset: FireRisk, tmp_path: Path) -> None:
+        FireRisk(root=str(tmp_path), download=True)
+
+    def test_already_downloaded_not_extracted(
+        self, dataset: FireRisk, tmp_path: Path
+    ) -> None:
+        shutil.rmtree(os.path.dirname(dataset.root))
+        download_url(dataset.url, root=str(tmp_path))
+        FireRisk(root=str(tmp_path), download=False)
+
+    def test_not_downloaded(self, tmp_path: Path) -> None:
+        with pytest.raises(RuntimeError, match="Dataset not found in"):
+            FireRisk(str(tmp_path))
+
+    def test_plot(self, dataset: FireRisk) -> None:
+        x = dataset[0].copy()
+        dataset.plot(x, suptitle="Test")
+        plt.close()
+        dataset.plot(x, show_titles=False)
+        plt.close()
+        x["prediction"] = x["label"].clone()
+        dataset.plot(x)
+        plt.close()

tests/trainers/test_classification.py

@@ -21,6 +21,7 @@
     BigEarthNetDataModule,
     EuroSAT100DataModule,
     EuroSATDataModule,
+    FireRiskDataModule,
     MisconfigurationException,
     RESISC45DataModule,
     So2SatDataModule,
@@ -77,6 +78,7 @@ class TestClassificationTask:
         [
             ("eurosat", EuroSATDataModule),
             ("eurosat", EuroSAT100DataModule),
+            ("fire_risk", FireRiskDataModule),
             ("resisc45", RESISC45DataModule),
             ("so2sat_all", So2SatDataModule),
             ("so2sat_s1", So2SatDataModule),

torchgeo/datamodules/__init__.py

@@ -11,6 +11,7 @@
 from .etci2021 import ETCI2021DataModule
 from .eurosat import EuroSAT100DataModule, EuroSATDataModule
 from .fair1m import FAIR1MDataModule
+from .fire_risk import FireRiskDataModule
 from .geo import GeoDataModule, NonGeoDataModule
 from .gid15 import GID15DataModule
 from .inria import InriaAerialImageLabelingDataModule
@@ -50,6 +51,7 @@
     "EuroSATDataModule",
     "EuroSAT100DataModule",
     "FAIR1MDataModule",
+    "FireRiskDataModule",
     "GID15DataModule",
     "InriaAerialImageLabelingDataModule",
     "LandCoverAIDataModule",

torchgeo/datamodules/fire_risk.py

@@ -0,0 +1,53 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""FireRisk datamodule."""
+
+from typing import Any
+
+import kornia.augmentation as K
+
+from ..datasets import FireRisk
+from ..transforms import AugmentationSequential
+from .geo import NonGeoDataModule
+
+
+class FireRiskDataModule(NonGeoDataModule):
+    """LightningDataModule implementation for the FireRisk dataset.
+
+    .. versionadded:: 0.5
+    """
+
+    def __init__(
+        self, batch_size: int = 64, num_workers: int = 0, **kwargs: Any
+    ) -> None:
+        """Initialize a new FireRiskDataModule instance.
+
+        Args:
+            batch_size: Size of each mini-batch.
+            num_workers: Number of workers for parallel data loading.
+            **kwargs: Additional keyword arguments passed to
+                :class:`~torchgeo.datasets.FireRisk`.
+        """
+        super().__init__(FireRisk, batch_size, num_workers, **kwargs)
+        self.train_aug = AugmentationSequential(
+            K.Normalize(mean=self.mean, std=self.std),
+            K.RandomRotation(p=0.5, degrees=90),
+            K.RandomHorizontalFlip(p=0.5),
+            K.RandomVerticalFlip(p=0.5),
+            K.RandomSharpness(p=0.5),
+            K.RandomErasing(p=0.1),
+            K.ColorJitter(p=0.5, brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1),
+            data_keys=["image"],
+        )
+
+    def setup(self, stage: str) -> None:
+        """Set up datasets.
+
+        Args:
+            stage: Either 'fit', 'validate', 'test', or 'predict'.
+        """
+        if stage in ["fit"]:
+            self.train_dataset = FireRisk(split="train", **self.kwargs)
+        if stage in ["fit", "validate"]:
+            self.val_dataset = FireRisk(split="val", **self.kwargs)

torchgeo/datasets/__init__.py

@@ -37,6 +37,7 @@
 from .eudem import EUDEM
 from .eurosat import EuroSAT, EuroSAT100
 from .fair1m import FAIR1M
+from .fire_risk import FireRisk
 from .forestdamage import ForestDamage
 from .gbif import GBIF
 from .geo import (
@@ -181,6 +182,7 @@
     "EuroSAT",
     "EuroSAT100",
     "FAIR1M",
+    "FireRisk",
     "ForestDamage",
     "GID15",
     "IDTReeS",