Copernicus-Pretrain#

class torchgeo.datasets.CopernicusPretrain(*args, **kwargs)[source]#

Bases: IterableDataset[dict[str, Any]]

Copernicus-Pretrain dataset.

Copernicus-Pretrain is an extension of the SSL4EO-S12 dataset to all major Sentinel missions (S1-S5P). The images are organized into ~310K regional grids (0.25°x0.25°, consistent with ERA5), densely covering the whole land surface and near-land ocean with time series from eight distinct Sentinel modalities.

This dataset class uses WebDataset for efficient data loading in distributed environments, which returns a PyTorch IterableDataset that is compatible with Pytorch DataLoader. Note: it is recommended to further use webdataset.WebLoader (a wrapper on DataLoader) for more features in data loading.

The full dataset has a varying number of modalities, S1/2 local patches, and timestamps for different grids. It also contains metadata including the filenames all images are derived from. For simplicity, the current dataset class provides a minimum example:

  • only use grids with all modalities (220k)

  • sample one local patch for S1 and S2

  • sample one timestamp for each modality

Therefore, each sample contains 8 tensors (S1, S2, S3, S5P NO2/CO/SO2/O3, DEM).

Example:

dataset = CopernicusPretrain(
    urls='data/example-{000000..000009}.tar', shardshuffle=True, resampled=True
)

# Check the first sample
sample = next(iter(dataset))
s1 = sample['s1_grd.pth']
s2 = sample['s2_toa.pth']
s3 = sample['s3_olci.pth']
s5p_co = sample['s5p_co.pth']
s5p_no2 = sample['s5p_no2.pth']
s5p_o3 = sample['s5p_o3.pth']
s5p_so2 = sample['s5p_so2.pth']
dem = sample['dem.pth']

# Create a DataLoader for distributed training on 2 GPUs
dataset = dataset.dataset.batched(10)  # batch size
dataloader = webdataset.WebLoader(dataset, batch_size=None, num_workers=2)
# Unbatch, shuffle, and rebatch to mix samples from different workers
dataloader = dataloader.unbatched().shuffle(100).batched(10)
# A resampled dataset is infinite size, but we can recreate a fixed epoch length
# Total number of samples / (batch size * world size)
number_of_batches = 1000 // (10 * 2)
dataloader = dataloader.with_epoch(number_of_batches)

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

Note

This dataset requires the following additional library to be installed:

Added in version 0.7.

__init__(*args, **kwargs)[source]#

Initialize a new CopernicusPretrain instance.

Parameters:

  • *args (Any) – Arguments passed to WebDataset base class.

  • **kwargs (Any) – Keyword arguments passed to WebDataset base class.

__iter__()[source]#

Iterate over images in the dataset.

Yields:

sample of images

Changed in version 0.10: Removed json metadata.

plot(sample, show_titles=True, suptitle=None)[source]#

Plot a sample from the dataset.

Parameters:

  • sample (dict[str, Any]) – A sample returned by __iter__().

  • show_titles (bool) – Flag indicating whether to show titles above each panel.

  • suptitle (str | None) – Optional string to use as a suptitle.

Returns:

A matplotlib Figure with the rendered sample.

Return type:

Figure

__annotate_func__()[source]#

The type of the None singleton.