Source code for torchgeo.datasets.gbm

# Copyright (c) TorchGeo Contributors. All rights reserved.
# Licensed under the MIT License.

"""GlobalBuildingMap."""

from typing import Any

import matplotlib.pyplot as plt
from matplotlib.figure import Figure

from .geo import RasterDataset


class GlobalBuildingMap(RasterDataset):
    """GlobalBuildingMap dataset.

    The GlobalBuildingMap (GBM) dataset provides the highest resolution and highest
    accuracy building footprint map on a global scale ever created. GBM was generated
    by training and applying modern deep neural networks on nearly 800,000 satellite
    images. The dataset is stored in 5 by 5 degree tiles in geotiff format.

    The GlobalBuildingMap is generated by applying an ensemble of deep neural networks
    on nearly 800,000 satellite images of about 3m resolution. The deep neural networks
    were trained with manually inspected training samples generated from OpenStreetMap.

    If you use this dataset in your research, please cite the following paper:

    * https://arxiv.org/abs/2404.13911

    .. versionadded:: 0.7
    """

    filename_glob = 'GBM_v1_*'
    is_image = False

    def plot(
        self,
        sample: dict[str, Any],
        show_titles: bool = True,
        suptitle: str | None = None,
    ) -> Figure:
        """Plot a sample from the dataset.

        Args:
            sample: A sample returned by :meth:`RasterDataset.__getitem__`.
            show_titles: Flag indicating whether to show titles above each panel.
            suptitle: Optional string to use as a suptitle.

        Returns:
            A matplotlib Figure with the rendered sample.
        """
        ncols = 2 if 'prediction' in sample else 1
        fig, axs = plt.subplots(ncols=ncols, squeeze=False)

        axs[0, 0].imshow(sample['mask'], cmap='gray')
        axs[0, 0].axis('off')
        if show_titles:
            axs[0, 0].set_title('Mask')

        if 'prediction' in sample:
            axs[0, 1].imshow(sample['prediction'], cmap='gray')
            axs[0, 1].axis('off')
            if show_titles:
                axs[0, 1].set_title('Prediction')

        if suptitle is not None:
            plt.suptitle(suptitle)

        return fig