dataset.sampling.sampler

sampler

HiRISE-aware geospatial sampler with train/val/test splitting for TorchGeo.

HiRISE satellite passes produce long, narrow, slightly rotated image strips. Sampling patches uniformly from axis-aligned bounding boxes wastes 60–90 % of samples on empty (zero-value) pixels that fall in the corners outside the actual strip.

:class:HiRISEGeoSampler avoids this by pre-computing a grid of patch centres that are confirmed to lie within each strip's polygon footprint, then sampling from that pre-computed set at each epoch. Construction runs once; iteration is O(1) per sample.

Two centre-placement strategies are available via center_mode:

• "simple" (default, backwards-compatible) — the original bbox-grid-then-filter algorithm. Generates an axis-aligned grid over each strip's bounding box and keeps centres whose patch achieves min_overlap with the footprint. Adjacent spacing is controlled by stride.

• "optimal" — the geometric packing algorithm. For each strip the valid centre region (the locus of centres where the patch is guaranteed to meet min_overlap) is computed analytically, then patches are packed inside it using a row/column sweep. Adjacent spacing is controlled by patch_overlap (fractional overlap between neighbouring patches, independent of min_overlap).

The "optimal" mode typically increases coverage by 20–60 % on narrow rotated strips because it places patches right up against the footprint edges. With patch_overlap > 0 it also supports dense augmentation-style sampling without wasted off-strip centres.

Split support ~~~~~~~~~~~~~

The sampler partitions the dataset's stereo pairs (index rows) into train / val / test subsets. Two splitting strategies are provided:

• "geographic" — sorts stereo pairs along a spatial axis (longitude or latitude) and assigns contiguous blocks to each split. This prevents spatial data leakage: nearby strips never appear on both sides of the split.

• "random" — assigns stereo pairs uniformly at random using a deterministic seed.

Both strategies support K-fold cross-validation. When n_folds is set, the data is partitioned into K equally-sized folds. fold_idx selects which fold is used as the test set; the remaining folds are re-split into train and val according to val_fraction.

Split assignments are cached to disk so that:

1. Every process in a distributed training run sees the same split.
2. Re-instantiating the sampler with identical parameters reuses the same assignment without recomputing.
3. The cache key incorporates all split-relevant parameters (dataset root, target, bbox, split method, seed, K, fold, fractions) so that changing any parameter produces a fresh split.

Cache files are written next to the dataset's spatial index cache (under <root>/.cache/) with a filename derived from the configuration hash.

Backwards compatibility. When center_mode="simple" the cache key is computed exactly as in previous versions of this module, so pre-existing cache files are transparently reused. The optimal-mode parameters (center_mode, patch_overlap, packing_phase_steps, valid_region_rays) are included in the cache key only when center_mode="optimal", keeping the two regimes isolated on disk.

HiRISEGeoSampler

HiRISEGeoSampler(dataset: GeoDataset, size: float | tuple[float, float], *, split: Literal['train', 'test', 'val', 'all'] | None = 'train', split_fractions: tuple[float, float, float] = (0.8, 0.1, 0.1), split_method: str = 'geographic', split_axis: str = 'longitude', n_folds: int | None = None, fold_idx: int = 0, seed: int = 42, length: int | None = None, stride: float | tuple[float, float] | None = None, roi=None, toi=None, units: Units = Units.CRS, generator: Generator | None = None, min_overlap: float = 0.5, replacement: bool = False, reuse_cache: bool = True, center_mode: Literal['simple', 'optimal'] = 'optimal', patch_overlap: float = 0.0, packing_phase_steps: int = 20, valid_region_rays: int = 3)

Bases: GeoSampler

Sampler that restricts patches to within HiRISE strip polygon footprints, with built-in train/val/test splitting and K-fold cross-validation.

Construction pre-computes a set of candidate patch centres for every strip in dataset.index. Two strategies are available:

• center_mode="simple" — classical grid-over-bbox followed by an intersection-area filter. Controlled by size and stride.

• center_mode="optimal" — computes each strip's valid centre region analytically (the locus of centres where the patch is guaranteed to meet min_overlap) and packs centres inside it with a row/column sweep. Adjacent spacing is controlled by patch_overlap.

The split is performed at the stereo-pair level — entire strips are assigned to train, val, or test. This prevents spatial data leakage (nearby terrain never appears on both sides of the split).

Parameters:

