torchgeo_sentinel2_dataset_with_random_geo_sampler_dataloader.py

1import torch
2from torch.utils.data import DataLoader
3from torchgeo.datasets import Sentinel2
4from torchgeo.samplers import RandomGeoSampler
5
6# 1. Initialize the dataset
7# Note: You should specify the root directory where data is stored or will be downloaded
8dataset = Sentinel2(root="data/sentinel2", bands=["B04", "B03", "B02"])
9
10# 2. Initialize a sampler
11# Because geospatial data is often too large to fit in memory, 
12# we sample small patches (e.g., 256x256 pixels) from the extent of the dataset.
13sampler = RandomGeoSampler(dataset, size=256, length=10)
14
15# 3. Initialize the DataLoader
16# The DataLoader combines the dataset and the sampler to yield batches of patches.
17dataloader = DataLoader(dataset, sampler=sampler, collate_fn=dataset.collate_fn)
18
19# 4. Iterate through the data
20for batch in dataloader:
21    image = batch["image"]
22    print(f"Batch image shape: {image.shape}")
23    # Perform training or inference here