torchgeo

import torch from torchgeo.datasets import EuroSAT from torchgeo.transforms import AugmentationSequential, RandomGrayscale from torchgeo.samplers import RandomBatchGeoSampler from torch.utils.data import DataLoader import tempfile import os # Initialize transforms transforms = AugmentationSequential( RandomGrayscale(p=0.5), data_keys=["image"] ) # Use a temporary directory for the dataset to avoid polluting the user's system with tempfile.TemporaryDirectory() as tmpdir: # Initialize EuroSAT dataset (will download if not present) dataset = EuroSAT(root=tmpdir, split="train", transforms=transforms, download=True) # Initialize a sampler to get patches sampler = RandomBatchGeoSampler(dataset, patch_size=(64, 64), batch_size=4, length=10) # Create a DataLoader dataloader = DataLoader(dataset, sampler=sampler, num_workers=0) # Iterate through one batch and print shapes for batch in dataloader: image = batch["image"] label = batch["label"] print(f"Batch image shape: {image.shape}, label shape: {label.shape}") break # Just one batch for quickstart