Name	Type	Description	Default
dataset	GeoDataset	The :class:~MarsHiRISEDTM dataset to sample from.	required
size	float | tuple[float, float]	Patch height and width in CRS units (degrees when units=Units.CRS) or pixels (when units=Units.PIXELS). A single float sets both dimensions equal.	required
split	Literal['train', 'test', 'val', 'all'] | None	Which split to sample from: "train", "val", or "test".	'train'
split_fractions	tuple[float, float, float]	(train, val, test) fractions summing to 1.0.	(0.8, 0.1, 0.1)
split_method	str	"geographic" (spatially contiguous blocks) or "random" (uniformly at random).	'geographic'
split_axis	str	For geographic splits: "longitude" or "latitude".	'longitude'
n_folds	int | None	Number of folds for K-fold cross-validation. None disables K-fold and uses split_fractions directly.	None
fold_idx	int	Which fold to use as the test set (0 to n_folds - 1).	0
seed	int	Random seed for reproducible split assignment.	42
length	int | None	Number of patches to yield per epoch. Defaults to the total number of pre-computed valid centres for this split.	None
stride	float | tuple[float, float] | None	(simple mode only) Centre-to-centre grid spacing in the same units as size. Defaults to size (non-overlapping grid). Ignored when center_mode="optimal".	None
roi		Optional Shapely Polygon to further restrict the spatial domain.	None
toi		Optional :class:pandas.Interval to restrict the temporal domain.	None
units	Units	Whether size and stride are given in CRS units or pixels.	CRS
generator	Generator | None	Optional :class:torch.Generator for reproducible sampling.	None
min_overlap	float	Minimum fraction of patch area that must overlap the strip footprint to be considered valid (default 0.5).	0.5
replacement	bool	Sample with replacement if True.	False
reuse_cache	bool	Reuse cached split assignment if available.	True
center_mode	Literal['simple', 'optimal']	"simple" (default, backwards-compatible) or "optimal" (geometric packing).	'optimal'
patch_overlap	float	(optimal mode only) Fractional overlap between adjacent patches in [0, 1). 0.0 is edge-to-edge, 0.5 means 50% overlap in both axes. Ignored when center_mode="simple".	0.0
packing_phase_steps	int	(optimal mode only) Number of phase offsets to try per sweep direction when packing. Default 20.	20
valid_region_rays	int	(optimal mode only) Number of extra rays per polygon edge when approximating the valid centre region. Default 3.	3

Example::

from hirise_sampler import HiRISEGeoSampler
from torchgeo.samplers import Units

# Legacy behaviour (unchanged; reuses existing caches)
train_sampler = HiRISEGeoSampler(
    dataset, size=0.005, split="train",
    split_fractions=(0.8, 0.1, 0.1),
    seed=42,
)

# Dense optimal packing with 50% patch overlap
train_sampler = HiRISEGeoSampler(
    dataset, size=0.005, split="train",
    split_fractions=(0.8, 0.1, 0.1),
    seed=42,
    center_mode="optimal",
    patch_overlap=0.5,
)

# 5-fold cross-validation, fold 0 as test
train_sampler = HiRISEGeoSampler(
    dataset, size=0.005, split="train",
    n_folds=5, fold_idx=0, seed=42,
)

Source code in src/dataset/sampling/sampler.py

def __init__(
        self,
        dataset: GeoDataset,
        size: float | tuple[float, float],
        *,
        split: Literal["train", "test", "val", "all"] | None = "train",
        split_fractions: tuple[float, float, float] = (0.8, 0.1, 0.1),
        split_method: str = "geographic",
        split_axis: str = "longitude",
        n_folds: int | None = None,
        fold_idx: int = 0,
        seed: int = 42,
        length: int | None = None,
        stride: float | tuple[float, float] | None = None,
        roi=None,
        toi=None,
        units: Units = Units.CRS,
        generator: torch.Generator | None = None,
        min_overlap: float = 0.5,
        replacement: bool = False,
        reuse_cache: bool = True,
        # ── New optimal-mode parameters ──
        center_mode: Literal["simple", "optimal"] = "optimal",
        patch_overlap: float = 0.0,
        packing_phase_steps: int = 20,
        valid_region_rays: int = 3,
) -> None:
    super().__init__(dataset, roi, toi)

    # Use the whole dataset
    if split is None:
        split = "all"

    if split == "all":
        split = "train"
        split_fractions = (1.0, 0.0, 0.0)

    # ── Validate split parameters ──
    if split not in VALID_SPLITS:
        raise ValueError(
            f"Invalid split '{split}'. Must be one of {sorted(VALID_SPLITS)}."
        )
    if split_method not in VALID_SPLIT_METHODS:
        raise ValueError(
            f"Invalid split_method '{split_method}'. "
            f"Must be one of {sorted(VALID_SPLIT_METHODS)}."
        )
    if center_mode not in VALID_CENTER_MODES:
        raise ValueError(
            f"Invalid center_mode '{center_mode}'. "
            f"Must be one of {sorted(VALID_CENTER_MODES)}."
        )
    if not (0.0 <= patch_overlap < 1.0):
        raise ValueError(
            f"patch_overlap must be in [0, 1); got {patch_overlap}"
        )

    train_f, val_f, test_f = split_fractions
    if n_folds is None and abs(train_f + val_f + test_f - 1.0) > 1e-6:
        raise ValueError(
            f"split_fractions must sum to 1.0, got "
            f"{train_f} + {val_f} + {test_f} = {train_f + val_f + test_f}"
        )

    self.split = split
    self.split_fractions = split_fractions
    self.split_method = split_method
    self.split_axis = split_axis
    self.n_folds = n_folds
    self.fold_idx = fold_idx
    self.seed = seed
    self.replacement = replacement
    self.min_overlap = min_overlap

    # Optimal-mode state
    self.center_mode = center_mode
    self.patch_overlap = float(patch_overlap)
    self.packing_phase_steps = int(packing_phase_steps)
    self.valid_region_rays = int(valid_region_rays)

    # ── Resolve size / stride to (height_deg, width_deg) ──
    size_h, size_w = _to_tuple(size)
    if units == Units.PIXELS:
        xres, yres = dataset.res
        size_h *= yres
        size_w *= xres

    if stride is None:
        stride_h = size_h * (1.0 - self.patch_overlap)
        stride_w = size_w * (1.0 - self.patch_overlap)
    else:
        stride_h, stride_w = _to_tuple(stride)
        if units == Units.PIXELS:
            stride_h *= yres
            stride_w *= xres

    self.size = (size_h, size_w)
    self.stride = (stride_h, stride_w)
    self.generator = generator

    # ── Compute or load split assignments ──
    self.cache_hash = ''
    self._assignments = self._get_split_assignments(dataset, reuse_cache)

    # Log split distribution
    split_counts = {}
    for s in VALID_SPLITS:
        split_counts[s] = sum(1 for v in self._assignments.values() if v == s)
    n_total = len(self._assignments)
    logger.info(
        "Split assignment: %d stereo pairs → train=%d, val=%d, test=%d "
        "(method=%s, seed=%d%s)",
        n_total,
        split_counts.get("train", 0),
        split_counts.get("val", 0),
        split_counts.get("test", 0),
        split_method,
        seed,
        f", fold={fold_idx}/{n_folds}" if n_folds else "",
    )

    # ── Pre-compute valid patch centres for this split only ──
    self._centers: list[tuple[float, float, pd.Interval]] = []
    # Per-strip diagnostics, populated during centre building.  Lets
    # downstream visualisation / evaluation code compare the two modes.
    self._per_strip_stats: list[dict[str, Any]] = []

    if self.center_mode == "simple":
        self._build_valid_centers_simple()
    else:
        self._build_valid_centers_optimal()

    n = len(self._centers)
    self.length = length if length is not None else max(1, n)

    if not self.replacement and self.length > n:
        logger.warning(
            "length (%d) > available centres (%d) with replacement=False; "
            "capping to %d. Use replacement=True for oversampling.",
            self.length, n, n,
        )
        self.length = n

    if not self._centers:
        logger.warning(
            "HiRISEGeoSampler: No valid patch centres found for split='%s'. "
            "Check that strip polygons are larger than the requested patch "
            "size (%.6f° × %.6f°) and that the split has data assigned.",
            self.split,
            size_h,
            size_w,
        )

split_summary property

split_summary: dict[str, str | int | dict[str, int] | Any]

Return the number of stereo pairs and patch centres per split.

per_strip_stats property

per_strip_stats: list[dict[str, Any]]

Per-strip diagnostics collected during centre building.

Useful for comparing "simple" vs "optimal" coverage on a per-strip basis. Each entry is a dict containing at minimum pair_idx and n_centers.