| |
|
|
| """Train S1 segmentation models with optional raw-AEF fusion. |
| |
| This script is the NERSC-ready training entrypoint for the crop-based wetland |
| segmentation workflow. It keeps the current loss, metrics, and schedule shape |
| while supporting two comparable input modes: |
| |
| - S1 only: a 3-band Sentinel-1 baseline for direct DW-target comparisons. |
| - S1 + raw AEF: a 67-band tensor where a learned 1x1 convolution projects the |
| 64 raw AlphaEarth bands before fusion with S1 inside the model. |
| - Training reads fixed 512x512 random windows directly from the source rasters. |
| - Epoch validation uses deterministic 512x512 windows, while final full-scene |
| metrics stream overlapping windows and blend them without loading all 67 |
| input bands for a scene at once. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import copy as copy_module |
| import csv |
| import gc |
| import importlib |
| import json |
| import math |
| import random |
| import shutil |
| import time |
| from collections import Counter |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Dict, Iterable, List, Optional, Sequence, Tuple |
|
|
| import numpy as np |
| import pandas as pd |
| import rasterio as rio |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch.utils.data import Dataset |
| from torchvision.transforms import InterpolationMode |
| from torchvision.transforms import functional as TVF |
|
|
| from fastai.imports import * |
| from fastai.torch_imports import * |
| from fastai.vision.all import * |
|
|
|
|
| DEFAULT_INPUT_DIR = Path("/pscratch/sd/r/rohit9/S1ML/training_mosaics/s1_multiband") |
| DEFAULT_AEF_DIR = Path("/pscratch/sd/r/rohit9/S1ML/training_embeddings/alphaearth_v1_annual_int8") |
| DEFAULT_LABEL_DIR = Path("/pscratch/sd/r/rohit9/S1ML/training_mosaics/dw_binary") |
| DEFAULT_ARTIFACT_DIR = Path( |
| "/pscratch/sd/r/rohit9/S1ML/training_runs/s1aef_bottleneck_resnet34_crop512" |
| ) |
| DEFAULT_HF_CACHE_DIR = Path("/pscratch/sd/r/rohit9/UFO/model_cache/huggingface") |
|
|
| SEGFORMER_MODEL_MAP = { |
| "segformer_b0": "nvidia/segformer-b0-finetuned-ade-512-512", |
| "segformer_b1": "nvidia/segformer-b1-finetuned-ade-512-512", |
| "segformer_b2": "nvidia/segformer-b2-finetuned-ade-512-512", |
| "segformer_b3": "nvidia/segformer-b3-finetuned-ade-512-512", |
| "segformer_b4": "nvidia/segformer-b4-finetuned-ade-512-512", |
| "segformer_b5": "nvidia/segformer-b5-finetuned-cityscapes-1024-1024", |
| } |
|
|
|
|
| @dataclass(frozen=True) |
| class TripletRecord: |
| name: str |
| s1_path: Path |
| aef_path: Optional[Path] |
| label_path: Path |
| sample_id: str = "" |
| group_id: str = "" |
| source_id: str = "" |
| prior_path: Optional[Path] = None |
|
|
|
|
| IGNORE_INDEX = 255 |
|
|
|
|
| def valid_target_mask(targ: torch.Tensor) -> torch.Tensor: |
| """True where Dynamic World observed the pixel (label is 0 or 1).""" |
| return targ != IGNORE_INDEX |
|
|
|
|
| def sanitized_target(targ: torch.Tensor, valid: torch.Tensor) -> torch.Tensor: |
| """Replace ignored labels with 0 so one-hot/indexing stays in range.""" |
| return torch.where(valid, targ, torch.zeros_like(targ)) |
|
|
|
|
| def masked_mean(per_pixel: torch.Tensor, valid: torch.Tensor) -> torch.Tensor: |
| weight = valid.float() |
| return (per_pixel * weight).sum() / weight.sum().clamp(min=1.0) |
|
|
|
|
| def masked_cross_entropy( |
| pred: torch.Tensor, |
| targ: torch.Tensor, |
| valid: torch.Tensor, |
| gamma: Optional[float] = None, |
| ) -> torch.Tensor: |
| """Cross-entropy (optionally focal) averaged over observed pixels only. |
| |
| With no ignored pixels this equals the fastai CrossEntropyLossFlat / |
| FocalLossFlat it replaces, so runs on older labels are unaffected. |
| """ |
| target = sanitized_target(targ.long(), valid) |
| per_pixel = F.cross_entropy(pred.float(), target, reduction="none") |
| if gamma is not None: |
| p_t = torch.exp(-per_pixel) |
| per_pixel = torch.pow(1.0 - p_t, gamma) * per_pixel |
| return masked_mean(per_pixel, valid) |
|
|
|
|
| def masked_dice_loss( |
| pred: torch.Tensor, |
| targ: torch.Tensor, |
| valid: torch.Tensor, |
| axis: int = 1, |
| smooth: float = 1.0, |
| reduction: str = "sum", |
| ) -> torch.Tensor: |
| """Dice over observed pixels only. |
| |
| fastai's DiceLoss one-hots the target, so a raw 255 raises IndexError, and it |
| has no ignore support. Zeroing logits does not work either: softmax(0, 0) is |
| 0.5, so ignored pixels would still contribute. Zero the probabilities and the |
| one-hot target instead - ignored pixels then add nothing to either term. The |
| reduction matches fastai's default ("sum" over batch x class), so with no |
| ignored pixels this is numerically identical to the loss it replaces. |
| """ |
| probability = F.softmax(pred.float(), dim=axis) |
| n_classes = probability.shape[axis] |
| target = sanitized_target(targ.long(), valid) |
| one_hot = F.one_hot(target, n_classes).permute(0, 3, 1, 2).float() |
| weight = valid.unsqueeze(1).float() |
| probability = probability * weight |
| one_hot = one_hot * weight |
| sum_dims = list(range(2, probability.ndim)) |
| intersection = torch.sum(probability * one_hot, dim=sum_dims) |
| union = torch.sum(probability + one_hot, dim=sum_dims) |
| dice_score = (2.0 * intersection + smooth) / (union + smooth) |
| loss = 1.0 - dice_score |
| return loss.sum() if reduction == "sum" else loss.mean() |
|
|
|
|
| def binarize_label_array(data) -> "np.ndarray": |
| """0 = dry, 1 = water, IGNORE_INDEX = not observed. |
| |
| Non-finite values are treated as unobserved rather than dry, so a gap never |
| becomes a training negative. |
| """ |
| values = np.asarray(data) |
| mask = np.full(values.shape, IGNORE_INDEX, dtype=np.int64) |
| observed = values != IGNORE_INDEX |
| if np.issubdtype(values.dtype, np.floating): |
| observed &= np.isfinite(values) |
| mask[observed] = (values[observed] > 0.5).astype(np.int64) |
| return mask |
|
|
|
|
| class CombinedLoss(nn.Module): |
| """Focal + Dice, carried over from the working baseline.""" |
|
|
| def __init__(self, axis=1, smooth=1.0, alpha=1.0, gamma=2.0): |
| super().__init__() |
| self.axis = axis |
| self.alpha = alpha |
| self.smooth = smooth |
| self.gamma = gamma |
|
|
| def forward(self, pred, targ): |
| pred = pred.float() |
| targ = targ.long() |
| valid = valid_target_mask(targ) |
| focal = masked_cross_entropy(pred, targ, valid, gamma=self.gamma) |
| dice = masked_dice_loss(pred, targ, valid, axis=self.axis, smooth=self.smooth) |
| return focal + (self.alpha * dice) |
|
|
| def decodes(self, x): |
| return x.argmax(dim=self.axis) |
|
|
| def activation(self, x): |
| return F.softmax(x, dim=self.axis) |
|
|
|
|
| def _lovasz_grad(gt_sorted: torch.Tensor) -> torch.Tensor: |
| p = len(gt_sorted) |
| gts = gt_sorted.sum() |
| intersection = gts - gt_sorted.float().cumsum(0) |
| union = gts + (1 - gt_sorted).float().cumsum(0) |
| jaccard = 1.0 - intersection / union |
| if p > 1: |
| jaccard[1:p] = jaccard[1:p] - jaccard[0:-1] |
| return jaccard |
|
|
|
|
| def _lovasz_softmax_flat( |
| probas: torch.Tensor, labels: torch.Tensor, classes: str = "present" |
| ) -> torch.Tensor: |
| if probas.numel() == 0: |
| return probas * 0.0 |
| num_classes = probas.size(1) |
| losses = [] |
| class_to_sum = list(range(num_classes)) if classes in ("all", "present") else classes |
| for class_idx in class_to_sum: |
| foreground = (labels == class_idx).float() |
| if classes == "present" and foreground.sum() == 0: |
| continue |
| errors = (foreground - probas[:, class_idx]).abs() |
| errors_sorted, permutation = torch.sort(errors, 0, descending=True) |
| foreground_sorted = foreground[permutation.data] |
| grad = _lovasz_grad(foreground_sorted) |
| losses.append((errors_sorted * grad).sum()) |
| if not losses: |
| return torch.tensor(0.0, device=probas.device, requires_grad=True) |
| return torch.stack(losses).mean() |
|
|
|
|
| class LovaszSoftmaxLoss(nn.Module): |
| """Lovasz-Softmax loss for IoU-oriented segmentation tuning.""" |
|
|
| def __init__(self, axis: int = 1): |
| super().__init__() |
| self.axis = axis |
|
|
| def forward(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor: |
| probas = F.softmax(pred.float(), dim=self.axis) |
| _, num_classes, _, _ = probas.shape |
| probas_flat = probas.permute(0, 2, 3, 1).reshape(-1, num_classes) |
| targ_flat = targ.long().reshape(-1) |
| keep = targ_flat != IGNORE_INDEX |
| probas_flat = probas_flat[keep] |
| targ_flat = targ_flat[keep] |
| return _lovasz_softmax_flat(probas_flat, targ_flat, classes="present") |
|
|
| def decodes(self, x: torch.Tensor) -> torch.Tensor: |
| return x.argmax(dim=self.axis) |
|
|
| def activation(self, x: torch.Tensor) -> torch.Tensor: |
| return F.softmax(x, dim=self.axis) |
|
|
|
|
| class CrossEntropyDiceLoss(nn.Module): |
| """Cross-entropy + Dice baseline for a less recall-heavy alternative.""" |
|
|
| def __init__(self, axis: int = 1, smooth: float = 1.0, alpha: float = 1.0): |
| super().__init__() |
| self.axis = axis |
| self.alpha = alpha |
| self.smooth = smooth |
|
|
| def forward(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor: |
| pred = pred.float() |
| targ = targ.long() |
| valid = valid_target_mask(targ) |
| cross_entropy = masked_cross_entropy(pred, targ, valid) |
| dice = masked_dice_loss(pred, targ, valid, axis=self.axis, smooth=self.smooth) |
| return cross_entropy + (self.alpha * dice) |
|
|
| def decodes(self, x: torch.Tensor) -> torch.Tensor: |
| return x.argmax(dim=self.axis) |
|
|
| def activation(self, x: torch.Tensor) -> torch.Tensor: |
| return F.softmax(x, dim=self.axis) |
|
|
|
|
| class TwoClassLossBase(nn.Module): |
| """Shared FastAI-compatible helpers for experimental water losses.""" |
|
|
| def __init__(self, axis: int = 1): |
| super().__init__() |
| self.axis = axis |
|
|
| def decodes(self, x: torch.Tensor) -> torch.Tensor: |
| return x.argmax(dim=self.axis) |
|
|
| def activation(self, x: torch.Tensor) -> torch.Tensor: |
| return F.softmax(x, dim=self.axis) |
|
|
| def water_probability_and_target( |
| self, pred: torch.Tensor, targ: torch.Tensor |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| """Water probability, water target and the observed-pixel mask. |
| |
| Probability and target are zeroed where the label is ignored, so any |
| sum over them skips unobserved pixels. |
| """ |
| valid = valid_target_mask(targ.long()) |
| weight = valid.float() |
| probability = F.softmax(pred.float(), dim=self.axis)[:, 1] * weight |
| target = (targ.long() == 1).float() * weight |
| return probability, target, valid |
|
|
|
|
| class CrossEntropySoftIoULoss(TwoClassLossBase): |
| """Cross-entropy plus a direct soft Jaccard/IoU objective.""" |
|
|
| def __init__(self, axis: int = 1, smooth: float = 1.0, aux_weight: float = 1.0): |
| super().__init__(axis=axis) |
| self.smooth = smooth |
| self.aux_weight = aux_weight |
|
|
| def forward(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor: |
| probability, target, valid = self.water_probability_and_target(pred, targ) |
| dims = tuple(range(1, probability.ndim)) |
| intersection = (probability * target).sum(dim=dims) |
| union = (probability + target - probability * target).sum(dim=dims) |
| soft_iou_loss = 1.0 - ((intersection + self.smooth) / (union + self.smooth)).mean() |
| return masked_cross_entropy(pred, targ, valid) + self.aux_weight * soft_iou_loss |
|
|
|
|
| class FocalTverskyCELoss(TwoClassLossBase): |
| """CE plus focal Tversky, exposing FP/FN asymmetry for water mapping.""" |
|
|
| def __init__( |
| self, |
| axis: int = 1, |
| alpha: float = 0.3, |
| beta: float = 0.7, |
| gamma: float = 1.33, |
| smooth: float = 1.0, |
| aux_weight: float = 1.0, |
| ): |
| super().__init__(axis=axis) |
| self.alpha = alpha |
| self.beta = beta |
| self.gamma = gamma |
| self.smooth = smooth |
| self.aux_weight = aux_weight |
|
|
| def forward(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor: |
| probability, target, valid = self.water_probability_and_target(pred, targ) |
| dims = tuple(range(1, probability.ndim)) |
| true_positive = (probability * target).sum(dim=dims) |
| false_positive = (probability * (1.0 - target)).sum(dim=dims) |
| false_negative = ((1.0 - probability) * target).sum(dim=dims) |
| tversky = (true_positive + self.smooth) / ( |
| true_positive |
| + self.alpha * false_positive |
| + self.beta * false_negative |
| + self.smooth |
| ) |
| focal_tversky = torch.pow(1.0 - tversky, self.gamma).mean() |
| return masked_cross_entropy(pred, targ, valid) + self.aux_weight * focal_tversky |
|
|
|
|
| class BoundaryAwareCEDiceLoss(TwoClassLossBase): |
| """CE+Dice with an auxiliary soft edge-Dice term for shorelines.""" |
|
|
| def __init__(self, axis: int = 1, smooth: float = 1.0, aux_weight: float = 0.2): |
| super().__init__(axis=axis) |
| self.smooth = smooth |
| self.aux_weight = aux_weight |
| self.dice_loss = DiceLoss(axis, smooth) |
| self.register_buffer( |
| "sobel_x", |
| torch.tensor( |
| [[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]] |
| ).view(1, 1, 3, 3), |
| ) |
| self.register_buffer( |
| "sobel_y", |
| torch.tensor( |
| [[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0]] |
| ).view(1, 1, 3, 3), |
| ) |
|
|
| def edge_strength(self, image: torch.Tensor) -> torch.Tensor: |
| image = image.unsqueeze(1) |
| grad_x = F.conv2d(image, self.sobel_x.to(image), padding=1) |
| grad_y = F.conv2d(image, self.sobel_y.to(image), padding=1) |
| return torch.sqrt(grad_x.square() + grad_y.square() + 1e-6).clamp(0.0, 1.0) |
|
|
| def forward(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor: |
| probability, target, valid = self.water_probability_and_target(pred, targ) |
| edge_weight = valid.float().unsqueeze(1) |
| pred_edge = self.edge_strength(probability) * edge_weight |
| target_edge = self.edge_strength(target) * edge_weight |
| dims = tuple(range(1, pred_edge.ndim)) |
| intersection = (pred_edge * target_edge).sum(dim=dims) |
| denominator = pred_edge.sum(dim=dims) + target_edge.sum(dim=dims) |
| boundary_loss = 1.0 - ((2.0 * intersection + self.smooth) / (denominator + self.smooth)).mean() |
| base = masked_cross_entropy(pred, targ, valid) + masked_dice_loss( |
| pred, targ, valid, axis=self.axis, smooth=self.smooth |
| ) |
| return base + self.aux_weight * boundary_loss |
|
|
|
|
| def _soft_erode(image: torch.Tensor) -> torch.Tensor: |
| vertical = -F.max_pool2d(-image, kernel_size=(3, 1), stride=1, padding=(1, 0)) |
| horizontal = -F.max_pool2d(-image, kernel_size=(1, 3), stride=1, padding=(0, 1)) |
| return torch.minimum(vertical, horizontal) |
|
|
|
|
| def _soft_dilate(image: torch.Tensor) -> torch.Tensor: |
| return F.max_pool2d(image, kernel_size=3, stride=1, padding=1) |
|
|
|
|
| def _soft_skeletonize(image: torch.Tensor, iterations: int) -> torch.Tensor: |
| opened = _soft_dilate(_soft_erode(image)) |
| skeleton = F.relu(image - opened) |
| for _ in range(iterations): |
| image = _soft_erode(image) |
| opened = _soft_dilate(_soft_erode(image)) |
| delta = F.relu(image - opened) |
| skeleton = skeleton + F.relu(delta - skeleton * delta) |
| return skeleton |
|
|
|
|
| class TopologyAwareCEDiceLoss(TwoClassLossBase): |
| """CE+Dice+soft-clDice for connected thin water networks.""" |
|
|
| def __init__( |
| self, |
| axis: int = 1, |
| smooth: float = 1.0, |
| aux_weight: float = 0.2, |
| skeleton_iterations: int = 10, |
| ): |
| super().__init__(axis=axis) |
| self.smooth = smooth |
| self.aux_weight = aux_weight |
| self.skeleton_iterations = skeleton_iterations |
| self.dice_loss = DiceLoss(axis, smooth) |
|
|
| def forward(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor: |
| probability, target, valid = self.water_probability_and_target(pred, targ) |
| probability = probability.unsqueeze(1) |
| target = target.unsqueeze(1) |
| pred_skeleton = _soft_skeletonize(probability, self.skeleton_iterations) |
| target_skeleton = _soft_skeletonize(target, self.skeleton_iterations) |
| dims = tuple(range(1, probability.ndim)) |
| topology_precision = ( |
| (pred_skeleton * target).sum(dim=dims) + self.smooth |
| ) / (pred_skeleton.sum(dim=dims) + self.smooth) |
| topology_sensitivity = ( |
| (target_skeleton * probability).sum(dim=dims) + self.smooth |
| ) / (target_skeleton.sum(dim=dims) + self.smooth) |
| cldice = ( |
| 2.0 * topology_precision * topology_sensitivity |
| / (topology_precision + topology_sensitivity + 1e-6) |
| ).mean() |
| base = masked_cross_entropy(pred, targ, valid) + masked_dice_loss( |
| pred, targ, valid, axis=self.axis, smooth=self.smooth |
| ) |
| return base + self.aux_weight * (1.0 - cldice) |
|
|
|
|
| class SymmetricCrossEntropyDiceLoss(TwoClassLossBase): |
| """Noise-robust symmetric cross-entropy plus Dice for weak DW labels.""" |
|
|
| def __init__(self, axis: int = 1, smooth: float = 1.0, aux_weight: float = 0.1): |
| super().__init__(axis=axis) |
| self.aux_weight = aux_weight |
| self.smooth = smooth |
|
|
| def forward(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor: |
| logits = pred.float() |
| labels = targ.long() |
| valid = valid_target_mask(labels) |
| safe_labels = sanitized_target(labels, valid) |
| probabilities = F.softmax(logits, dim=self.axis).clamp(1e-7, 1.0) |
| one_hot = F.one_hot(safe_labels, num_classes=logits.shape[1]).permute(0, 3, 1, 2).float() |
| reverse_ce_map = -(probabilities * torch.log(one_hot.clamp(1e-4, 1.0))).sum(dim=1) |
| reverse_ce = masked_mean(reverse_ce_map, valid) |
| return ( |
| masked_cross_entropy(logits, labels, valid) |
| + self.aux_weight * reverse_ce |
| + masked_dice_loss(logits, labels, valid, axis=self.axis, smooth=self.smooth) |
| ) |
|
|
|
|
| class MaskedDice(Metric): |
| """fastai Dice restricted to observed pixels. |
| |
| The stock fastai Dice/JaccardCoeff multiply and add raw target values, so a |
| label of 255 corrupts both the intersection and the union. |
| """ |
|
|
| def __init__(self, axis: int = 1): |
| self.axis = axis |
| self.reset() |
|
|
| def reset(self): |
| self.inter = 0.0 |
| self.union = 0.0 |
|
|
| def accumulate(self, learn): |
| pred = learn.pred.argmax(dim=self.axis) |
| targ = learn.y |
| valid = targ != IGNORE_INDEX |
| water_pred = (pred == 1) & valid |
| water_targ = (targ == 1) & valid |
| self.inter += (water_pred & water_targ).float().sum().item() |
| self.union += water_pred.float().sum().item() + water_targ.float().sum().item() |
|
|
| @property |
| def value(self): |
| return 2.0 * self.inter / self.union if self.union > 0 else None |
|
|
| @property |
| def name(self): |
| return "dice" |
|
|
|
|
| class MaskedJaccardCoeff(MaskedDice): |
| """fastai JaccardCoeff restricted to observed pixels.""" |
|
|
| @property |
| def value(self): |
| denominator = self.union - self.inter |
| return self.inter / denominator if denominator > 0 else None |
|
|
| @property |
| def name(self): |
| return "jaccard_coeff" |
|
|
|
|
| class WaterIoU(Metric): |
| """Per-class IoU for the water class (index 1).""" |
|
|
| def __init__(self): |
| self.inter = 0 |
| self.union = 0 |
|
|
| def reset(self): |
| self.inter = 0 |
| self.union = 0 |
|
|
| def accumulate(self, learn): |
| pred = learn.pred.argmax(dim=1) |
| targ = learn.y |
| valid = targ != IGNORE_INDEX |
| water_pred = (pred == 1) & valid |
| water_targ = (targ == 1) & valid |
| self.inter += (water_pred & water_targ).float().sum().item() |
| self.union += (water_pred | water_targ).float().sum().item() |
|
|
| @property |
| def value(self): |
| return self.inter / self.union if self.union > 0 else 0.0 |
|
|
| @property |
| def name(self): |
| return "water_iou" |
|
|
|
|
| class SourceBalancedWaterIoU(Metric): |
| """Mean water IoU across validation sources, giving each source equal weight.""" |
|
|
| def __init__(self, source_ids: Sequence[str]): |
| self.source_ids = list(source_ids) |
| self.reset() |
|
|
| def reset(self): |
| self.cursor = 0 |
| self.intersections: Dict[str, float] = {} |
| self.unions: Dict[str, float] = {} |
|
|
| def accumulate(self, learn): |
| pred = learn.pred.argmax(dim=1) |
| targ = learn.y |
| batch_size = int(pred.shape[0]) |
| if self.cursor + batch_size > len(self.source_ids): |
| raise RuntimeError( |
| "Validation source tracking exceeded the configured source sequence." |
| ) |
| for batch_idx in range(batch_size): |
| source_id = self.source_ids[self.cursor + batch_idx] or "unspecified" |
| valid = targ[batch_idx] != IGNORE_INDEX |
| water_pred = (pred[batch_idx] == 1) & valid |
| water_targ = (targ[batch_idx] == 1) & valid |
| self.intersections[source_id] = self.intersections.get(source_id, 0.0) + ( |
| water_pred & water_targ |
| ).float().sum().item() |
| self.unions[source_id] = self.unions.get(source_id, 0.0) + ( |
| water_pred | water_targ |
| ).float().sum().item() |
| self.cursor += batch_size |
|
|
| @property |
| def value(self): |
| source_ious = [ |
| self.intersections.get(source_id, 0.0) / union |
| for source_id, union in self.unions.items() |
| if union > 0 |
| ] |
| return float(np.mean(source_ious)) if source_ious else 0.0 |
|
|
| @property |
| def name(self): |
| return "source_balanced_water_iou" |
|
|
|
|
| class TargetSourceWaterIoU(Metric): |
| """Pooled water IoU for one validation source used as the checkpoint objective.""" |
|
|
| def __init__(self, source_ids: Sequence[str], target_source_id: str): |
| self.source_ids = list(source_ids) |
| self.target_source_id = target_source_id |
| self.reset() |
|
|
| def reset(self): |
| self.cursor = 0 |
| self.intersection = 0.0 |
| self.union = 0.0 |
|
|
| def accumulate(self, learn): |
| pred = learn.pred.argmax(dim=1) |
| targ = learn.y |
| batch_size = int(pred.shape[0]) |
| if self.cursor + batch_size > len(self.source_ids): |
| raise RuntimeError( |
| "Validation source tracking exceeded the configured source sequence." |
| ) |
| for batch_idx in range(batch_size): |
| source_id = self.source_ids[self.cursor + batch_idx] or "unspecified" |
| if source_id != self.target_source_id: |
| continue |
| valid = targ[batch_idx] != IGNORE_INDEX |
| water_pred = (pred[batch_idx] == 1) & valid |
| water_targ = (targ[batch_idx] == 1) & valid |
| self.intersection += (water_pred & water_targ).float().sum().item() |
| self.union += (water_pred | water_targ).float().sum().item() |
| self.cursor += batch_size |
|
|
| @property |
| def value(self): |
| return self.intersection / self.union if self.union > 0 else 0.0 |
|
|
| @property |
| def name(self): |
| return "target_source_water_iou" |
|
|
|
|
| class NonFiniteLossCallback(Callback): |
| """Stop loudly on NaN/Inf instead of silently cancelling the fit.""" |
|
|
| order = TerminateOnNaNCallback.order |
|
|
| def after_batch(self): |
| if not self.training or self.loss is None: |
| return |
|
|
| loss = self.loss.detach().float() |
| if torch.isfinite(loss).all(): |
| return |
|
|
| lr = None |
| if getattr(self.learn, "opt", None) is not None and len(self.opt.hypers) > 0: |
| lr = self.opt.hypers[-1].get("lr") |
|
|
| debug_summary = { |
| "phase": getattr(self.learn, "training_phase", "unknown"), |
| "epoch": int(getattr(self.learn, "epoch", -1)), |
| "iter": int(getattr(self.learn, "iter", -1)), |
| "loss": float(loss.item()), |
| "lr": None if lr is None else float(lr), |
| } |
| print({"non_finite_loss": debug_summary}) |
| raise RuntimeError("Non-finite loss detected during {phase}.".format(**debug_summary)) |
|
|
|
|
| class EpochTimerCallback(Callback): |
| """Track wall-clock epoch durations for benchmark summaries.""" |
|
|
| def before_fit(self): |
| self.epoch_times = [] |
| self._epoch_start = None |
|
|
| def before_epoch(self): |
| self._epoch_start = time.perf_counter() |
|
|
| def after_epoch(self): |
| if self._epoch_start is None: |
| return |
| self.epoch_times.append(time.perf_counter() - self._epoch_start) |
| self._epoch_start = None |
|
|
|
|
| def load_segformer_components(): |
| try: |
| transformers = importlib.import_module("transformers") |
| except ModuleNotFoundError as error: |
| raise ModuleNotFoundError( |
| "SegFormer support requires the 'transformers' package. " |
| "Use a SegFormer-capable environment or keep --model_family=resnet34_unet." |
| ) from error |
| return transformers.SegformerConfig, transformers.SegformerForSemanticSegmentation |
|
|
|
|
| class AEFBottleneckUNet(nn.Module): |
| """Wrap a U-Net with a learned raw-AEF bottleneck projection.""" |
|
|
| def __init__( |
| self, |
| unet_model: nn.Module, |
| n_s1: int = 3, |
| n_aef: int = 64, |
| n_proj: int = 4, |
| n_prior: int = 0, |
| prior_dropout_p: float = 0.0, |
| aef_dropout_p: float = 0.0, |
| ): |
| super().__init__() |
| self.aef_dropout_p = aef_dropout_p |
| self.n_s1 = n_s1 |
| self.n_aef = n_aef |
| self.n_proj = n_proj |
| self.n_prior = n_prior |
| self.prior_dropout_p = prior_dropout_p |
| self.aef_proj = nn.Conv2d(n_aef, n_proj, kernel_size=1, bias=False) |
| nn.init.xavier_uniform_(self.aef_proj.weight, gain=0.1) |
| self.unet = unet_model |
|
|
| def forward(self, x): |
| s1 = x[:, : self.n_s1] |
| aef = x[:, self.n_s1 : self.n_s1 + self.n_aef] |
| aef_projected = self.aef_proj(aef) |
| if self.training and self.aef_dropout_p > 0: |
| |
| keep = ( |
| torch.rand((aef_projected.shape[0], 1, 1, 1), device=aef_projected.device) |
| >= self.aef_dropout_p |
| ).to(aef_projected.dtype) |
| aef_projected = aef_projected * keep |
| components = [s1, aef_projected] |
| if self.n_prior: |
| prior = x[ |
| :, |
| self.n_s1 + self.n_aef : self.n_s1 + self.n_aef + self.n_prior, |
| ] |
| if self.training and self.prior_dropout_p > 0: |
| keep = ( |
| torch.rand((prior.shape[0], 1, 1, 1), device=prior.device) |
| >= self.prior_dropout_p |
| ).to(prior.dtype) |
| prior = prior * keep |
| components.append(prior) |
| fused = torch.cat(components, dim=1) |
| return self.unet(fused) |
|
|
|
|
| class AEFBottleneckSegFormer(nn.Module): |
| """Wrap a SegFormer with the same learned AEF projection used by the U-Net.""" |
|
|
| def __init__( |
| self, |
| model_name: str, |
| num_labels: int, |
| n_s1: int = 3, |
| n_aef: int = 64, |
| n_proj: int = 4, |
| hf_cache_dir: Optional[Path] = None, |
| pretrained: bool = True, |
| local_files_only: bool = False, |
| ): |
| super().__init__() |
| self.n_s1 = n_s1 |
| self.n_aef = n_aef |
| self.n_proj = n_proj |
| self.aef_proj = nn.Conv2d(n_aef, n_proj, kernel_size=1, bias=False) |
| nn.init.xavier_uniform_(self.aef_proj.weight, gain=0.1) |
|
|
| SegformerConfig, SegformerForSemanticSegmentation = load_segformer_components() |
| cache_dir = str(hf_cache_dir) if hf_cache_dir is not None else None |
| if pretrained: |
| config = SegformerConfig.from_pretrained( |
| model_name, |
| cache_dir=cache_dir, |
| local_files_only=local_files_only, |
| ) |
| config.num_labels = num_labels |
| self.segformer = SegformerForSemanticSegmentation.from_pretrained( |
| model_name, |
| config=config, |
| ignore_mismatched_sizes=True, |
| cache_dir=cache_dir, |
| local_files_only=local_files_only, |
| ) |
| else: |
| config = SegformerConfig.from_pretrained( |
| model_name, |
| cache_dir=cache_dir, |
| local_files_only=local_files_only, |
| ) |
| config.num_labels = num_labels |
| self.segformer = SegformerForSemanticSegmentation(config) |
|
|
| self._adapt_patch_embedding(in_channels=n_s1 + n_proj) |
|
|
| def _adapt_patch_embedding(self, in_channels: int) -> None: |
| old_proj = self.segformer.segformer.encoder.patch_embeddings[0].proj |
| new_proj = nn.Conv2d( |
| in_channels, |
| old_proj.out_channels, |
| kernel_size=old_proj.kernel_size, |
| stride=old_proj.stride, |
| padding=old_proj.padding, |
| bias=old_proj.bias is not None, |
| ) |
| with torch.no_grad(): |
| new_proj.weight.zero_() |
| copied_channels = min(old_proj.in_channels, in_channels) |
| if copied_channels > 0: |
| new_proj.weight[:, :copied_channels] = old_proj.weight[:, :copied_channels] |
| if in_channels > copied_channels: |
| mean_weight = old_proj.weight.mean(dim=1, keepdim=True) |
| new_proj.weight[:, copied_channels:] = mean_weight |
| if old_proj.bias is not None and new_proj.bias is not None: |
| new_proj.bias.copy_(old_proj.bias) |
| self.segformer.segformer.encoder.patch_embeddings[0].proj = new_proj |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| s1 = x[:, : self.n_s1] |
| aef = x[:, self.n_s1 : self.n_s1 + self.n_aef] |
| aef_projected = self.aef_proj(aef) |
| fused = torch.cat([s1, aef_projected], dim=1) |
| logits = self.segformer(pixel_values=fused).logits |
| if logits.shape[-2:] != fused.shape[-2:]: |
| logits = F.interpolate( |
| logits, |
| size=fused.shape[-2:], |
| mode="bilinear", |
| align_corners=False, |
| ) |
| return logits |
|
|
|
|
| class S1AEFCropDataset(Dataset): |
| """Training dataset that samples aligned random crops from triplet tiles.""" |
|
|
| def __init__( |
| self, |
| records: Sequence[TripletRecord], |
| crop_size: int, |
| band_mean: np.ndarray, |
| band_std: np.ndarray, |
| n_s1_bands: int, |
| rotate_max_deg: float, |
| rotate_p: float, |
| scale_jitter_min: float, |
| scale_jitter_max: float, |
| s1_gain_jitter: float, |
| s1_bias_jitter: float, |
| gaussian_noise_std: float, |
| gaussian_noise_p: float, |
| positive_crop_p: float, |
| min_water_fraction: float, |
| positive_crop_attempts: int, |
| ): |
| self.records = list(records) |
| self.crop_size = crop_size |
| self.band_mean = torch.from_numpy(band_mean.astype(np.float32)).view(-1, 1, 1) |
| self.band_std = torch.from_numpy(band_std.astype(np.float32)).view(-1, 1, 1) |
| self.n_s1_bands = n_s1_bands |
| self.rotate_max_deg = rotate_max_deg |
| self.rotate_p = rotate_p |
| self.scale_jitter_min = scale_jitter_min |
| self.scale_jitter_max = scale_jitter_max |
| self.s1_gain_jitter = s1_gain_jitter |
| self.s1_bias_jitter = s1_bias_jitter |
| self.gaussian_noise_std = gaussian_noise_std |
| self.gaussian_noise_p = gaussian_noise_p |
| self.positive_crop_p = positive_crop_p |
| self.min_water_fraction = min_water_fraction |
| self.positive_crop_attempts = positive_crop_attempts |
|
|
| def __len__(self): |
| return len(self.records) |
|
|
| def __getitem__(self, idx): |
| image, mask = load_random_triplet_window( |
| self.records[idx], |
| self.crop_size, |
| seek_positive=random.random() < self.positive_crop_p, |
| min_water_fraction=self.min_water_fraction, |
| attempts=self.positive_crop_attempts, |
| ) |
| image, mask = apply_train_augmentations( |
| image=image, |
| mask=mask, |
| crop_size=self.crop_size, |
| n_s1_bands=self.n_s1_bands, |
| rotate_max_deg=self.rotate_max_deg, |
| rotate_p=self.rotate_p, |
| scale_jitter_min=self.scale_jitter_min, |
| scale_jitter_max=self.scale_jitter_max, |
| s1_gain_jitter=self.s1_gain_jitter, |
| s1_bias_jitter=self.s1_bias_jitter, |
| gaussian_noise_std=self.gaussian_noise_std, |
| gaussian_noise_p=self.gaussian_noise_p, |
| normalized_band_count=int(self.band_mean.shape[0]), |
| ) |
| image = normalize_image(image, self.band_mean, self.band_std) |
| return image, mask.long() |
|
|
|
|
| class S1AEFValidationPatchDataset(Dataset): |
| """Stable validation patches used by fastai during each training epoch.""" |
|
|
| def __init__( |
| self, |
| records: Sequence[TripletRecord], |
| crop_size: int, |
| band_mean: np.ndarray, |
| band_std: np.ndarray, |
| ): |
| self.records = list(records) |
| self.crop_size = crop_size |
| self.band_mean = torch.from_numpy(band_mean.astype(np.float32)).view(-1, 1, 1) |
| self.band_std = torch.from_numpy(band_std.astype(np.float32)).view(-1, 1, 1) |
|
|
| def __len__(self): |
| return len(self.records) |
|
|
| def __getitem__(self, idx): |
| image, mask = load_center_triplet_window(self.records[idx], self.crop_size) |
| image = normalize_image(image, self.band_mean, self.band_std) |
| return image, mask.long() |
|
|
|
|
| def read_tile_index(path: Path, records: Sequence[TripletRecord], split: str) -> List[Tuple[TripletRecord, int, int]]: |
| """full_coverage_v1: tiles (record, top, left) of the given split with at least one observed label pixel.""" |
| table = pd.read_csv(path) |
| table = table[(table["split"] == split) & (table["observed_frac"] > 0)] |
| by_name = {record.name: record for record in records} |
| table = table[table["tile_name"].isin(by_name)] |
| missing = sorted(set(by_name) - set(table["tile_name"])) |
| if missing: |
| raise ValueError("{n} {split} records have no tile in {path}, e.g. {ex}".format(n=len(missing), split=split, path=path, ex=missing[:3])) |
| return [(by_name[name], int(top), int(left)) for name, top, left in zip(table["tile_name"], table["top"], table["left"])] |
|
|
|
|
| class S1AEFTileDataset(Dataset): |
| """full_coverage_v1 training set: every tile of every training chip once per epoch (no random crop).""" |
|
|
| def __init__(self, tiles, crop_size, band_mean, band_std, n_s1_bands, rotate_max_deg, rotate_p, scale_jitter_min, |
| scale_jitter_max, s1_gain_jitter, s1_bias_jitter, gaussian_noise_std, gaussian_noise_p): |
| self.tiles = list(tiles) |
| self.records = sorted({tile[0].name: tile[0] for tile in self.tiles}.values(), key=lambda r: r.name) |
| self.crop_size = crop_size |
| self.band_mean = torch.from_numpy(band_mean.astype(np.float32)).view(-1, 1, 1) |
| self.band_std = torch.from_numpy(band_std.astype(np.float32)).view(-1, 1, 1) |
| self.aug = dict(n_s1_bands=n_s1_bands, rotate_max_deg=rotate_max_deg, rotate_p=rotate_p, scale_jitter_min=scale_jitter_min, |
| scale_jitter_max=scale_jitter_max, s1_gain_jitter=s1_gain_jitter, s1_bias_jitter=s1_bias_jitter, |
| gaussian_noise_std=gaussian_noise_std, gaussian_noise_p=gaussian_noise_p) |
|
|
| def __len__(self): |
| return len(self.tiles) |
|
|
| def __getitem__(self, idx): |
| record, top, left = self.tiles[idx] |
| image, mask = load_triplet_window(record, top=top, left=left, height=self.crop_size, width=self.crop_size) |
| image, mask = apply_train_augmentations(image=image, mask=mask, crop_size=self.crop_size, |
| normalized_band_count=int(self.band_mean.shape[0]), **self.aug) |
| image = normalize_image(image, self.band_mean, self.band_std) |
| return image, mask.long() |
|
|
|
|
| class S1AEFValidationTileDataset(Dataset): |
| """full_coverage_v1 validation: every tile of every validation chip, fixed order. Exposes the whole-chip records |
| (for the windowed end-of-training evaluation) and one source id per tile (for the per-epoch metrics).""" |
|
|
| def __init__(self, tiles, records, crop_size, band_mean, band_std): |
| self.tiles = list(tiles) |
| self.records = list(records) |
| self.item_source_ids = [tile[0].source_id for tile in self.tiles] |
| self.crop_size = crop_size |
| self.band_mean = torch.from_numpy(band_mean.astype(np.float32)).view(-1, 1, 1) |
| self.band_std = torch.from_numpy(band_std.astype(np.float32)).view(-1, 1, 1) |
|
|
| def __len__(self): |
| return len(self.tiles) |
|
|
| def __getitem__(self, idx): |
| record, top, left = self.tiles[idx] |
| image, mask = load_triplet_window(record, top=top, left=left, height=self.crop_size, width=self.crop_size) |
| image = normalize_image(image, self.band_mean, self.band_std) |
| return image, mask.long() |
|
|
|
|
| def add_bool_arg(parser: argparse.ArgumentParser, name: str, default: bool, help_text: str) -> None: |
| parser.add_argument("--{name}".format(name=name), dest=name, action="store_true", help=help_text) |
| parser.add_argument( |
| "--no-{name}".format(name=name), |
| dest=name, |
| action="store_false", |
| help="Disable {name}".format(name=name.replace("_", " ")), |
| ) |
| parser.set_defaults(**{name: default}) |
|
|
|
|
| def using_aef(args: argparse.Namespace) -> bool: |
| return args.n_aef_bands > 0 |
|
|
|
|
| def using_prior(args: argparse.Namespace) -> bool: |
| return bool(getattr(args, "use_permanent_prior", False)) |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser( |
| description="Train S1 segmentation models with optional raw-AEF fusion.", |
| formatter_class=argparse.ArgumentDefaultsHelpFormatter, |
| ) |
| parser.add_argument("--input_dir", type=Path, default=DEFAULT_INPUT_DIR) |
| parser.add_argument("--aef_dir", type=Path, default=DEFAULT_AEF_DIR) |
| parser.add_argument("--label_dir", type=Path, default=DEFAULT_LABEL_DIR) |
| parser.add_argument( |
| "--triplet_manifest", |
| type=Path, |
| default=None, |
| help="CSV with explicit tile_name/name, s1_path, label_path, and optional aef_path rows.", |
| ) |
| parser.add_argument( |
| "--triplet_sample_id", |
| default=None, |
| help="Optional sample_id filter for --triplet_manifest. Required when it contains multiple labels.", |
| ) |
| parser.add_argument( |
| "--fixed_split_manifest", |
| type=Path, |
| default=None, |
| help="CSV with split and tile_name columns. Locks train/validation membership across experiments.", |
| ) |
| parser.add_argument("--artifact_dir", type=Path, default=DEFAULT_ARTIFACT_DIR) |
| parser.add_argument( |
| "--warm_start_run_dir", |
| type=Path, |
| default=None, |
| help="Parent run directory with a compatible best checkpoint and band_stats.npz.", |
| ) |
| parser.add_argument( |
| "--warm_start_checkpoint", |
| type=Path, |
| default=None, |
| help="Explicit checkpoint for strict warm-start loading.", |
| ) |
| parser.add_argument( |
| "--band_stats_path", |
| type=Path, |
| default=None, |
| help="Precomputed band_stats.npz to copy into the new artifact directory.", |
| ) |
| parser.add_argument( |
| "--model_family", |
| choices=["resnet34_unet", "resnet50_unet", "resnet101_unet", "segformer"], |
| default="resnet34_unet", |
| ) |
| parser.add_argument( |
| "--segformer_variant", |
| choices=list(SEGFORMER_MODEL_MAP.keys()), |
| default="segformer_b4", |
| ) |
| parser.add_argument("--hf_cache_dir", type=Path, default=DEFAULT_HF_CACHE_DIR) |
| parser.add_argument("--benchmark_summary_path", type=Path, default=None) |
| parser.add_argument("--mode", choices=["benchmark", "train"], default="train") |
| parser.add_argument("--crop_size", type=int, default=512) |
| parser.add_argument("--tile_overlap", type=int, default=64) |
| parser.add_argument("--batch_size", type=int, default=4) |
| parser.add_argument("--num_workers", type=int, default=8) |
| parser.add_argument( |
| "--disable_pin_memory", |
| action="store_true", |
| help="Disable DataLoader pinned memory for clusters where CUDA pinning is unreliable.", |
| ) |
| parser.add_argument("--device_id", type=int, default=0) |
| parser.add_argument("--seed", type=int, default=42) |
| parser.add_argument( |
| "--split_seed", |
| type=int, |
| default=None, |
| help="Seed for the train/validation split. Defaults to --seed.", |
| ) |
| parser.add_argument( |
| "--train_subset_seed", |
| type=int, |
| default=None, |
| help="Seed for deterministic train-record subsampling. Defaults to --seed.", |
| ) |
| parser.add_argument("--valid_pct", type=float, default=0.2) |
| parser.add_argument( |
| "--max_train_records", |
| type=int, |
| default=None, |
| help="Cap the number of training records after the validation split.", |
| ) |
| parser.add_argument("--stats_sample_size", type=int, default=256) |
| parser.add_argument("--freeze_epochs", type=int, default=2) |
| parser.add_argument("--partial_unfreeze_epochs", type=int, default=2) |
| parser.add_argument("--finetune_epochs", type=int, default=12) |
| parser.add_argument("--finetune_tail_epochs", type=int, default=10) |
| parser.add_argument("--warmup_lr", type=float, default=1e-4) |
| parser.add_argument("--partial_unfreeze_lr", type=float, default=1e-5) |
| parser.add_argument("--unfreeze_lr", type=float, default=3e-6) |
| parser.add_argument("--finetune_tail_lr", type=float, default=1.5e-6) |
| parser.add_argument( |
| "--tuning_strategy", |
| choices=["staged", "head_only", "decoder_aef", "last_encoder_decoder_aef", "full_low_lr"], |
| default="staged", |
| help="Explicit trainable scope for transfer-learning strategy comparisons.", |
| ) |
| parser.add_argument("--strategy_epochs", type=int, default=6) |
| parser.add_argument("--strategy_lr", type=float, default=1e-5) |
| parser.add_argument("--weight_decay", type=float, default=1e-2) |
| parser.add_argument("--grad_clip", type=float, default=1.0) |
| parser.add_argument("--grad_accum_steps", type=int, default=1) |
| parser.add_argument("--optimizer", choices=["ranger", "adamw"], default="ranger") |
| parser.add_argument("--schedule_policy", choices=["one_cycle", "flat_cos"], default="one_cycle") |
| parser.add_argument("--patience", type=int, default=4) |
| parser.add_argument( |
| "--selection_metric", |
| choices=["dice", "water_iou", "source_balanced_water_iou", "target_source_water_iou"], |
| default="water_iou", |
| ) |
| parser.add_argument( |
| "--target_source_id", |
| default="global_postflood", |
| help="Source to oversample when --target_source_fraction is set.", |
| ) |
| parser.add_argument( |
| "--target_source_fraction", |
| type=float, |
| default=None, |
| help="Expected fraction of each training epoch drawn from --target_source_id.", |
| ) |
| parser.add_argument("--min_delta", type=float, default=0.001) |
| parser.add_argument( |
| "--loss", |
| choices=[ |
| "focal_dice", |
| "ce_dice", |
| "lovasz", |
| "ce_soft_iou", |
| "focal_tversky", |
| "ce_dice_boundary", |
| "ce_dice_cldice", |
| "sce_dice", |
| ], |
| default="focal_dice", |
| ) |
| parser.add_argument("--loss_aux_weight", type=float, default=None) |
| parser.add_argument("--tversky_alpha", type=float, default=0.3) |
| parser.add_argument("--tversky_beta", type=float, default=0.7) |
| parser.add_argument("--tversky_gamma", type=float, default=1.33) |
| parser.add_argument("--cldice_iterations", type=int, default=10) |
| parser.add_argument("--rotate_max_deg", type=float, default=5.0) |
| parser.add_argument("--rotate_p", type=float, default=0.5) |
| parser.add_argument("--scale_jitter_min", type=float, default=1.0) |
| parser.add_argument("--scale_jitter_max", type=float, default=1.0) |
| parser.add_argument("--s1_gain_jitter", type=float, default=0.0) |
| parser.add_argument("--s1_bias_jitter", type=float, default=0.0) |
| parser.add_argument("--gaussian_noise_std", type=float, default=0.0) |
| parser.add_argument("--gaussian_noise_p", type=float, default=0.0) |
| parser.add_argument( |
| "--positive_crop_p", |
| type=float, |
| default=0.0, |
| help="Probability of trying to draw a crop with at least --min_water_fraction water.", |
| ) |
| parser.add_argument("--min_water_fraction", type=float, default=0.005) |
| parser.add_argument("--positive_crop_attempts", type=int, default=8) |
| parser.add_argument("--train_tile_index", type=Path, default=None, |
| help="full_coverage_v1: tile index CSV; train on every tile of every chip each epoch.") |
| parser.add_argument("--valid_tile_index", type=Path, default=None, |
| help="full_coverage_v1: tile index CSV; per-epoch validation on every tile of every validation chip.") |
| parser.add_argument("--max_train_batches", type=int, default=None) |
| parser.add_argument("--max_valid_batches", type=int, default=None) |
| parser.add_argument("--probe_batch_sizes", type=int, nargs="+", default=[12, 8, 4, 2]) |
| parser.add_argument("--max_benchmark_batch_size", type=int, default=None) |
| parser.add_argument("--eval_thresholds", type=float, nargs="+", default=None) |
| parser.add_argument("--worst_tile_count", type=int, default=25) |
| parser.add_argument("--n_s1_bands", type=int, default=3) |
| parser.add_argument("--n_aef_bands", type=int, default=64) |
| parser.add_argument("--n_proj_bands", type=int, default=4) |
| parser.add_argument( |
| "--aef_dropout_p", |
| type=float, |
| default=0.0, |
| help="factorial_v1: per-sample probability of zeroing the projected AEF input during training.", |
| ) |
| parser.add_argument( |
| "--prior_dropout_p", |
| type=float, |
| default=0.0, |
| help="Per-sample probability of zeroing the permanent-water prior during training.", |
| ) |
| add_bool_arg( |
| parser, |
| "use_permanent_prior", |
| False, |
| "Append prior_path probability maps from the triplet manifest as a raw input channel.", |
| ) |
| add_bool_arg(parser, "use_fp16", True, "Enable mixed precision on GPU.") |
| add_bool_arg(parser, "pretrained", True, "Use ImageNet-pretrained encoder weights.") |
| add_bool_arg(parser, "hf_local_files_only", False, "Restrict SegFormer loading to the local Hugging Face cache.") |
| add_bool_arg(parser, "force_tiled_validation", False, "Force tiled inference during validation and threshold sweeps.") |
| args = parser.parse_args() |
|
|
| args.codes = ["other", "water"] |
| if args.n_aef_bands < 0: |
| parser.error("--n_aef_bands must be >= 0") |
| if args.n_aef_bands == 0 and args.n_proj_bands != 0: |
| parser.error("--n_proj_bands must be 0 when --n_aef_bands=0") |
| if args.n_aef_bands == 0 and args.model_family == "segformer": |
| parser.error("--model_family=segformer is not yet supported with --n_aef_bands=0") |
| if using_prior(args) and args.model_family == "segformer": |
| parser.error("--use_permanent_prior is currently supported only with ResNet U-Nets") |
| if not 0 <= args.prior_dropout_p < 1: |
| parser.error("--prior_dropout_p must be in [0, 1)") |
| if args.prior_dropout_p > 0 and not using_prior(args): |
| parser.error("--prior_dropout_p requires --use_permanent_prior") |
|
|
| if args.model_family == "segformer": |
| variant_name = args.segformer_variant.replace("segformer_", "") |
| if using_aef(args): |
| args.best_model_name = "s1aef_bottleneck_segformer_{variant}_best".format( |
| variant=variant_name |
| ) |
| args.history_csv = "s1aef_bottleneck_segformer_{variant}_history.csv".format( |
| variant=variant_name |
| ) |
| args.threshold_sweep_csv = "s1aef_bottleneck_segformer_{variant}_threshold_sweep.csv".format( |
| variant=variant_name |
| ) |
| args.eval_summary_json = "s1aef_bottleneck_segformer_{variant}_eval_summary.json".format( |
| variant=variant_name |
| ) |
| args.tile_metrics_csv = "s1aef_bottleneck_segformer_{variant}_tile_metrics.csv".format( |
| variant=variant_name |
| ) |
| args.worst_tiles_csv = "s1aef_bottleneck_segformer_{variant}_worst_tiles.csv".format( |
| variant=variant_name |
| ) |
| else: |
| encoder_name = args.model_family.replace("_unet", "") |
| if using_aef(args): |
| args.best_model_name = "s1aef_bottleneck_{encoder}_best".format(encoder=encoder_name) |
| args.history_csv = "s1aef_bottleneck_{encoder}_history.csv".format(encoder=encoder_name) |
| args.threshold_sweep_csv = "s1aef_bottleneck_{encoder}_threshold_sweep.csv".format(encoder=encoder_name) |
| args.eval_summary_json = "s1aef_bottleneck_{encoder}_eval_summary.json".format(encoder=encoder_name) |
| args.tile_metrics_csv = "s1aef_bottleneck_{encoder}_tile_metrics.csv".format(encoder=encoder_name) |
| args.worst_tiles_csv = "s1aef_bottleneck_{encoder}_worst_tiles.csv".format(encoder=encoder_name) |
| else: |
| args.best_model_name = "s1dw_{encoder}_best".format(encoder=encoder_name) |
| args.history_csv = "s1dw_{encoder}_history.csv".format(encoder=encoder_name) |
| args.threshold_sweep_csv = "s1dw_{encoder}_threshold_sweep.csv".format(encoder=encoder_name) |
| args.eval_summary_json = "s1dw_{encoder}_eval_summary.json".format(encoder=encoder_name) |
| args.tile_metrics_csv = "s1dw_{encoder}_tile_metrics.csv".format(encoder=encoder_name) |
| args.worst_tiles_csv = "s1dw_{encoder}_worst_tiles.csv".format(encoder=encoder_name) |
| if using_prior(args): |
| encoder_name = args.model_family.replace("_unet", "") |
| args.best_model_name = "s1aef_prior_{encoder}_best".format(encoder=encoder_name) |
| args.history_csv = "s1aef_prior_{encoder}_history.csv".format(encoder=encoder_name) |
| args.threshold_sweep_csv = "s1aef_prior_{encoder}_threshold_sweep.csv".format(encoder=encoder_name) |
| args.eval_summary_json = "s1aef_prior_{encoder}_eval_summary.json".format(encoder=encoder_name) |
| args.tile_metrics_csv = "s1aef_prior_{encoder}_tile_metrics.csv".format(encoder=encoder_name) |
| args.worst_tiles_csv = "s1aef_prior_{encoder}_worst_tiles.csv".format(encoder=encoder_name) |
| args.benchmark_summary_json = "benchmark_summary.json" |
| if args.eval_thresholds is None: |
| coarse_low = np.arange(0.30, 0.55, 0.05) |
| focused_mid = np.arange(0.55, 0.71, 0.01) |
| coarse_high = np.arange(0.75, 0.81, 0.05) |
| args.eval_thresholds = sorted( |
| {round(float(v), 2) for v in np.concatenate([coarse_low, focused_mid, coarse_high])} |
| ) |
| args.use_tta_eval = True |
| args.tta_eval_mode_count = 4 |
| args.input_dir = args.input_dir.expanduser().resolve() |
| args.aef_dir = args.aef_dir.expanduser().resolve() |
| args.label_dir = args.label_dir.expanduser().resolve() |
| if args.triplet_manifest is not None: |
| args.triplet_manifest = args.triplet_manifest.expanduser().resolve() |
| if args.fixed_split_manifest is not None: |
| args.fixed_split_manifest = args.fixed_split_manifest.expanduser().resolve() |
| args.artifact_dir = args.artifact_dir.expanduser().resolve() |
| if args.warm_start_run_dir is not None: |
| args.warm_start_run_dir = args.warm_start_run_dir.expanduser().resolve() |
| if args.warm_start_checkpoint is not None: |
| args.warm_start_checkpoint = args.warm_start_checkpoint.expanduser().resolve() |
| if args.band_stats_path is not None: |
| args.band_stats_path = args.band_stats_path.expanduser().resolve() |
| args.hf_cache_dir = args.hf_cache_dir.expanduser().resolve() |
| if args.benchmark_summary_path is not None: |
| args.benchmark_summary_path = args.benchmark_summary_path.expanduser().resolve() |
| if args.crop_size < 1: |
| parser.error("--crop_size must be >= 1") |
| if args.tile_overlap < 0 or args.tile_overlap >= args.crop_size: |
| parser.error("--tile_overlap must be >= 0 and < --crop_size") |
| if args.scale_jitter_min <= 0 or args.scale_jitter_max <= 0: |
| parser.error("--scale_jitter_min and --scale_jitter_max must be > 0") |
| if args.scale_jitter_min > args.scale_jitter_max: |
| parser.error("--scale_jitter_min must be <= --scale_jitter_max") |
| if args.batch_size < 1: |
| parser.error("--batch_size must be >= 1") |
| if args.grad_accum_steps < 1: |
| parser.error("--grad_accum_steps must be >= 1") |
| if args.target_source_fraction is not None and not 0 < args.target_source_fraction <= 1: |
| parser.error("--target_source_fraction must be greater than 0 and at most 1") |
| if not 0 <= args.positive_crop_p <= 1: |
| parser.error("--positive_crop_p must be in [0, 1]") |
| if not 0 <= args.min_water_fraction <= 1: |
| parser.error("--min_water_fraction must be in [0, 1]") |
| if args.positive_crop_attempts < 1: |
| parser.error("--positive_crop_attempts must be >= 1") |
| if args.loss_aux_weight is not None and args.loss_aux_weight < 0: |
| parser.error("--loss_aux_weight must be non-negative") |
| if args.tversky_alpha < 0 or args.tversky_beta < 0: |
| parser.error("--tversky_alpha and --tversky_beta must be non-negative") |
| if args.tversky_gamma <= 0: |
| parser.error("--tversky_gamma must be > 0") |
| if args.cldice_iterations < 1: |
| parser.error("--cldice_iterations must be >= 1") |
| if args.strategy_epochs < 1: |
| parser.error("--strategy_epochs must be >= 1") |
| if args.strategy_lr <= 0: |
| parser.error("--strategy_lr must be > 0") |
| if args.max_benchmark_batch_size is not None and args.max_benchmark_batch_size < 1: |
| parser.error("--max_benchmark_batch_size must be >= 1") |
| if args.worst_tile_count < 1: |
| parser.error("--worst_tile_count must be >= 1") |
| if args.triplet_sample_id and args.triplet_manifest is None: |
| parser.error("--triplet_sample_id requires --triplet_manifest") |
| return args |
|
|
|
|
| def seed_everything(seed: int) -> None: |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(seed) |
|
|
|
|
| def seed_worker(worker_id: Optional[int] = None) -> None: |
| worker_seed = torch.initial_seed() % (2**32) |
| random.seed(worker_seed) |
| np.random.seed(worker_seed) |
|
|
|
|
| def setup_environment(args: argparse.Namespace) -> torch.device: |
| seed_everything(args.seed) |
| args.artifact_dir.mkdir(parents=True, exist_ok=True) |
| (args.artifact_dir / "models").mkdir(parents=True, exist_ok=True) |
|
|
| if torch.cuda.is_available(): |
| device_count = torch.cuda.device_count() |
| actual_device = args.device_id if args.device_id < device_count else 0 |
| torch.cuda.set_device(actual_device) |
| device = torch.device("cuda:{idx}".format(idx=actual_device)) |
| torch.backends.cudnn.benchmark = True |
| print( |
| { |
| "device": str(device), |
| "device_name": torch.cuda.get_device_name(actual_device), |
| "device_count": device_count, |
| } |
| ) |
| else: |
| device = torch.device("cpu") |
| print({"device": "cpu"}) |
|
|
| n_input = args.n_s1_bands + args.n_aef_bands |
| n_unet_input = args.n_s1_bands + (args.n_proj_bands if using_aef(args) else 0) |
| if using_aef(args): |
| dataloader_bands = "{s1} S1 + {aef} AEF = {n}".format( |
| s1=args.n_s1_bands, aef=args.n_aef_bands, n=n_input |
| ) |
| model_input_bands = "{s1} S1 + {proj} projected = {n}".format( |
| s1=args.n_s1_bands, proj=args.n_proj_bands, n=n_unet_input |
| ) |
| else: |
| dataloader_bands = "{s1} S1 = {n}".format(s1=args.n_s1_bands, n=n_input) |
| model_input_bands = "{s1} S1 = {n}".format(s1=args.n_s1_bands, n=n_unet_input) |
| print( |
| { |
| "artifact_dir": str(args.artifact_dir), |
| "mode": args.mode, |
| "model_family": args.model_family, |
| "segformer_variant": args.segformer_variant if args.model_family == "segformer" else None, |
| "dataloader_bands": dataloader_bands, |
| "model_input_bands": model_input_bands, |
| "crop_size": args.crop_size, |
| "batch_size": args.batch_size, |
| "num_workers": args.num_workers, |
| "loss": args.loss, |
| "tuning_strategy": args.tuning_strategy, |
| "grad_accum_steps": args.grad_accum_steps, |
| "selection_metric": args.selection_metric, |
| "target_source_id": args.target_source_id, |
| "target_source_fraction": args.target_source_fraction, |
| "eval_thresholds": args.eval_thresholds, |
| "force_tiled_validation": args.force_tiled_validation, |
| "use_fp16": args.use_fp16, |
| "pretrained": args.pretrained, |
| } |
| ) |
| return device |
|
|
|
|
| def build_aef_index(aef_dir: Path) -> Dict[str, Path]: |
| if not aef_dir.exists(): |
| raise FileNotFoundError("AEF directory does not exist: {path}".format(path=aef_dir)) |
| index: Dict[str, Path] = {} |
| duplicates: List[str] = [] |
| for path in sorted(aef_dir.rglob("*.tif")): |
| rel_parts = path.relative_to(aef_dir).parts |
| if any(part.startswith(".") for part in rel_parts): |
| continue |
| if path.name in index: |
| duplicates.append(path.name) |
| else: |
| index[path.name] = path |
| if duplicates: |
| raise RuntimeError( |
| "Found duplicate AEF filenames across recursive folders. Examples: {examples}".format( |
| examples=duplicates[:5] |
| ) |
| ) |
| if not index: |
| raise FileNotFoundError("No AEF GeoTIFFs found under {path}".format(path=aef_dir)) |
| return index |
|
|
|
|
| def collect_triplets( |
| input_dir: Path, |
| aef_dir: Optional[Path], |
| label_dir: Path, |
| ) -> List[TripletRecord]: |
| if not input_dir.exists(): |
| raise FileNotFoundError("Input directory does not exist: {path}".format(path=input_dir)) |
| if not label_dir.exists(): |
| raise FileNotFoundError("Label directory does not exist: {path}".format(path=label_dir)) |
|
|
| use_aef = aef_dir is not None |
| aef_index = build_aef_index(aef_dir) if use_aef else {} |
| |
| |
| |
| label_index = {path.name: path for path in label_dir.glob("*.tif")} |
| s1_paths = sorted(input_dir.glob("*.tif")) |
| triplets: List[TripletRecord] = [] |
| missing_aef: List[str] = [] |
| missing_labels: List[str] = [] |
|
|
| for s1_path in s1_paths: |
| aef_path = aef_index.get(s1_path.name) if use_aef else None |
| label_path = label_index.get(s1_path.name) |
| if use_aef and aef_path is None: |
| missing_aef.append(s1_path.name) |
| continue |
| if label_path is None: |
| missing_labels.append(s1_path.name) |
| continue |
| triplets.append( |
| TripletRecord( |
| name=s1_path.name, |
| s1_path=s1_path, |
| aef_path=aef_path, |
| label_path=label_path, |
| group_id=s1_path.name, |
| ) |
| ) |
|
|
| if not triplets: |
| raise FileNotFoundError( |
| "No aligned S1/AEF/label triplets found. Missing AEF: {ma}, missing labels: {ml}".format( |
| ma=len(missing_aef), ml=len(missing_labels) |
| ) |
| ) |
|
|
| print( |
| { |
| "triplet_count": len(triplets), |
| "missing_aef_count": len(missing_aef) if use_aef else 0, |
| "missing_label_count": len(missing_labels), |
| "missing_aef_examples": missing_aef[:3] if use_aef else [], |
| "missing_label_examples": missing_labels[:3], |
| } |
| ) |
| return triplets |
|
|
|
|
| def collect_triplets_from_manifest( |
| manifest_path: Path, |
| sample_id: Optional[str], |
| expect_aef: bool, |
| expect_prior: bool = False, |
| ) -> List[TripletRecord]: |
| if not manifest_path.exists(): |
| raise FileNotFoundError("Triplet manifest does not exist: {path}".format(path=manifest_path)) |
| with manifest_path.open(newline="") as handle: |
| rows = list(csv.DictReader(handle)) |
| required = {"s1_path", "label_path"} |
| if expect_aef: |
| required.add("aef_path") |
| if expect_prior: |
| required.add("prior_path") |
| missing = required.difference(rows[0].keys() if rows else []) |
| if missing: |
| raise ValueError("Triplet manifest is missing columns: {cols}".format(cols=", ".join(sorted(missing)))) |
|
|
| manifest_sample_ids = sorted({row.get("sample_id", "") for row in rows if row.get("sample_id", "")}) |
| if sample_id is None and len(manifest_sample_ids) > 1: |
| raise ValueError( |
| "Triplet manifest has multiple sample_id values; pass --triplet_sample_id. Examples: {ids}".format( |
| ids=", ".join(manifest_sample_ids[:5]) |
| ) |
| ) |
| selected_rows = [row for row in rows if sample_id is None or row.get("sample_id", "") == sample_id] |
| if not selected_rows: |
| raise ValueError("Triplet manifest has no rows for sample_id {sample}".format(sample=sample_id or "<all>")) |
|
|
| records: List[TripletRecord] = [] |
| missing_paths: List[str] = [] |
| for row in selected_rows: |
| name = row.get("tile_name") or row.get("name") or Path(row["s1_path"]).name |
| s1_path = Path(row["s1_path"]).expanduser() |
| label_path = Path(row["label_path"]).expanduser() |
| aef_value = row.get("aef_path", "") |
| aef_path = Path(aef_value).expanduser() if expect_aef and aef_value else None |
| prior_value = row.get("prior_path", "") |
| prior_path = Path(prior_value).expanduser() if expect_prior and prior_value else None |
| path_checks = [("s1", s1_path), ("label", label_path)] |
| if expect_aef: |
| if aef_path is None: |
| missing_paths.append("{name}: empty aef_path".format(name=name)) |
| continue |
| path_checks.append(("aef", aef_path)) |
| if expect_prior: |
| if prior_path is None: |
| missing_paths.append("{name}: empty prior_path".format(name=name)) |
| continue |
| path_checks.append(("prior", prior_path)) |
| missing_for_row = ["{kind}={path}".format(kind=kind, path=path) for kind, path in path_checks if not path.exists()] |
| if missing_for_row: |
| missing_paths.append("{name}: {paths}".format(name=name, paths=", ".join(missing_for_row))) |
| continue |
| records.append( |
| TripletRecord( |
| name=name, |
| s1_path=s1_path.resolve(), |
| aef_path=aef_path.resolve() if aef_path is not None else None, |
| label_path=label_path.resolve(), |
| sample_id=row.get("sample_id", sample_id or ""), |
| group_id=row.get("grid_id") or row.get("group_id") or name, |
| source_id=row.get("source_id", ""), |
| prior_path=prior_path.resolve() if prior_path is not None else None, |
| ) |
| ) |
|
|
| if missing_paths: |
| raise FileNotFoundError( |
| "Triplet manifest has missing or incomplete paths for {n} row(s). Examples: {examples}".format( |
| n=len(missing_paths), |
| examples="; ".join(missing_paths[:5]), |
| ) |
| ) |
| if not records: |
| raise FileNotFoundError("Triplet manifest did not produce any usable records: {}".format(manifest_path)) |
| print( |
| { |
| "triplet_manifest": str(manifest_path), |
| "triplet_sample_id": sample_id or (manifest_sample_ids[0] if len(manifest_sample_ids) == 1 else None), |
| "triplet_count": len(records), |
| } |
| ) |
| return sorted(records, key=lambda record: record.name) |
|
|
|
|
| def split_grouped_triplets( |
| records: Sequence[TripletRecord], |
| valid_pct: float, |
| seed: int, |
| ) -> Tuple[List[TripletRecord], List[TripletRecord]]: |
| if not records: |
| raise ValueError("Cannot split an empty record list.") |
| grouped: Dict[str, List[TripletRecord]] = {} |
| for record in records: |
| grouped.setdefault(record.group_id or record.name, []).append(record) |
| if len(grouped) < 2: |
| raise ValueError("At least two independent groups are required for a train/validation split.") |
|
|
| rng = random.Random(seed) |
| shuffled_groups = list(grouped.items()) |
| rng.shuffle(shuffled_groups) |
| target_valid_count = max(1, int(round(len(records) * valid_pct))) |
| valid_group_ids = set() |
| valid_count = 0 |
| for group_id, group_records in shuffled_groups: |
| if valid_count >= target_valid_count: |
| break |
| valid_group_ids.add(group_id) |
| valid_count += len(group_records) |
|
|
| valid_records = sorted( |
| [record for record in records if (record.group_id or record.name) in valid_group_ids], |
| key=lambda r: r.name, |
| ) |
| train_records = sorted( |
| [record for record in records if (record.group_id or record.name) not in valid_group_ids], |
| key=lambda r: r.name, |
| ) |
| if not train_records: |
| raise ValueError("Train split is empty; reduce --valid_pct.") |
| train_group_ids = {record.group_id or record.name for record in train_records} |
| if train_group_ids.intersection(valid_group_ids): |
| raise RuntimeError("Group leakage detected between train and validation splits.") |
| return train_records, valid_records |
|
|
|
|
| def split_triplets(records: Sequence[TripletRecord], valid_pct: float, seed: int) -> Tuple[List[TripletRecord], List[TripletRecord]]: |
| source_ids = sorted({record.source_id for record in records if record.source_id}) |
| if len(source_ids) <= 1: |
| return split_grouped_triplets(records, valid_pct=valid_pct, seed=seed) |
|
|
| train_records: List[TripletRecord] = [] |
| valid_records: List[TripletRecord] = [] |
| for source_id in source_ids: |
| source_records = [record for record in records if record.source_id == source_id] |
| source_train, source_valid = split_grouped_triplets( |
| source_records, |
| valid_pct=valid_pct, |
| seed=seed, |
| ) |
| train_records.extend(source_train) |
| valid_records.extend(source_valid) |
|
|
| assigned = len(train_records) + len(valid_records) |
| if assigned != len(records): |
| raise RuntimeError( |
| "Source-stratified split assigned {assigned} of {total} records.".format( |
| assigned=assigned, |
| total=len(records), |
| ) |
| ) |
| return sorted(train_records, key=lambda r: r.name), sorted(valid_records, key=lambda r: r.name) |
|
|
|
|
| def split_triplets_from_manifest( |
| records: Sequence[TripletRecord], manifest_path: Path |
| ) -> Tuple[List[TripletRecord], List[TripletRecord]]: |
| """Apply an immutable train/validation membership manifest by tile name.""" |
| if not manifest_path.exists(): |
| raise FileNotFoundError("Fixed split manifest does not exist: {path}".format(path=manifest_path)) |
| with manifest_path.open(newline="") as handle: |
| rows = list(csv.DictReader(handle)) |
| required = {"split", "tile_name"} |
| missing_columns = required.difference(rows[0].keys() if rows else []) |
| if missing_columns: |
| raise ValueError( |
| "Fixed split manifest is missing columns: {columns}".format( |
| columns=", ".join(sorted(missing_columns)) |
| ) |
| ) |
|
|
| assignments: Dict[str, str] = {} |
| for row in rows: |
| name = row["tile_name"].strip() |
| split = row["split"].strip().lower() |
| if split not in {"train", "valid", "validation"}: |
| raise ValueError("Unknown split {split!r} for {name}".format(split=split, name=name)) |
| split = "valid" if split == "validation" else split |
| if name in assignments: |
| raise ValueError("Duplicate tile in fixed split manifest: {name}".format(name=name)) |
| assignments[name] = split |
|
|
| records_by_name = {record.name: record for record in records} |
| if len(records_by_name) != len(records): |
| raise ValueError("Training records contain duplicate tile names; fixed split is ambiguous.") |
| missing_records = sorted(set(assignments).difference(records_by_name)) |
| unassigned_records = sorted(set(records_by_name).difference(assignments)) |
| if missing_records or unassigned_records: |
| raise ValueError( |
| "Fixed split population mismatch: missing_records={missing}, unassigned_records={unassigned}, " |
| "missing_examples={missing_examples}, unassigned_examples={unassigned_examples}".format( |
| missing=len(missing_records), |
| unassigned=len(unassigned_records), |
| missing_examples=missing_records[:5], |
| unassigned_examples=unassigned_records[:5], |
| ) |
| ) |
|
|
| train_records = sorted( |
| [records_by_name[name] for name, split in assignments.items() if split == "train"], |
| key=lambda record: record.name, |
| ) |
| valid_records = sorted( |
| [records_by_name[name] for name, split in assignments.items() if split == "valid"], |
| key=lambda record: record.name, |
| ) |
| if not train_records or not valid_records: |
| raise ValueError("Fixed split manifest must contain at least one train and one valid tile.") |
| print( |
| { |
| "fixed_split_manifest": str(manifest_path), |
| "fixed_train_tiles": len(train_records), |
| "fixed_valid_tiles": len(valid_records), |
| } |
| ) |
| return train_records, valid_records |
|
|
|
|
| def limit_train_records( |
| records: Sequence[TripletRecord], |
| max_records: Optional[int], |
| subset_seed: int, |
| ) -> List[TripletRecord]: |
| selected = list(records) |
| if max_records is None: |
| return selected |
| if max_records <= 0: |
| raise ValueError("--max_train_records must be positive when set.") |
| rng = random.Random(subset_seed) |
| shuffled = list(selected) |
| rng.shuffle(shuffled) |
| return sorted(shuffled[: min(max_records, len(shuffled))], key=lambda r: r.name) |
|
|
|
|
| def write_split_manifest( |
| artifact_dir: Path, |
| train_records: Sequence[TripletRecord], |
| valid_records: Sequence[TripletRecord], |
| ) -> Path: |
| split_path = artifact_dir / "triplet_split_manifest.csv" |
| fields = ( |
| "split", |
| "sample_id", |
| "source_id", |
| "group_id", |
| "name", |
| "s1_path", |
| "aef_path", |
| "label_path", |
| "prior_path", |
| ) |
| rows = [("train", record) for record in train_records] + [("valid", record) for record in valid_records] |
| with split_path.open("w", newline="") as handle: |
| writer = csv.DictWriter(handle, fieldnames=fields) |
| writer.writeheader() |
| for split, record in rows: |
| writer.writerow( |
| { |
| "split": split, |
| "sample_id": record.sample_id, |
| "source_id": record.source_id, |
| "group_id": record.group_id or record.name, |
| "name": record.name, |
| "s1_path": str(record.s1_path), |
| "aef_path": "" if record.aef_path is None else str(record.aef_path), |
| "label_path": str(record.label_path), |
| "prior_path": "" if record.prior_path is None else str(record.prior_path), |
| } |
| ) |
| print( |
| { |
| "triplet_split_manifest": str(split_path), |
| "train_tiles": len(train_records), |
| "valid_tiles": len(valid_records), |
| "train_groups": len({record.group_id or record.name for record in train_records}), |
| "valid_groups": len({record.group_id or record.name for record in valid_records}), |
| } |
| ) |
| return split_path |
|
|
|
|
| def subset_records(records: Sequence[TripletRecord], max_batches: Optional[int], batch_size: int) -> List[TripletRecord]: |
| selected = list(records) |
| if max_batches is None: |
| return selected |
| max_items = max(batch_size * max_batches, batch_size) |
| max_items = min(len(selected), max_items) |
| records_by_source: Dict[str, List[TripletRecord]] = {} |
| for record in selected: |
| records_by_source.setdefault(record.source_id or "unspecified", []).append(record) |
| if len(records_by_source) <= 1: |
| return selected[:max_items] |
|
|
| |
| |
| source_ids = sorted(records_by_source) |
| source_positions = {source_id: 0 for source_id in source_ids} |
| source_stratified: List[TripletRecord] = [] |
| while len(source_stratified) < max_items: |
| added_record = False |
| for source_id in source_ids: |
| source_records = records_by_source[source_id] |
| source_position = source_positions[source_id] |
| if source_position >= len(source_records): |
| continue |
| source_stratified.append(source_records[source_position]) |
| source_positions[source_id] += 1 |
| added_record = True |
| if len(source_stratified) >= max_items: |
| break |
| if not added_record: |
| break |
| return source_stratified |
|
|
|
|
| def open_geotiff_float(path: Path) -> torch.Tensor: |
| with rio.open(str(path)) as src: |
| data = src.read().astype(np.float32) |
| data = np.nan_to_num(data, nan=0.0, posinf=0.0, neginf=0.0) |
| return torch.from_numpy(data) |
|
|
|
|
| def open_mask_binary(path: Path) -> torch.Tensor: |
| """0 = dry, 1 = water, IGNORE_INDEX = Dynamic World did not observe.""" |
| with rio.open(str(path)) as src: |
| data = src.read(1).astype(np.float32) |
| return torch.from_numpy(binarize_label_array(data)) |
|
|
|
|
| def open_prior_probability(path: Path) -> torch.Tensor: |
| with rio.open(str(path)) as src: |
| data = src.read(1).astype(np.float32) |
| nodata = src.nodata |
| invalid = ~np.isfinite(data) |
| if nodata is not None: |
| invalid |= data == nodata |
| data = np.where(invalid, 0.0, data) |
| data = np.clip(data, 0.0, 1.0) |
| return torch.from_numpy(data).unsqueeze(0) |
|
|
|
|
| def triplet_shape(record: TripletRecord) -> Tuple[int, int]: |
| with rio.open(str(record.s1_path)) as s1_src: |
| height, width = s1_src.height, s1_src.width |
| with rio.open(str(record.label_path)) as label_src: |
| label_shape = (label_src.height, label_src.width) |
| if label_shape != (height, width): |
| raise ValueError( |
| "Triplet {name} has mismatched S1/label shapes: s1={s1}, label={label}".format( |
| name=record.name, |
| s1=(height, width), |
| label=label_shape, |
| ) |
| ) |
| if record.aef_path is not None: |
| with rio.open(str(record.aef_path)) as aef_src: |
| aef_shape = (aef_src.height, aef_src.width) |
| if aef_shape != (height, width): |
| raise ValueError( |
| "Triplet {name} has mismatched S1/AEF shapes: s1={s1}, aef={aef}".format( |
| name=record.name, |
| s1=(height, width), |
| aef=aef_shape, |
| ) |
| ) |
| if record.prior_path is not None: |
| with rio.open(str(record.prior_path)) as prior_src: |
| prior_shape = (prior_src.height, prior_src.width) |
| prior_count = prior_src.count |
| if prior_count != 1 or prior_shape != (height, width): |
| raise ValueError( |
| "Triplet {name} has mismatched prior: s1={s1}, prior_shape={prior}, prior_bands={bands}".format( |
| name=record.name, |
| s1=(height, width), |
| prior=prior_shape, |
| bands=prior_count, |
| ) |
| ) |
| return height, width |
|
|
|
|
| def load_triplet_window( |
| record: TripletRecord, |
| top: int, |
| left: int, |
| height: int, |
| width: int, |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| """Read one aligned source window instead of materializing a full 67-band tile.""" |
| window = rio.windows.Window(col_off=left, row_off=top, width=width, height=height) |
| with rio.open(str(record.s1_path)) as src: |
| s1_data = src.read(window=window).astype(np.float32) |
| with rio.open(str(record.label_path)) as src: |
| mask_data = src.read(1, window=window).astype(np.float32) |
| s1_missing = ~np.isfinite(s1_data[:2]).all(axis=0) | (s1_data[:2] == 0).any(axis=0) |
| mask_data = np.where(s1_missing, float(IGNORE_INDEX), mask_data).astype(np.float32) |
|
|
| arrays = [s1_data] |
| if record.aef_path is not None: |
| with rio.open(str(record.aef_path)) as src: |
| arrays.append(src.read(window=window).astype(np.float32)) |
| if record.prior_path is not None: |
| with rio.open(str(record.prior_path)) as src: |
| prior_data = src.read(1, window=window).astype(np.float32) |
| prior_nodata = src.nodata |
| prior_invalid = ~np.isfinite(prior_data) |
| if prior_nodata is not None: |
| prior_invalid |= prior_data == prior_nodata |
| prior_data = np.where(prior_invalid, 0.0, prior_data) |
| arrays.append(np.clip(prior_data, 0.0, 1.0)[None, :, :]) |
| image_data = np.concatenate(arrays, axis=0) if len(arrays) > 1 else arrays[0] |
| image_data = np.nan_to_num(image_data, nan=0.0, posinf=0.0, neginf=0.0) |
| image = torch.from_numpy(image_data).float() |
| mask = torch.from_numpy(binarize_label_array(mask_data)).long() |
| return image, mask |
|
|
|
|
| def load_random_triplet_window( |
| record: TripletRecord, |
| crop_size: int, |
| seek_positive: bool = False, |
| min_water_fraction: float = 0.005, |
| attempts: int = 8, |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| height, width = triplet_shape(record) |
| if height < crop_size or width < crop_size: |
| raise ValueError( |
| "Crop size {crop} exceeds tile shape ({height}, {width}) for {name}".format( |
| crop=crop_size, |
| height=height, |
| width=width, |
| name=record.name, |
| ) |
| ) |
| top = random.randint(0, height - crop_size) |
| left = random.randint(0, width - crop_size) |
| if seek_positive: |
| with rio.open(str(record.label_path)) as label_src: |
| for _ in range(attempts): |
| candidate_top = random.randint(0, height - crop_size) |
| candidate_left = random.randint(0, width - crop_size) |
| window = rio.windows.Window( |
| col_off=candidate_left, |
| row_off=candidate_top, |
| width=crop_size, |
| height=crop_size, |
| ) |
| candidate_mask = binarize_label_array(label_src.read(1, window=window)) |
| top, left = candidate_top, candidate_left |
| water_pixels = float(np.count_nonzero(candidate_mask == 1)) |
| if water_pixels / candidate_mask.size >= min_water_fraction: |
| break |
| return load_triplet_window(record, top=top, left=left, height=crop_size, width=crop_size) |
|
|
|
|
| def load_center_triplet_window(record: TripletRecord, crop_size: int) -> Tuple[torch.Tensor, torch.Tensor]: |
| height, width = triplet_shape(record) |
| if height < crop_size or width < crop_size: |
| raise ValueError( |
| "Crop size {crop} exceeds tile shape ({height}, {width}) for {name}".format( |
| crop=crop_size, |
| height=height, |
| width=width, |
| name=record.name, |
| ) |
| ) |
| top = (height - crop_size) // 2 |
| left = (width - crop_size) // 2 |
| return load_triplet_window(record, top=top, left=left, height=crop_size, width=crop_size) |
|
|
|
|
| def load_triplet_tensors(record: TripletRecord) -> Tuple[torch.Tensor, torch.Tensor]: |
| s1 = open_geotiff_float(record.s1_path) |
| mask = open_mask_binary(record.label_path) |
| mask[(s1[:2] == 0).any(dim=0)] = IGNORE_INDEX |
| if record.aef_path is None: |
| if s1.shape[1:] != mask.shape: |
| raise ValueError( |
| "Triplet {name} has mismatched shapes: s1={s1}, mask={mask}".format( |
| name=record.name, s1=tuple(s1.shape), mask=tuple(mask.shape) |
| ) |
| ) |
| image = s1 |
| else: |
| aef = open_geotiff_float(record.aef_path) |
| if s1.shape[1:] != aef.shape[1:] or s1.shape[1:] != mask.shape: |
| raise ValueError( |
| "Triplet {name} has mismatched shapes: s1={s1}, aef={aef}, mask={mask}".format( |
| name=record.name, s1=tuple(s1.shape), aef=tuple(aef.shape), mask=tuple(mask.shape) |
| ) |
| ) |
| image = torch.cat([s1, aef], dim=0) |
| if record.prior_path is not None: |
| prior = open_prior_probability(record.prior_path) |
| if prior.shape[1:] != mask.shape: |
| raise ValueError( |
| "Triplet {name} has mismatched prior/mask shapes: prior={prior}, mask={mask}".format( |
| name=record.name, |
| prior=tuple(prior.shape), |
| mask=tuple(mask.shape), |
| ) |
| ) |
| image = torch.cat([image, prior], dim=0) |
| return image.float(), mask.long() |
|
|
|
|
| def compute_band_stats( |
| records: Sequence[TripletRecord], |
| max_files: Optional[int], |
| normalized_band_count: Optional[int] = None, |
| ) -> Tuple[np.ndarray, np.ndarray]: |
| all_records = list(records) |
| source_ids = sorted({record.source_id for record in all_records if record.source_id}) |
| if max_files is None or max_files >= len(all_records): |
| selected = all_records |
| elif len(source_ids) <= 1: |
| selected = all_records[:max_files] |
| else: |
| selected = [] |
| base_count, remainder = divmod(max_files, len(source_ids)) |
| for source_idx, source_id in enumerate(source_ids): |
| source_records = [record for record in all_records if record.source_id == source_id] |
| source_limit = base_count + (1 if source_idx < remainder else 0) |
| selected.extend(source_records[: min(source_limit, len(source_records))]) |
| if len(selected) < max_files: |
| selected_names = {record.name for record in selected} |
| selected.extend( |
| record for record in all_records if record.name not in selected_names |
| ) |
| selected = selected[:max_files] |
| if not selected: |
| raise ValueError("Cannot compute stats without at least one training record.") |
|
|
| print( |
| { |
| "band_stats_record_count": len(selected), |
| "band_stats_source_counts": dict(Counter(record.source_id or "unspecified" for record in selected)), |
| } |
| ) |
|
|
| channel_sum = None |
| channel_sumsq = None |
| pixel_count = 0 |
|
|
| for record in selected: |
| image, _ = load_triplet_tensors(record) |
| if normalized_band_count is not None: |
| image = image[:normalized_band_count] |
| flat = image.numpy().reshape(image.shape[0], -1) |
| if channel_sum is None: |
| channel_sum = flat.sum(axis=1) |
| channel_sumsq = np.square(flat).sum(axis=1) |
| else: |
| channel_sum += flat.sum(axis=1) |
| channel_sumsq += np.square(flat).sum(axis=1) |
| pixel_count += flat.shape[1] |
|
|
| mean = channel_sum / pixel_count |
| variance = np.maximum((channel_sumsq / pixel_count) - np.square(mean), 1e-8) |
| std = np.sqrt(variance) |
| return mean.astype(np.float32), std.astype(np.float32) |
|
|
|
|
| def get_band_stats( |
| train_records: Sequence[TripletRecord], |
| artifact_dir: Path, |
| stats_sample_size: int, |
| source_stats_path: Optional[Path] = None, |
| expected_bands: Optional[int] = None, |
| ) -> Tuple[np.ndarray, np.ndarray]: |
| stats_cache = artifact_dir / "band_stats.npz" |
| if source_stats_path is not None: |
| if not source_stats_path.exists(): |
| raise FileNotFoundError("Band stats source does not exist: {path}".format(path=source_stats_path)) |
| if source_stats_path.resolve() != stats_cache.resolve(): |
| shutil.copy2(source_stats_path, stats_cache) |
| save_json( |
| artifact_dir / "band_stats_source.json", |
| {"source_band_stats": str(source_stats_path.resolve()), "artifact_band_stats": str(stats_cache)}, |
| ) |
| print("Copied band stats from {source} to {dest}".format(source=source_stats_path, dest=stats_cache)) |
| if stats_cache.exists(): |
| cached = np.load(stats_cache) |
| validate_band_stats(cached["mean"], cached["std"], expected_bands, stats_cache) |
| print("Loaded cached band stats from {path}".format(path=stats_cache)) |
| return cached["mean"], cached["std"] |
| band_mean, band_std = compute_band_stats( |
| train_records, |
| max_files=stats_sample_size, |
| normalized_band_count=expected_bands, |
| ) |
| validate_band_stats(band_mean, band_std, expected_bands, stats_cache) |
| np.savez(stats_cache, mean=band_mean, std=band_std) |
| print("Saved band stats to {path}".format(path=stats_cache)) |
| return band_mean, band_std |
|
|
|
|
| def validate_band_stats(mean: np.ndarray, std: np.ndarray, expected_bands: Optional[int], path: Path) -> None: |
| if mean.shape != std.shape: |
| raise ValueError("Band stats mean/std shape mismatch at {path}: {mean} vs {std}".format(path=path, mean=mean.shape, std=std.shape)) |
| if expected_bands is not None and len(mean) != expected_bands: |
| raise ValueError( |
| "Band stats at {path} have {actual} bands; expected {expected}".format( |
| path=path, |
| actual=len(mean), |
| expected=expected_bands, |
| ) |
| ) |
|
|
|
|
| def normalize_image(image: torch.Tensor, band_mean: torch.Tensor, band_std: torch.Tensor) -> torch.Tensor: |
| normalized_band_count = int(band_mean.shape[0]) |
| if image.shape[0] < normalized_band_count: |
| raise ValueError( |
| "Image has {actual} bands but normalization expects {expected}.".format( |
| actual=image.shape[0], |
| expected=normalized_band_count, |
| ) |
| ) |
| normalized = image.clone() |
| normalized[:normalized_band_count] = ( |
| image[:normalized_band_count] - band_mean |
| ) / band_std |
| return normalized |
|
|
|
|
| def random_crop(image: torch.Tensor, mask: torch.Tensor, crop_size: int) -> Tuple[torch.Tensor, torch.Tensor]: |
| _, height, width = image.shape |
| if height < crop_size or width < crop_size: |
| raise ValueError( |
| "Crop size {crop} exceeds tile shape ({h}, {w})".format( |
| crop=crop_size, h=height, w=width |
| ) |
| ) |
| top = random.randint(0, height - crop_size) |
| left = random.randint(0, width - crop_size) |
| return image[:, top : top + crop_size, left : left + crop_size], mask[top : top + crop_size, left : left + crop_size] |
|
|
|
|
| def apply_train_augmentations( |
| image: torch.Tensor, |
| mask: torch.Tensor, |
| crop_size: int, |
| n_s1_bands: int, |
| rotate_max_deg: float, |
| rotate_p: float, |
| scale_jitter_min: float, |
| scale_jitter_max: float, |
| s1_gain_jitter: float, |
| s1_bias_jitter: float, |
| gaussian_noise_std: float, |
| gaussian_noise_p: float, |
| normalized_band_count: int, |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| if random.random() < 0.5: |
| image = torch.flip(image, dims=(-1,)) |
| mask = torch.flip(mask, dims=(-1,)) |
|
|
| rotations = random.randint(0, 3) |
| if rotations: |
| image = torch.rot90(image, rotations, dims=(-2, -1)) |
| mask = torch.rot90(mask, rotations, dims=(-2, -1)) |
|
|
| if scale_jitter_min != 1.0 or scale_jitter_max != 1.0: |
| scale = random.uniform(scale_jitter_min, scale_jitter_max) |
| scaled_size = max(32, int(round(crop_size * scale))) |
| if scaled_size != crop_size: |
| image = F.interpolate( |
| image.unsqueeze(0), |
| size=(scaled_size, scaled_size), |
| mode="bilinear", |
| align_corners=False, |
| ).squeeze(0) |
| mask = F.interpolate( |
| mask.unsqueeze(0).unsqueeze(0).float(), |
| size=(scaled_size, scaled_size), |
| mode="nearest", |
| ).squeeze(0).squeeze(0) |
| if scaled_size > crop_size: |
| image, mask = random_crop(image, mask, crop_size) |
| else: |
| pad_h = crop_size - scaled_size |
| pad_w = crop_size - scaled_size |
| top = random.randint(0, pad_h) |
| left = random.randint(0, pad_w) |
| bottom = pad_h - top |
| right = pad_w - left |
| image = F.pad(image, (left, right, top, bottom), mode="reflect") |
| mask = F.pad(mask, (left, right, top, bottom), mode="constant", value=0) |
|
|
| if rotate_max_deg > 0 and random.random() < rotate_p: |
| angle = random.uniform(-rotate_max_deg, rotate_max_deg) |
| image = TVF.rotate(image, angle=angle, interpolation=InterpolationMode.BILINEAR) |
| mask = TVF.rotate( |
| mask.unsqueeze(0).float(), |
| angle=angle, |
| interpolation=InterpolationMode.NEAREST, |
| ).squeeze(0) |
|
|
| if s1_gain_jitter > 0 or s1_bias_jitter > 0: |
| gain = 1.0 + random.uniform(-s1_gain_jitter, s1_gain_jitter) |
| bias = random.uniform(-s1_bias_jitter, s1_bias_jitter) |
| image[:n_s1_bands] = (image[:n_s1_bands] * gain) + bias |
|
|
| if gaussian_noise_std > 0 and random.random() < gaussian_noise_p: |
| image[:normalized_band_count] = image[:normalized_band_count] + ( |
| torch.randn_like(image[:normalized_band_count]) * gaussian_noise_std |
| ) |
|
|
| return image.float(), mask.long() |
|
|
|
|
| def build_dataloaders( |
| args: argparse.Namespace, |
| train_records: Sequence[TripletRecord], |
| valid_records: Sequence[TripletRecord], |
| band_mean: np.ndarray, |
| band_std: np.ndarray, |
| device: torch.device, |
| ) -> DataLoaders: |
| train_dataset = S1AEFCropDataset( |
| records=train_records, |
| crop_size=args.crop_size, |
| band_mean=band_mean, |
| band_std=band_std, |
| n_s1_bands=args.n_s1_bands, |
| rotate_max_deg=args.rotate_max_deg, |
| rotate_p=args.rotate_p, |
| scale_jitter_min=args.scale_jitter_min, |
| scale_jitter_max=args.scale_jitter_max, |
| s1_gain_jitter=args.s1_gain_jitter, |
| s1_bias_jitter=args.s1_bias_jitter, |
| gaussian_noise_std=args.gaussian_noise_std, |
| gaussian_noise_p=args.gaussian_noise_p, |
| positive_crop_p=args.positive_crop_p, |
| min_water_fraction=args.min_water_fraction, |
| positive_crop_attempts=args.positive_crop_attempts, |
| ) |
| valid_dataset = S1AEFValidationPatchDataset( |
| valid_records, |
| crop_size=args.crop_size, |
| band_mean=band_mean, |
| band_std=band_std, |
| ) |
|
|
| if args.train_tile_index is not None: |
| train_dataset = S1AEFTileDataset( |
| read_tile_index(args.train_tile_index, train_records, "train"), crop_size=args.crop_size, band_mean=band_mean, |
| band_std=band_std, n_s1_bands=args.n_s1_bands, rotate_max_deg=args.rotate_max_deg, rotate_p=args.rotate_p, |
| scale_jitter_min=args.scale_jitter_min, scale_jitter_max=args.scale_jitter_max, s1_gain_jitter=args.s1_gain_jitter, |
| s1_bias_jitter=args.s1_bias_jitter, gaussian_noise_std=args.gaussian_noise_std, gaussian_noise_p=args.gaussian_noise_p) |
| print({"full_coverage": "train", "tiles": len(train_dataset), "chips": len(train_records), |
| "tiles_by_source": dict(Counter(tile[0].source_id for tile in train_dataset.tiles))}) |
| if args.valid_tile_index is not None: |
| valid_dataset = S1AEFValidationTileDataset( |
| read_tile_index(args.valid_tile_index, valid_records, "valid"), valid_records, crop_size=args.crop_size, |
| band_mean=band_mean, band_std=band_std) |
| print({"full_coverage": "valid", "tiles": len(valid_dataset), "chips": len(valid_records), |
| "tiles_by_source": dict(Counter(valid_dataset.item_source_ids))}) |
| pin_memory = device.type == "cuda" and not args.disable_pin_memory |
| train_loader_kwargs = { |
| "dataset": train_dataset, |
| "bs": args.batch_size, |
| "shuffle": True, |
| "num_workers": args.num_workers, |
| "pin_memory": pin_memory, |
| "persistent_workers": args.num_workers > 0, |
| "device": None, |
| "wif": seed_worker, |
| } |
| if args.target_source_fraction is not None: |
| target_indices = [ |
| idx for idx, record in enumerate(train_records) if record.source_id == args.target_source_id |
| ] |
| other_indices = [ |
| idx for idx, record in enumerate(train_records) if record.source_id != args.target_source_id |
| ] |
| if not target_indices or not other_indices: |
| raise ValueError( |
| "Source-aware sampling requires both target source {target} and non-target records. " |
| "Observed counts: {counts}".format( |
| target=args.target_source_id, |
| counts=dict(Counter(record.source_id or "unspecified" for record in train_records)), |
| ) |
| ) |
| target_weight = args.target_source_fraction / len(target_indices) |
| other_weight = (1.0 - args.target_source_fraction) / len(other_indices) |
| record_weights = [ |
| target_weight if record.source_id == args.target_source_id else other_weight |
| for record in train_records |
| ] |
|
|
| source_sampler_rng = random.Random(args.seed) |
|
|
| def source_weighted_indices(): |
| return source_sampler_rng.choices( |
| list(range(len(record_weights))), |
| weights=record_weights, |
| k=len(record_weights), |
| ) |
|
|
| train_loader_kwargs["shuffle"] = False |
| train_loader_kwargs["get_idxs"] = source_weighted_indices |
| print( |
| { |
| "source_sampling": "weighted_with_replacement", |
| "target_source_id": args.target_source_id, |
| "target_source_fraction": args.target_source_fraction, |
| "epoch_samples": len(train_dataset), |
| "source_record_counts": dict( |
| Counter(record.source_id or "unspecified" for record in train_records) |
| ), |
| } |
| ) |
| else: |
| print( |
| { |
| "source_sampling": "natural_without_replacement", |
| "epoch_samples": len(train_dataset), |
| "source_record_counts": dict( |
| Counter(record.source_id or "unspecified" for record in train_records) |
| ), |
| } |
| ) |
| train_dl = DataLoader(**train_loader_kwargs) |
| valid_dl = DataLoader( |
| valid_dataset, |
| bs=1, |
| shuffle=False, |
| num_workers=args.num_workers, |
| pin_memory=pin_memory, |
| persistent_workers=args.num_workers > 0, |
| device=None, |
| wif=seed_worker, |
| ) |
| dls = DataLoaders(train_dl, valid_dl, path=args.artifact_dir, device=device) |
| dls.c = len(args.codes) |
| dls.vocab = args.codes |
| return dls |
|
|
|
|
| def format_lr(lr_max): |
| if isinstance(lr_max, slice): |
| return {"low": lr_max.start, "high": lr_max.stop} |
| return lr_max |
|
|
|
|
| class GlobalBestSaveModelCallback(SaveModelCallback): |
| """full_coverage_v1: best checkpoint over all training stages (the stock callback restarts its best at each stage).""" |
|
|
| def before_fit(self): |
| super().before_fit() |
| best = getattr(self.learn, "global_best_metric", None) |
| if best is not None: |
| self.best = best |
|
|
| def after_epoch(self): |
| super().after_epoch() |
| if self.new_best: |
| self.learn.global_best_metric = self.best |
| self.learn.global_best_where = {"phase": getattr(self.learn, "training_phase", None), "epoch_in_phase": int(self.epoch), |
| "value": float(self.best)} |
| print({"global_best": self.learn.global_best_where}) |
|
|
|
|
| def build_callbacks(args: argparse.Namespace, epoch_timer: Optional[EpochTimerCallback] = None) -> List[Callback]: |
| callbacks: List[Callback] = [ |
| GradientClip(args.grad_clip), |
| NonFiniteLossCallback(), |
| GlobalBestSaveModelCallback( |
| monitor=args.selection_metric, |
| comp=np.greater, |
| min_delta=args.min_delta, |
| fname=args.best_model_name, |
| ), |
| EarlyStoppingCallback( |
| monitor=args.selection_metric, |
| comp=np.greater, |
| min_delta=args.min_delta, |
| patience=args.patience, |
| ), |
| ] |
| if args.grad_accum_steps > 1: |
| callbacks.append(GradientAccumulation(n_acc=args.grad_accum_steps)) |
| if epoch_timer is not None: |
| callbacks.append(epoch_timer) |
| return callbacks |
|
|
|
|
| def recorder_history_to_df(learn: Learner, phase: str) -> pd.DataFrame: |
| values = list(learn.recorder.values) |
| if not values: |
| return pd.DataFrame(columns=["phase"]) |
|
|
| column_count = len(values[0]) |
| metric_names = list(learn.recorder.metric_names[1:]) |
| if metric_names and metric_names[-1] == "time" and len(metric_names) == column_count + 1: |
| metric_names = metric_names[:-1] |
| if len(metric_names) != column_count: |
| metric_names = ["value_{idx}".format(idx=idx) for idx in range(column_count)] |
| return pd.DataFrame(values, columns=metric_names).assign(phase=phase) |
|
|
|
|
| def run_phase( |
| learn: Learner, |
| args: argparse.Namespace, |
| phase: str, |
| epochs: int, |
| lr_max, |
| epoch_timer: Optional[EpochTimerCallback] = None, |
| ) -> pd.DataFrame: |
| if epochs <= 0: |
| return pd.DataFrame(columns=["phase"]) |
|
|
| print({"phase": phase, "epochs": epochs, "lr_max": format_lr(lr_max)}) |
| learn.training_phase = phase |
| try: |
| fit_kwargs = { |
| "wd": args.weight_decay, |
| "cbs": build_callbacks(args, epoch_timer=epoch_timer), |
| "reset_opt": True, |
| } |
| if args.schedule_policy == "flat_cos": |
| learn.fit_flat_cos(epochs, lr=lr_max, **fit_kwargs) |
| else: |
| learn.fit_one_cycle(epochs, lr_max=lr_max, **fit_kwargs) |
| finally: |
| learn.training_phase = None |
| return recorder_history_to_df(learn, phase=phase) |
|
|
|
|
| def build_schedule(args: argparse.Namespace) -> List[Tuple[str, int, object]]: |
| if args.mode == "benchmark": |
| return [("freeze", 2, args.warmup_lr)] |
| if args.tuning_strategy != "staged": |
| return [(args.tuning_strategy, args.strategy_epochs, args.strategy_lr)] |
| schedule: List[Tuple[str, int, object]] = [ |
| ("freeze", args.freeze_epochs, args.warmup_lr), |
| ] |
| bridge_epochs = min(args.partial_unfreeze_epochs, args.finetune_epochs) |
| remaining_finetune_epochs = max(args.finetune_epochs - bridge_epochs, 0) |
| if bridge_epochs > 0: |
| schedule.append( |
| ( |
| "partial_unfreeze", |
| bridge_epochs, |
| slice(args.partial_unfreeze_lr / 20, args.partial_unfreeze_lr), |
| ) |
| ) |
| if remaining_finetune_epochs > 0: |
| schedule.append( |
| ( |
| "finetune", |
| remaining_finetune_epochs, |
| slice(args.unfreeze_lr / 50, args.unfreeze_lr), |
| ) |
| ) |
| if args.finetune_tail_epochs > 0: |
| schedule.append( |
| ( |
| "finetune_tail", |
| args.finetune_tail_epochs, |
| slice(args.finetune_tail_lr / 10, args.finetune_tail_lr), |
| ) |
| ) |
| return schedule |
|
|
|
|
| def set_requires_grad(module: nn.Module, trainable: bool) -> None: |
| for parameter in module.parameters(): |
| parameter.requires_grad = trainable |
|
|
|
|
| def trainable_parameter_summary(model: nn.Module) -> Dict[str, int]: |
| trainable = sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad) |
| total = sum(parameter.numel() for parameter in model.parameters()) |
| return {"trainable_parameters": int(trainable), "total_parameters": int(total)} |
|
|
|
|
| def resnet_unet_children(model: nn.Module) -> List[nn.Module]: |
| if not isinstance(model, AEFBottleneckUNet): |
| raise ValueError("Explicit tuning strategies currently require AEFBottleneckUNet.") |
| children = list(model.unet.layers) if hasattr(model.unet, "layers") else list(model.unet.children()) |
| if len(children) < 3: |
| raise ValueError("Unexpected DynamicUnet layout for tuning strategy.") |
| return children |
|
|
|
|
| def enable_head_only(model: nn.Module) -> None: |
| children = resnet_unet_children(model) |
| parameterized_children = [module for module in children if any(True for _ in module.parameters())] |
| if not parameterized_children: |
| raise ValueError("DynamicUnet head did not expose trainable parameters.") |
| head = parameterized_children[-1] |
| set_requires_grad(head, True) |
|
|
|
|
| def enable_decoder(model: nn.Module) -> None: |
| children = resnet_unet_children(model) |
| for module in children[1:]: |
| set_requires_grad(module, True) |
|
|
|
|
| def enable_aef_projection(model: nn.Module) -> None: |
| if not isinstance(model, AEFBottleneckUNet): |
| raise ValueError("AEF projection strategy requires AEFBottleneckUNet.") |
| set_requires_grad(model.aef_proj, True) |
|
|
|
|
| def enable_last_encoder_stage(model: nn.Module) -> None: |
| encoder = resnet_unet_children(model)[0] |
| encoder_children = list(encoder.children()) |
| if not encoder_children: |
| raise ValueError("DynamicUnet encoder did not expose child stages.") |
| set_requires_grad(encoder_children[-1], True) |
|
|
|
|
| def apply_explicit_tuning_strategy(learn: Learner, strategy: str) -> None: |
| set_requires_grad(learn.model, False) |
| if strategy == "head_only": |
| enable_head_only(learn.model) |
| elif strategy == "decoder_aef": |
| enable_decoder(learn.model) |
| enable_aef_projection(learn.model) |
| elif strategy == "last_encoder_decoder_aef": |
| enable_decoder(learn.model) |
| enable_aef_projection(learn.model) |
| enable_last_encoder_stage(learn.model) |
| elif strategy == "full_low_lr": |
| set_requires_grad(learn.model, True) |
| else: |
| raise ValueError("Unknown explicit tuning strategy: {}".format(strategy)) |
| summary = trainable_parameter_summary(learn.model) |
| if summary["trainable_parameters"] == 0: |
| raise ValueError("Tuning strategy {} left no trainable parameters.".format(strategy)) |
| print({"tuning_strategy_scope": strategy, **summary}) |
|
|
|
|
| def apply_phase_freeze(learn: Learner, phase: str) -> None: |
| if phase in {"head_only", "decoder_aef", "last_encoder_decoder_aef", "full_low_lr"}: |
| apply_explicit_tuning_strategy(learn, phase) |
| return |
| if phase == "freeze": |
| learn.freeze() |
| elif phase == "partial_unfreeze": |
| learn.freeze_to(-2) |
| else: |
| learn.unfreeze() |
|
|
|
|
| def build_wrapped_splitter(base_splitter): |
| def wrapped_splitter(model: AEFBottleneckUNet): |
| groups = list(base_splitter(model.unet)) |
| if not groups: |
| return L([L(model.parameters())]) |
| final_group = L(groups[-1]) + L(model.aef_proj.parameters()) |
| return L(groups[:-1]) + L([final_group]) |
|
|
| return wrapped_splitter |
|
|
|
|
| def build_segformer_splitter(): |
| def splitter(model: AEFBottleneckSegFormer): |
| encoder = model.segformer.segformer.encoder |
| groups: List[L] = [] |
| for patch_embedding, blocks, norm in zip( |
| encoder.patch_embeddings, |
| encoder.block, |
| encoder.layer_norm, |
| ): |
| groups.append( |
| L(patch_embedding.parameters()) + L(blocks.parameters()) + L(norm.parameters()) |
| ) |
| groups.append(L(model.segformer.decode_head.parameters()) + L(model.aef_proj.parameters())) |
| return L(groups) |
|
|
| return splitter |
|
|
|
|
| def build_loss(args: argparse.Namespace) -> nn.Module: |
| loss_name = args.loss |
| if loss_name == "focal_dice": |
| return CombinedLoss() |
| if loss_name == "ce_dice": |
| return CrossEntropyDiceLoss() |
| if loss_name == "lovasz": |
| return LovaszSoftmaxLoss() |
| if loss_name == "ce_soft_iou": |
| weight = 1.0 if args.loss_aux_weight is None else args.loss_aux_weight |
| return CrossEntropySoftIoULoss(aux_weight=weight) |
| if loss_name == "focal_tversky": |
| weight = 1.0 if args.loss_aux_weight is None else args.loss_aux_weight |
| return FocalTverskyCELoss( |
| alpha=args.tversky_alpha, |
| beta=args.tversky_beta, |
| gamma=args.tversky_gamma, |
| aux_weight=weight, |
| ) |
| if loss_name == "ce_dice_boundary": |
| weight = 0.2 if args.loss_aux_weight is None else args.loss_aux_weight |
| return BoundaryAwareCEDiceLoss(aux_weight=weight) |
| if loss_name == "ce_dice_cldice": |
| weight = 0.2 if args.loss_aux_weight is None else args.loss_aux_weight |
| return TopologyAwareCEDiceLoss( |
| aux_weight=weight, |
| skeleton_iterations=args.cldice_iterations, |
| ) |
| if loss_name == "sce_dice": |
| weight = 0.1 if args.loss_aux_weight is None else args.loss_aux_weight |
| return SymmetricCrossEntropyDiceLoss(aux_weight=weight) |
| raise ValueError("Unknown loss: {name}".format(name=loss_name)) |
|
|
|
|
| def build_optimizer(args: argparse.Namespace): |
| if args.optimizer == "ranger": |
| return ranger |
| if args.optimizer == "adamw": |
| return Adam |
| raise ValueError("Unknown optimizer: {name}".format(name=args.optimizer)) |
|
|
|
|
| def build_learner( |
| args: argparse.Namespace, |
| dls: DataLoaders, |
| valid_records: Sequence[TripletRecord], |
| ) -> Learner: |
| n_prior_bands = 1 if using_prior(args) else 0 |
| n_total_in = args.n_s1_bands + args.n_aef_bands + n_prior_bands |
| n_model_in = ( |
| args.n_s1_bands |
| + (args.n_proj_bands if using_aef(args) else 0) |
| + n_prior_bands |
| ) |
| loss_func = build_loss(args) |
| opt_func = build_optimizer(args) |
| metrics = [ |
| MaskedJaccardCoeff(), |
| MaskedDice(), |
| WaterIoU(), |
| SourceBalancedWaterIoU(getattr(dls.valid.dataset, "item_source_ids", [record.source_id for record in valid_records])), |
| TargetSourceWaterIoU( |
| getattr(dls.valid.dataset, "item_source_ids", [record.source_id for record in valid_records]), |
| target_source_id=args.target_source_id, |
| ), |
| ] |
| if args.model_family == "segformer": |
| wrapped_model = AEFBottleneckSegFormer( |
| model_name=SEGFORMER_MODEL_MAP[args.segformer_variant], |
| num_labels=len(args.codes), |
| n_s1=args.n_s1_bands, |
| n_aef=args.n_aef_bands, |
| n_proj=args.n_proj_bands, |
| hf_cache_dir=args.hf_cache_dir, |
| pretrained=args.pretrained, |
| local_files_only=args.hf_local_files_only, |
| ) |
| wrapped_model = wrapped_model.to(dls.device) |
| learn = Learner( |
| dls, |
| wrapped_model, |
| loss_func=loss_func, |
| metrics=metrics, |
| opt_func=opt_func, |
| splitter=build_segformer_splitter(), |
| path=args.artifact_dir, |
| model_dir="models", |
| ) |
| model_name = "AEFBottleneckSegFormer({variant})".format( |
| variant=args.segformer_variant |
| ) |
| else: |
| encoder_arches = { |
| "resnet34_unet": resnet34, |
| "resnet50_unet": resnet50, |
| "resnet101_unet": resnet101, |
| } |
| encoder_arch = encoder_arches[args.model_family] |
| learn = unet_learner( |
| dls, |
| encoder_arch, |
| normalize=False, |
| n_in=n_model_in, |
| n_out=len(args.codes), |
| metrics=metrics, |
| loss_func=loss_func, |
| opt_func=opt_func, |
| act_cls=Mish, |
| path=args.artifact_dir, |
| model_dir="models", |
| pretrained=args.pretrained, |
| ) |
| if using_aef(args): |
| base_splitter = learn.splitter |
| wrapped_model = AEFBottleneckUNet( |
| unet_model=learn.model, |
| n_s1=args.n_s1_bands, |
| n_aef=args.n_aef_bands, |
| n_proj=args.n_proj_bands, |
| n_prior=n_prior_bands, |
| prior_dropout_p=args.prior_dropout_p, |
| aef_dropout_p=getattr(args, "aef_dropout_p", 0.0), |
| ) |
| learn.model = wrapped_model.to(dls.device) |
| learn.splitter = build_wrapped_splitter(base_splitter) |
| model_name = "AEFBottleneckUNet({family})".format(family=args.model_family) |
| else: |
| model_name = "S1UNet({family})".format(family=args.model_family) |
| learn.model = learn.model.to(dls.device) |
| if args.use_fp16 and torch.cuda.is_available(): |
| learn = learn.to_fp16() |
|
|
| model_summary = { |
| "model": model_name, |
| "dataloader_input_bands": n_total_in, |
| "backbone_input_bands": n_model_in, |
| "loss": args.loss, |
| "loss_aux_weight": args.loss_aux_weight, |
| "optimizer": args.optimizer, |
| "schedule_policy": args.schedule_policy, |
| "selection_metric": args.selection_metric, |
| "tuning_strategy": args.tuning_strategy, |
| "grad_accum_steps": args.grad_accum_steps, |
| "use_permanent_prior": using_prior(args), |
| "prior_dropout_p": args.prior_dropout_p, |
| "aef_dropout_p": getattr(args, "aef_dropout_p", 0.0), |
| "precision": "fp16" if args.use_fp16 and torch.cuda.is_available() else "fp32", |
| "best_model_path": str(args.artifact_dir / "models" / "{name}.pth".format(name=args.best_model_name)), |
| } |
| if using_aef(args): |
| model_summary["bottleneck"] = "{aef}->{proj}".format( |
| aef=args.n_aef_bands, proj=args.n_proj_bands |
| ) |
| print( |
| model_summary |
| ) |
| return learn |
|
|
|
|
| def resolve_warm_start_checkpoint(args: argparse.Namespace) -> Optional[Path]: |
| if args.warm_start_checkpoint is not None: |
| return args.warm_start_checkpoint |
| if args.warm_start_run_dir is None: |
| return None |
| return args.warm_start_run_dir / "models" / "{name}.pth".format(name=args.best_model_name) |
|
|
|
|
| def resolve_band_stats_source(args: argparse.Namespace) -> Optional[Path]: |
| if args.band_stats_path is not None: |
| return args.band_stats_path |
| if args.warm_start_run_dir is not None: |
| return args.warm_start_run_dir / "band_stats.npz" |
| return None |
|
|
|
|
| def unwrap_state_dict(checkpoint: object) -> Dict[str, torch.Tensor]: |
| if not isinstance(checkpoint, dict): |
| raise ValueError("Warm-start checkpoint must contain a state dict mapping.") |
| for key in ("model", "state_dict"): |
| value = checkpoint.get(key) |
| if isinstance(value, dict): |
| checkpoint = value |
| break |
| if not checkpoint or not all(isinstance(key, str) for key in checkpoint.keys()): |
| raise ValueError("Warm-start checkpoint does not expose string state-dict keys.") |
| return checkpoint |
|
|
|
|
| def validate_warm_start_projection( |
| state_dict: Dict[str, torch.Tensor], |
| model: nn.Module, |
| args: argparse.Namespace, |
| checkpoint_path: Path, |
| ) -> None: |
| if not using_aef(args): |
| return |
| key = "aef_proj.weight" |
| if key not in state_dict: |
| raise ValueError("Warm-start checkpoint is missing {key}: {path}".format(key=key, path=checkpoint_path)) |
| checkpoint_shape = tuple(state_dict[key].shape) |
| model_shape = tuple(model.state_dict()[key].shape) |
| expected_shape = (args.n_proj_bands, args.n_aef_bands, 1, 1) |
| if checkpoint_shape != expected_shape or model_shape != expected_shape: |
| raise ValueError( |
| "Warm-start AEF projection shape mismatch at {path}: checkpoint={ckpt}, model={model}, expected={expected}".format( |
| path=checkpoint_path, |
| ckpt=checkpoint_shape, |
| model=model_shape, |
| expected=expected_shape, |
| ) |
| ) |
|
|
|
|
| def maybe_warm_start(learn: Learner, args: argparse.Namespace) -> Optional[Path]: |
| checkpoint_path = resolve_warm_start_checkpoint(args) |
| if checkpoint_path is None: |
| return None |
| if not checkpoint_path.exists(): |
| raise FileNotFoundError("Warm-start checkpoint does not exist: {path}".format(path=checkpoint_path)) |
| state_dict = unwrap_state_dict(torch.load(checkpoint_path, map_location="cpu")) |
| validate_warm_start_projection(state_dict, learn.model, args, checkpoint_path) |
| try: |
| learn.model.load_state_dict(state_dict, strict=True) |
| except RuntimeError as exc: |
| raise RuntimeError("Warm-start checkpoint is incompatible with this model: {path}\n{error}".format(path=checkpoint_path, error=exc)) from exc |
| save_json( |
| args.artifact_dir / "warm_start.json", |
| { |
| "checkpoint": str(checkpoint_path.resolve()), |
| "parent_run_dir": "" if args.warm_start_run_dir is None else str(args.warm_start_run_dir), |
| "n_aef_bands": args.n_aef_bands, |
| "n_proj_bands": args.n_proj_bands, |
| }, |
| ) |
| print({"warm_start_checkpoint": str(checkpoint_path), "warm_start_strict": True}) |
| return checkpoint_path |
|
|
|
|
| def save_json(path: Path, payload: dict) -> None: |
| path.write_text(json.dumps(payload, indent=2)) |
|
|
|
|
| def inspect_bottleneck(model: nn.Module, artifact_dir: Path) -> dict: |
| weight = model.aef_proj.weight.detach().cpu().squeeze(-1).squeeze(-1) |
| abs_weight = weight.abs() |
| top_k = 5 |
| report = { |
| "projection_shape": list(weight.shape), |
| "weight_norm_per_proj_band": [float(v) for v in weight.norm(dim=1).tolist()], |
| "weight_norm_per_aef_band": [float(v) for v in weight.norm(dim=0).tolist()], |
| } |
| for idx in range(weight.shape[0]): |
| top_vals, top_idx = abs_weight[idx].topk(top_k) |
| report["proj_band_{idx}_top{top}_aef_bands".format(idx=idx, top=top_k)] = top_idx.tolist() |
| report["proj_band_{idx}_top{top}_weights".format(idx=idx, top=top_k)] = [ |
| round(v, 6) for v in top_vals.tolist() |
| ] |
|
|
| np.save(artifact_dir / "aef_bottleneck_weights.npy", weight.numpy()) |
| save_json(artifact_dir / "aef_bottleneck_report.json", report) |
| print( |
| { |
| "bottleneck_weight_file": str(artifact_dir / "aef_bottleneck_weights.npy"), |
| "bottleneck_report_file": str(artifact_dir / "aef_bottleneck_report.json"), |
| } |
| ) |
| return report |
|
|
|
|
| def get_tta_transforms(mode_count: int): |
| transforms = [ |
| ("identity", lambda x: x), |
| ("flip_horizontal", lambda x: torch.flip(x, dims=(-1,))), |
| ("flip_vertical", lambda x: torch.flip(x, dims=(-2,))), |
| ("flip_both", lambda x: torch.flip(x, dims=(-2, -1))), |
| ] |
| return transforms[: max(1, min(mode_count, len(transforms)))] |
|
|
|
|
| def is_cuda_oom(error: RuntimeError) -> bool: |
| message = str(error).lower() |
| return "out of memory" in message or "cuda error: out of memory" in message |
|
|
|
|
| def sliding_positions(size: int, tile_size: int, overlap: int) -> List[int]: |
| if size <= tile_size: |
| return [0] |
| stride = max(1, tile_size - overlap) |
| positions = list(range(0, size - tile_size + 1, stride)) |
| if positions[-1] != size - tile_size: |
| positions.append(size - tile_size) |
| return positions |
|
|
|
|
| def predict_record_logits_windowed( |
| model: nn.Module, |
| record: TripletRecord, |
| band_mean: torch.Tensor, |
| band_std: torch.Tensor, |
| tile_size: int, |
| overlap: int, |
| device: torch.device, |
| use_tta: bool = False, |
| tta_mode_count: int = 4, |
| ) -> torch.Tensor: |
| """Blend window predictions on CPU while only one input patch resides on the GPU.""" |
| height, width = triplet_shape(record) |
| y_positions = sliding_positions(height, tile_size, overlap) |
| x_positions = sliding_positions(width, tile_size, overlap) |
| logits_sum = None |
| weight_sum = torch.zeros((height, width), dtype=torch.float32) |
| transforms = get_tta_transforms(tta_mode_count) if use_tta else [("identity", lambda x: x)] |
|
|
| for top in y_positions: |
| for left in x_positions: |
| window_height = min(tile_size, height - top) |
| window_width = min(tile_size, width - left) |
| image, _ = load_triplet_window( |
| record, |
| top=top, |
| left=left, |
| height=window_height, |
| width=window_width, |
| ) |
| image = normalize_image(image, band_mean, band_std) |
| xb = image.unsqueeze(0).to(device, non_blocking=False) |
| patch_logits_sum = None |
| for _, transform in transforms: |
| augmented_x = transform(xb) |
| augmented_logits = model(augmented_x).float() |
| restored_logits = transform(augmented_logits) |
| patch_logits_sum = ( |
| restored_logits if patch_logits_sum is None else patch_logits_sum + restored_logits |
| ) |
| patch_logits = (patch_logits_sum / len(transforms))[0].detach().cpu() |
| if logits_sum is None: |
| logits_sum = torch.zeros( |
| (patch_logits.shape[0], height, width), |
| dtype=torch.float32, |
| ) |
| bottom = top + window_height |
| right = left + window_width |
| logits_sum[:, top:bottom, left:right] += patch_logits[:, :window_height, :window_width] |
| weight_sum[top:bottom, left:right] += 1.0 |
| del image, xb, patch_logits, patch_logits_sum |
|
|
| if logits_sum is None: |
| raise RuntimeError("No inference windows were produced for {name}".format(name=record.name)) |
| return logits_sum / weight_sum.clamp_min(1.0).unsqueeze(0) |
|
|
|
|
| def windowed_validation_context(learn: Learner) -> Tuple[List[TripletRecord], torch.Tensor, torch.Tensor]: |
| dataset = learn.dls.valid.dataset |
| if not hasattr(dataset, "records") or not hasattr(dataset, "band_mean") or not hasattr(dataset, "band_std"): |
| raise TypeError("Validation dataset does not expose records and normalization statistics.") |
| return list(dataset.records), dataset.band_mean, dataset.band_std |
|
|
|
|
| def predict_logits_tiled(model: nn.Module, xb: torch.Tensor, tile_size: int, overlap: int) -> torch.Tensor: |
| batch_size, _, height, width = xb.shape |
| if batch_size != 1: |
| raise ValueError("Tiled inference expects batch size 1, got {bs}".format(bs=batch_size)) |
|
|
| y_positions = sliding_positions(height, tile_size, overlap) |
| x_positions = sliding_positions(width, tile_size, overlap) |
| logits_sum = None |
| weight_sum = torch.zeros((1, 1, height, width), device=xb.device, dtype=torch.float32) |
|
|
| for top in y_positions: |
| for left in x_positions: |
| bottom = min(top + tile_size, height) |
| right = min(left + tile_size, width) |
| tile = xb[:, :, top:bottom, left:right] |
| logits = model(tile).float() |
| if logits_sum is None: |
| logits_sum = torch.zeros( |
| (1, logits.shape[1], height, width), device=xb.device, dtype=torch.float32 |
| ) |
| logits_sum[:, :, top:bottom, left:right] += logits |
| weight_sum[:, :, top:bottom, left:right] += 1.0 |
|
|
| return logits_sum / weight_sum.clamp_min(1.0) |
|
|
|
|
| def predict_logits_with_fallback( |
| model: nn.Module, |
| xb: torch.Tensor, |
| tile_size: int, |
| overlap: int, |
| force_tiled: bool = False, |
| ) -> Tuple[torch.Tensor, bool]: |
| if force_tiled: |
| return predict_logits_tiled(model, xb, tile_size=tile_size, overlap=overlap), True |
| try: |
| return model(xb).float(), False |
| except RuntimeError as error: |
| if not is_cuda_oom(error): |
| raise |
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
| print({"eval_fallback": "switching to tiled inference"}) |
| return predict_logits_tiled(model, xb, tile_size=tile_size, overlap=overlap), True |
|
|
|
|
| def predict_water_probabilities( |
| model: nn.Module, |
| xb: torch.Tensor, |
| tile_size: int, |
| overlap: int, |
| use_tta: bool = False, |
| tta_mode_count: int = 4, |
| force_tiled: bool = False, |
| ) -> Tuple[torch.Tensor, bool]: |
| if not use_tta: |
| logits, used_tiled = predict_logits_with_fallback( |
| model, xb, tile_size=tile_size, overlap=overlap, force_tiled=force_tiled |
| ) |
| return F.softmax(logits, dim=1)[:, 1], used_tiled |
|
|
| transforms = get_tta_transforms(tta_mode_count) |
| probs_sum = None |
| used_tiled_any = force_tiled |
| for _, transform in transforms: |
| augmented_x = transform(xb) |
| logits, used_tiled = predict_logits_with_fallback( |
| model, augmented_x, tile_size=tile_size, overlap=overlap, force_tiled=used_tiled_any |
| ) |
| augmented_probs = F.softmax(logits, dim=1)[:, 1:2] |
| restored_probs = transform(augmented_probs) |
| probs_sum = restored_probs if probs_sum is None else probs_sum + restored_probs |
| used_tiled_any = used_tiled_any or used_tiled |
| return (probs_sum / len(transforms))[:, 0], used_tiled_any |
|
|
|
|
| def metrics_from_binary_confusion(tp: int, fp: int, tn: int, fn: int) -> dict: |
| precision = tp / max(tp + fp, 1) |
| recall = tp / max(tp + fn, 1) |
| water_iou = tp / max(tp + fp + fn, 1) |
| dice = (2 * tp) / max((2 * tp) + fp + fn, 1) |
| accuracy = (tp + tn) / max(tp + tn + fp + fn, 1) |
| other_iou = tn / max(tn + fp + fn, 1) |
| other_dice = (2 * tn) / max((2 * tn) + fp + fn, 1) |
| mean_iou = (other_iou + water_iou) / 2.0 |
| mean_dice = (other_dice + dice) / 2.0 |
| return { |
| "precision": float(precision), |
| "recall": float(recall), |
| "water_iou": float(water_iou), |
| "dice": float(mean_dice), |
| "accuracy": float(accuracy), |
| "jaccard": float(mean_iou), |
| } |
|
|
|
|
| def validate_model( |
| learn: Learner, |
| tile_size: int, |
| overlap: int, |
| force_tiled: bool = False, |
| max_valid_batches: Optional[int] = None, |
| ) -> dict: |
| device = next(learn.model.parameters()).device |
| loss_func = learn.loss_func |
| records, band_mean, band_std = windowed_validation_context(learn) |
| total_loss = 0.0 |
| batch_count = 0 |
| tp = fp = tn = fn = 0 |
| use_tiled = False |
|
|
| was_training = learn.model.training |
| learn.model.eval() |
| with torch.inference_mode(): |
| for batch_idx, record in enumerate(records): |
| if max_valid_batches is not None and batch_idx >= max_valid_batches: |
| break |
| logits = predict_record_logits_windowed( |
| learn.model, |
| record, |
| band_mean=band_mean, |
| band_std=band_std, |
| tile_size=tile_size, |
| overlap=overlap, |
| device=device, |
| ) |
| yb = open_mask_binary(record.label_path).long() |
| batch_loss = loss_func( |
| logits.unsqueeze(0).to(device), |
| yb.unsqueeze(0).to(device), |
| ).detach().float().item() |
| pred = logits.argmax(dim=0) |
| observed = yb != IGNORE_INDEX |
| total_loss += batch_loss |
| batch_count += 1 |
| tp += int(((pred == 1) & (yb == 1) & observed).sum().item()) |
| fp += int(((pred == 1) & (yb != 1) & observed).sum().item()) |
| tn += int(((pred != 1) & (yb != 1) & observed).sum().item()) |
| fn += int(((pred != 1) & (yb == 1) & observed).sum().item()) |
| use_tiled = True |
| del logits, yb, pred, observed |
| learn.model.train(was_training) |
|
|
| metric_summary = metrics_from_binary_confusion(tp=tp, fp=fp, tn=tn, fn=fn) |
| metric_summary.update( |
| { |
| "loss": float(total_loss / max(batch_count, 1)), |
| "used_tiled_inference": use_tiled, |
| "tp": tp, |
| "fp": fp, |
| "tn": tn, |
| "fn": fn, |
| } |
| ) |
| return metric_summary |
|
|
|
|
| def evaluate_threshold_sweep( |
| learn: Learner, |
| thresholds: Sequence[float], |
| tile_size: int, |
| overlap: int, |
| use_tta: bool = False, |
| tta_mode_count: int = 4, |
| force_tiled: bool = False, |
| max_valid_batches: Optional[int] = None, |
| ) -> pd.DataFrame: |
| device = next(learn.model.parameters()).device |
| records, band_mean, band_std = windowed_validation_context(learn) |
| counts = {name: torch.zeros(len(thresholds), dtype=torch.long) for name in ["tp", "fp", "tn", "fn"]} |
| eval_mode = "tta" if use_tta else "plain" |
| if force_tiled: |
| eval_mode = "{mode}_tiled".format(mode=eval_mode) |
| use_tiled = force_tiled |
|
|
| was_training = learn.model.training |
| learn.model.eval() |
| with torch.inference_mode(): |
| for batch_idx, record in enumerate(records): |
| if max_valid_batches is not None and batch_idx >= max_valid_batches: |
| break |
| logits = predict_record_logits_windowed( |
| learn.model, |
| record, |
| band_mean=band_mean, |
| band_std=band_std, |
| tile_size=tile_size, |
| overlap=overlap, |
| device=device, |
| use_tta=use_tta, |
| tta_mode_count=tta_mode_count, |
| ) |
| probs = F.softmax(logits, dim=0)[1] |
| labels = open_mask_binary(record.label_path) |
| observed = labels != IGNORE_INDEX |
| targ = labels == 1 |
| for threshold_idx, threshold in enumerate(thresholds): |
| pred = probs >= float(threshold) |
| counts["tp"][threshold_idx] += (pred & targ & observed).sum() |
| counts["fp"][threshold_idx] += (pred & ~targ & observed).sum() |
| counts["tn"][threshold_idx] += ((~pred) & ~targ & observed).sum() |
| counts["fn"][threshold_idx] += ((~pred) & targ & observed).sum() |
| use_tiled = True |
| del logits, probs, targ, labels, observed |
| learn.model.train(was_training) |
|
|
| records: List[dict] = [] |
| for idx, threshold in enumerate(thresholds): |
| tp = int(counts["tp"][idx].item()) |
| fp = int(counts["fp"][idx].item()) |
| tn = int(counts["tn"][idx].item()) |
| fn = int(counts["fn"][idx].item()) |
| metrics = metrics_from_binary_confusion(tp=tp, fp=fp, tn=tn, fn=fn) |
| records.append( |
| { |
| "eval_mode": eval_mode, |
| "threshold": float(threshold), |
| "used_tiled_inference": use_tiled, |
| **metrics, |
| "tp": tp, |
| "fp": fp, |
| "tn": tn, |
| "fn": fn, |
| } |
| ) |
| return pd.DataFrame(records) |
|
|
|
|
| def evaluate_tiles_at_threshold( |
| learn: Learner, |
| valid_records: Sequence[TripletRecord], |
| threshold: float, |
| tile_size: int, |
| overlap: int, |
| use_tta: bool = False, |
| tta_mode_count: int = 4, |
| force_tiled: bool = False, |
| max_valid_batches: Optional[int] = None, |
| ) -> pd.DataFrame: |
| device = next(learn.model.parameters()).device |
| _, band_mean, band_std = windowed_validation_context(learn) |
| rows: List[dict] = [] |
| eval_mode = "tta" if use_tta else "plain" |
| if force_tiled: |
| eval_mode = "{mode}_tiled".format(mode=eval_mode) |
|
|
| was_training = learn.model.training |
| learn.model.eval() |
| with torch.inference_mode(): |
| for batch_idx, record in enumerate(valid_records): |
| if max_valid_batches is not None and batch_idx >= max_valid_batches: |
| break |
| logits = predict_record_logits_windowed( |
| learn.model, |
| record, |
| band_mean=band_mean, |
| band_std=band_std, |
| tile_size=tile_size, |
| overlap=overlap, |
| device=device, |
| use_tta=use_tta, |
| tta_mode_count=tta_mode_count, |
| ) |
| probs = F.softmax(logits, dim=0)[1] |
| yb = open_mask_binary(record.label_path).long() |
| pred = probs >= threshold |
| observed = yb != IGNORE_INDEX |
| targ = yb == 1 |
| tp = int((pred & targ & observed).sum().item()) |
| fp = int((pred & ~targ & observed).sum().item()) |
| tn = int(((~pred) & ~targ & observed).sum().item()) |
| fn = int(((~pred) & targ & observed).sum().item()) |
| metrics = metrics_from_binary_confusion(tp=tp, fp=fp, tn=tn, fn=fn) |
| rows.append( |
| { |
| "tile_name": record.name, |
| "source_id": record.source_id, |
| "eval_mode": eval_mode, |
| "threshold": float(threshold), |
| "used_tiled_inference": True, |
| "water_pixels": int(targ.sum().item()), |
| **metrics, |
| "tp": tp, |
| "fp": fp, |
| "tn": tn, |
| "fn": fn, |
| } |
| ) |
| del logits, probs, yb, pred, targ, observed |
| learn.model.train(was_training) |
| return pd.DataFrame(rows) |
|
|
|
|
| def summarize_threshold_results(threshold_df: pd.DataFrame) -> dict: |
| summary: dict = {} |
| if threshold_df.empty: |
| return summary |
|
|
| for eval_mode in threshold_df["eval_mode"].unique(): |
| mode_df = threshold_df[threshold_df["eval_mode"] == eval_mode].sort_values( |
| ["water_iou", "dice", "threshold"], |
| ascending=[False, False, True], |
| ) |
| best = mode_df.iloc[0] |
| summary.update( |
| { |
| "{mode}_best_threshold".format(mode=eval_mode): round(float(best["threshold"]), 4), |
| "{mode}_best_water_iou".format(mode=eval_mode): round(float(best["water_iou"]), 6), |
| "{mode}_best_dice".format(mode=eval_mode): round(float(best["dice"]), 6), |
| "{mode}_best_precision".format(mode=eval_mode): round(float(best["precision"]), 6), |
| "{mode}_best_recall".format(mode=eval_mode): round(float(best["recall"]), 6), |
| } |
| ) |
|
|
| best_row = threshold_df.sort_values(["water_iou", "dice", "threshold"], ascending=[False, False, True]).iloc[0] |
| summary.update( |
| { |
| "best_eval_mode": str(best_row["eval_mode"]), |
| "best_eval_threshold": round(float(best_row["threshold"]), 4), |
| "best_eval_water_iou": round(float(best_row["water_iou"]), 6), |
| "best_eval_dice": round(float(best_row["dice"]), 6), |
| "best_eval_precision": round(float(best_row["precision"]), 6), |
| "best_eval_recall": round(float(best_row["recall"]), 6), |
| } |
| ) |
| return summary |
|
|
|
|
| def summarize_worst_tiles(tile_df: pd.DataFrame, worst_tile_count: int) -> pd.DataFrame: |
| if tile_df.empty: |
| return tile_df |
| return tile_df.sort_values( |
| ["water_iou", "precision", "recall", "tile_name"], |
| ascending=[True, True, True, True], |
| ).head(worst_tile_count) |
|
|
|
|
| def summarize_source_metrics(tile_df: pd.DataFrame) -> Dict[str, dict]: |
| if tile_df.empty or "source_id" not in tile_df.columns: |
| return {} |
| summaries: Dict[str, dict] = {} |
| for source_id, source_df in tile_df.groupby("source_id", dropna=False): |
| source_name = str(source_id) if str(source_id) else "unspecified" |
| tp = int(source_df["tp"].sum()) |
| fp = int(source_df["fp"].sum()) |
| tn = int(source_df["tn"].sum()) |
| fn = int(source_df["fn"].sum()) |
| metrics = metrics_from_binary_confusion(tp=tp, fp=fp, tn=tn, fn=fn) |
| metrics["water_f1"] = float((2 * tp) / max((2 * tp) + fp + fn, 1)) |
| summaries[source_name] = { |
| "tile_count": int(len(source_df)), |
| **{key: round(float(value), 6) for key, value in metrics.items()}, |
| "tp": tp, |
| "fp": fp, |
| "tn": tn, |
| "fn": fn, |
| } |
| return summaries |
|
|
|
|
| def run_training_schedule( |
| learn: Learner, |
| args: argparse.Namespace, |
| schedule: Sequence[Tuple[str, int, object]], |
| epoch_timer: Optional[EpochTimerCallback] = None, |
| ) -> pd.DataFrame: |
| training_history = [] |
| for phase_name, epochs, lr_max in schedule: |
| apply_phase_freeze(learn, phase_name) |
| training_history.append( |
| run_phase( |
| learn=learn, |
| args=args, |
| phase=phase_name, |
| epochs=epochs, |
| lr_max=lr_max, |
| epoch_timer=epoch_timer, |
| ) |
| ) |
| if epochs > 0: |
| learn.load(args.best_model_name, with_opt=False) |
| save_json(args.artifact_dir / "global_best.json", getattr(learn, "global_best_where", {})) |
| return pd.concat(training_history, ignore_index=True) if training_history else pd.DataFrame(columns=["phase"]) |
|
|
|
|
| def effective_epoch_count(args: argparse.Namespace) -> int: |
| return int(args.freeze_epochs + args.partial_unfreeze_epochs + args.finetune_epochs + args.finetune_tail_epochs) |
|
|
|
|
| def probe_batch_size( |
| args: argparse.Namespace, |
| train_records: Sequence[TripletRecord], |
| valid_records: Sequence[TripletRecord], |
| band_mean: np.ndarray, |
| band_std: np.ndarray, |
| device: torch.device, |
| ) -> dict: |
| if device.type != "cuda": |
| return { |
| "selected_batch_size": args.batch_size, |
| "probe_results": [], |
| "probe_skipped": "CUDA unavailable", |
| } |
|
|
| results: List[dict] = [] |
| selected_batch_size = None |
| train_probe_records = list(train_records[: max(max(args.probe_batch_sizes), 16)]) |
| valid_probe_records = list(valid_records[:1]) |
|
|
| for candidate in args.probe_batch_sizes: |
| candidate_args = copy_module.deepcopy(args) |
| candidate_args.batch_size = candidate |
| candidate_args.num_workers = min(candidate_args.num_workers, 2) |
| try: |
| dls = build_dataloaders( |
| candidate_args, |
| train_records=train_probe_records, |
| valid_records=valid_probe_records, |
| band_mean=band_mean, |
| band_std=band_std, |
| device=device, |
| ) |
| learn = build_learner(candidate_args, dls, valid_records=valid_probe_records) |
| learn.model.train() |
| optimizer = torch.optim.AdamW(learn.model.parameters(), lr=args.warmup_lr) |
| xb, yb = dls.one_batch() |
| optimizer.zero_grad(set_to_none=True) |
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
| torch.cuda.reset_peak_memory_stats(device) |
| start = time.perf_counter() |
| if args.use_fp16: |
| scaler = torch.cuda.amp.GradScaler(enabled=True) |
| with torch.cuda.amp.autocast(enabled=True): |
| pred = learn.model(xb) |
| loss = learn.loss_func(pred, yb) |
| scaler.scale(loss).backward() |
| scaler.step(optimizer) |
| scaler.update() |
| else: |
| pred = learn.model(xb) |
| loss = learn.loss_func(pred, yb) |
| loss.backward() |
| optimizer.step() |
| torch.cuda.synchronize(device) |
| elapsed = time.perf_counter() - start |
| peak_mem_gb = torch.cuda.max_memory_allocated(device) / (1024**3) |
| result = { |
| "batch_size": candidate, |
| "success": True, |
| "loss": float(loss.detach().float().item()), |
| "step_seconds": round(elapsed, 4), |
| "max_memory_gb": round(peak_mem_gb, 4), |
| } |
| results.append(result) |
| selected_batch_size = candidate |
| del pred, loss, xb, yb, learn, dls, optimizer |
| gc.collect() |
| torch.cuda.empty_cache() |
| break |
| except RuntimeError as error: |
| if not is_cuda_oom(error): |
| raise |
| results.append({"batch_size": candidate, "success": False, "error": "cuda_oom"}) |
| gc.collect() |
| torch.cuda.empty_cache() |
|
|
| if selected_batch_size is None: |
| raise RuntimeError("Batch size probe failed for all candidates: {candidates}".format(candidates=args.probe_batch_sizes)) |
|
|
| probe_summary = {"selected_batch_size": selected_batch_size, "probe_results": results} |
| print(probe_summary) |
| return probe_summary |
|
|
|
|
| def maybe_override_batch_size_from_summary(args: argparse.Namespace) -> None: |
| if args.benchmark_summary_path is None: |
| return |
| if not args.benchmark_summary_path.exists(): |
| raise FileNotFoundError( |
| "Benchmark summary path does not exist: {path}".format(path=args.benchmark_summary_path) |
| ) |
| summary = json.loads(args.benchmark_summary_path.read_text()) |
| selected = summary.get("selected_batch_size") |
| if not selected: |
| raise ValueError( |
| "Benchmark summary at {path} is missing selected_batch_size".format(path=args.benchmark_summary_path) |
| ) |
| args.batch_size = int(selected) |
| print( |
| { |
| "batch_size_override": args.batch_size, |
| "benchmark_summary_path": str(args.benchmark_summary_path), |
| } |
| ) |
|
|
|
|
| def build_audit_summary( |
| args: argparse.Namespace, |
| train_records: Sequence[TripletRecord], |
| valid_records: Sequence[TripletRecord], |
| band_mean: np.ndarray, |
| band_std: np.ndarray, |
| dls: DataLoaders, |
| ) -> dict: |
| sample_image, sample_mask = load_center_triplet_window(train_records[0], args.crop_size) |
| batch_images, batch_masks = dls.one_batch() |
| return { |
| "paired_tiles": len(train_records) + len(valid_records), |
| "train_tiles": len(train_records), |
| "valid_tiles": len(valid_records), |
| "sample_image_shape": tuple(sample_image.shape), |
| "sample_mask_shape": tuple(sample_mask.shape), |
| "sample_mask_values": sorted(torch.unique(sample_mask).tolist()), |
| "band_mean_s1": band_mean[:3].round(6).tolist(), |
| "band_std_s1": band_std[:3].round(6).tolist(), |
| "band_mean_aef_first5": band_mean[3:8].round(6).tolist() if using_aef(args) else [], |
| "band_std_aef_first5": band_std[3:8].round(6).tolist() if using_aef(args) else [], |
| "train_batch_images": tuple(batch_images.shape), |
| "train_batch_masks": tuple(batch_masks.shape), |
| "model_family": args.model_family, |
| "segformer_variant": args.segformer_variant if args.model_family == "segformer" else None, |
| "loss": args.loss, |
| "loss_aux_weight": args.loss_aux_weight, |
| "optimizer": args.optimizer, |
| "schedule_policy": args.schedule_policy, |
| "tuning_strategy": args.tuning_strategy, |
| "augmentation_rotate_max_deg": getattr(dls.train.dataset, "rotate_max_deg", None), |
| "augmentation_scale_jitter": [ |
| getattr(dls.train.dataset, "scale_jitter_min", None), |
| getattr(dls.train.dataset, "scale_jitter_max", None), |
| ], |
| "augmentation_s1_gain_jitter": getattr(dls.train.dataset, "s1_gain_jitter", None), |
| "augmentation_s1_bias_jitter": getattr(dls.train.dataset, "s1_bias_jitter", None), |
| "augmentation_gaussian_noise_std": getattr(dls.train.dataset, "gaussian_noise_std", None), |
| "positive_crop_p": getattr(dls.train.dataset, "positive_crop_p", None), |
| "min_water_fraction": getattr(dls.train.dataset, "min_water_fraction", None), |
| "positive_crop_attempts": getattr(dls.train.dataset, "positive_crop_attempts", None), |
| "target_source_id": args.target_source_id, |
| "target_source_fraction": args.target_source_fraction, |
| "use_permanent_prior": using_prior(args), |
| "prior_dropout_p": args.prior_dropout_p, |
| "aef_dropout_p": getattr(args, "aef_dropout_p", 0.0), |
| } |
|
|
|
|
| def run_benchmark_mode( |
| args: argparse.Namespace, |
| device: torch.device, |
| train_records: Sequence[TripletRecord], |
| valid_records: Sequence[TripletRecord], |
| band_mean: np.ndarray, |
| band_std: np.ndarray, |
| ) -> dict: |
| planned_full_epochs = effective_epoch_count(args) |
| probe_summary = probe_batch_size(args, train_records, valid_records, band_mean, band_std, device) |
| args.batch_size = int(probe_summary["selected_batch_size"]) |
| if args.max_benchmark_batch_size is not None: |
| args.batch_size = min(args.batch_size, args.max_benchmark_batch_size) |
| probe_summary["selected_batch_size"] = args.batch_size |
| train_records = subset_records(train_records, args.max_train_batches, args.batch_size) |
| valid_records = subset_records(valid_records, args.max_valid_batches, 1) |
| dls = build_dataloaders(args, train_records, valid_records, band_mean, band_std, device) |
| print(build_audit_summary(args, train_records, valid_records, band_mean, band_std, dls)) |
|
|
| learn = build_learner(args, dls, valid_records=valid_records) |
| maybe_warm_start(learn, args) |
| epoch_timer = EpochTimerCallback() |
| history_df = run_training_schedule( |
| learn=learn, |
| args=args, |
| schedule=build_schedule(args), |
| epoch_timer=epoch_timer, |
| ) |
| history_path = args.artifact_dir / args.history_csv |
| history_df.to_csv(history_path, index=False) |
|
|
| validation = validate_model( |
| learn, |
| tile_size=args.crop_size, |
| overlap=args.tile_overlap, |
| force_tiled=args.force_tiled_validation, |
| max_valid_batches=args.max_valid_batches, |
| ) |
| seconds_per_epoch = sum(epoch_timer.epoch_times) / max(len(epoch_timer.epoch_times), 1) |
| samples_per_second = ( |
| (len(train_records) / seconds_per_epoch) if seconds_per_epoch > 0 else None |
| ) |
| benchmark_summary = { |
| "mode": "benchmark", |
| "model_family": args.model_family, |
| "segformer_variant": args.segformer_variant if args.model_family == "segformer" else None, |
| "selected_batch_size": args.batch_size, |
| "probe_results": probe_summary["probe_results"], |
| "epoch_times_seconds": [round(v, 4) for v in epoch_timer.epoch_times], |
| "avg_seconds_per_epoch": round(seconds_per_epoch, 4), |
| "samples_per_second": None if samples_per_second is None else round(samples_per_second, 4), |
| "projected_full_train_seconds": round(seconds_per_epoch * planned_full_epochs, 1), |
| "projected_full_train_hours": round((seconds_per_epoch * planned_full_epochs) / 3600.0, 3), |
| "train_records_used": len(train_records), |
| "valid_records_used": len(valid_records), |
| "validation": validation, |
| "history_csv": str(history_path), |
| "loss": args.loss, |
| "selection_metric": args.selection_metric, |
| "tuning_strategy": args.tuning_strategy, |
| "force_tiled_validation": bool(args.force_tiled_validation), |
| "split_seed": args.seed if args.split_seed is None else args.split_seed, |
| "train_subset_seed": args.seed if args.train_subset_seed is None else args.train_subset_seed, |
| "max_train_records": args.max_train_records, |
| "target_source_id": args.target_source_id, |
| "target_source_fraction": args.target_source_fraction, |
| "use_permanent_prior": using_prior(args), |
| "prior_dropout_p": args.prior_dropout_p, |
| "aef_dropout_p": getattr(args, "aef_dropout_p", 0.0), |
| } |
| save_json(args.artifact_dir / args.benchmark_summary_json, benchmark_summary) |
| print(benchmark_summary) |
| return benchmark_summary |
|
|
|
|
| def run_train_mode( |
| args: argparse.Namespace, |
| device: torch.device, |
| train_records: Sequence[TripletRecord], |
| valid_records: Sequence[TripletRecord], |
| band_mean: np.ndarray, |
| band_std: np.ndarray, |
| ) -> dict: |
| maybe_override_batch_size_from_summary(args) |
| train_records = subset_records(train_records, args.max_train_batches, args.batch_size) |
| valid_records = subset_records(valid_records, args.max_valid_batches, 1) |
| dls = build_dataloaders(args, train_records, valid_records, band_mean, band_std, device) |
| print(build_audit_summary(args, train_records, valid_records, band_mean, band_std, dls)) |
|
|
| learn = build_learner(args, dls, valid_records=valid_records) |
| maybe_warm_start(learn, args) |
| history_df = run_training_schedule(learn=learn, args=args, schedule=build_schedule(args)) |
| history_path = args.artifact_dir / args.history_csv |
| history_df.to_csv(history_path, index=False) |
| print(history_df.tail()) |
|
|
| if using_aef(args): |
| bottleneck_report = inspect_bottleneck(learn.model, args.artifact_dir) |
| print({"bottleneck_weight_norms": bottleneck_report["weight_norm_per_proj_band"]}) |
|
|
| validation = validate_model( |
| learn, |
| tile_size=args.crop_size, |
| overlap=args.tile_overlap, |
| force_tiled=args.force_tiled_validation, |
| max_valid_batches=args.max_valid_batches, |
| ) |
| threshold_results = [ |
| evaluate_threshold_sweep( |
| learn, |
| thresholds=args.eval_thresholds, |
| tile_size=args.crop_size, |
| overlap=args.tile_overlap, |
| use_tta=False, |
| force_tiled=args.force_tiled_validation, |
| max_valid_batches=args.max_valid_batches, |
| ) |
| ] |
| if args.use_tta_eval: |
| threshold_results.append( |
| evaluate_threshold_sweep( |
| learn, |
| thresholds=args.eval_thresholds, |
| tile_size=args.crop_size, |
| overlap=args.tile_overlap, |
| use_tta=True, |
| tta_mode_count=args.tta_eval_mode_count, |
| force_tiled=args.force_tiled_validation, |
| max_valid_batches=args.max_valid_batches, |
| ) |
| ) |
| threshold_df = pd.concat(threshold_results, ignore_index=True) |
| threshold_path = args.artifact_dir / args.threshold_sweep_csv |
| threshold_df.to_csv(threshold_path, index=False) |
| print( |
| threshold_df.sort_values( |
| ["eval_mode", "water_iou", "precision", "threshold"], |
| ascending=[True, False, False, True], |
| ).groupby("eval_mode").head(5) |
| ) |
|
|
| best_eval_row = threshold_df.sort_values( |
| ["water_iou", "dice", "threshold"], |
| ascending=[False, False, True], |
| ).iloc[0] |
| best_eval_mode = str(best_eval_row["eval_mode"]) |
| tile_metrics_df = evaluate_tiles_at_threshold( |
| learn, |
| valid_records=valid_records, |
| threshold=float(best_eval_row["threshold"]), |
| tile_size=args.crop_size, |
| overlap=args.tile_overlap, |
| use_tta=best_eval_mode.startswith("tta"), |
| tta_mode_count=args.tta_eval_mode_count, |
| force_tiled=best_eval_mode.endswith("_tiled"), |
| max_valid_batches=args.max_valid_batches, |
| ) |
| tile_metrics_path = args.artifact_dir / args.tile_metrics_csv |
| tile_metrics_df.to_csv(tile_metrics_path, index=False) |
| worst_tiles_df = summarize_worst_tiles(tile_metrics_df, worst_tile_count=args.worst_tile_count) |
| worst_tiles_path = args.artifact_dir / args.worst_tiles_csv |
| worst_tiles_df.to_csv(worst_tiles_path, index=False) |
|
|
| summary = { |
| "model_family": args.model_family, |
| "segformer_variant": args.segformer_variant if args.model_family == "segformer" else None, |
| "final_valid_loss": round(float(validation["loss"]), 6), |
| "final_valid_jaccard": round(float(validation["jaccard"]), 6), |
| "final_valid_dice": round(float(validation["dice"]), 6), |
| "final_valid_water_iou": round(float(validation["water_iou"]), 6), |
| "final_valid_precision": round(float(validation["precision"]), 6), |
| "final_valid_recall": round(float(validation["recall"]), 6), |
| "used_tiled_validation": bool(validation["used_tiled_inference"]), |
| "best_model_path": str(args.artifact_dir / "models" / "{name}.pth".format(name=args.best_model_name)), |
| "history_csv": str(history_path), |
| "threshold_sweep_csv": str(threshold_path), |
| "tile_metrics_csv": str(tile_metrics_path), |
| "worst_tiles_csv": str(worst_tiles_path), |
| "loss": args.loss, |
| "selection_metric": args.selection_metric, |
| "tuning_strategy": args.tuning_strategy, |
| "force_tiled_validation": bool(args.force_tiled_validation), |
| "split_seed": args.seed if args.split_seed is None else args.split_seed, |
| "train_subset_seed": args.seed if args.train_subset_seed is None else args.train_subset_seed, |
| "max_train_records": args.max_train_records, |
| "target_source_id": args.target_source_id, |
| "target_source_fraction": args.target_source_fraction, |
| "use_permanent_prior": using_prior(args), |
| "prior_dropout_p": args.prior_dropout_p, |
| "aef_dropout_p": getattr(args, "aef_dropout_p", 0.0), |
| "train_records_used": len(train_records), |
| "valid_records_used": len(valid_records), |
| "source_metrics_at_best_threshold": summarize_source_metrics(tile_metrics_df), |
| } |
| if using_aef(args): |
| summary["bottleneck_proj_shape"] = "{aef}->{proj}".format( |
| aef=args.n_aef_bands, proj=args.n_proj_bands |
| ) |
| else: |
| summary["input_bands"] = args.n_s1_bands |
| summary.update(summarize_threshold_results(threshold_df)) |
| if not worst_tiles_df.empty: |
| summary["worst_tile_examples"] = worst_tiles_df[ |
| ["tile_name", "water_iou", "precision", "recall"] |
| ].round(6).to_dict(orient="records") |
| save_json(args.artifact_dir / args.eval_summary_json, summary) |
| print(summary) |
| return summary |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| device = setup_environment(args) |
|
|
| if args.triplet_manifest is not None: |
| triplets = collect_triplets_from_manifest( |
| args.triplet_manifest, |
| args.triplet_sample_id, |
| expect_aef=using_aef(args), |
| expect_prior=using_prior(args), |
| ) |
| else: |
| triplets = collect_triplets( |
| args.input_dir, |
| args.aef_dir if using_aef(args) else None, |
| args.label_dir, |
| ) |
| split_seed = args.seed if args.split_seed is None else args.split_seed |
| train_subset_seed = args.seed if args.train_subset_seed is None else args.train_subset_seed |
| if args.fixed_split_manifest is not None: |
| full_train_records, valid_records = split_triplets_from_manifest( |
| triplets, args.fixed_split_manifest |
| ) |
| split_protocol = "fixed_manifest" |
| else: |
| full_train_records, valid_records = split_triplets( |
| triplets, valid_pct=args.valid_pct, seed=split_seed |
| ) |
| split_protocol = "seeded" |
| train_records = limit_train_records( |
| full_train_records, |
| max_records=args.max_train_records, |
| subset_seed=train_subset_seed, |
| ) |
| print( |
| { |
| "split_seed": split_seed, |
| "split_protocol": split_protocol, |
| "fixed_split_manifest": None if args.fixed_split_manifest is None else str(args.fixed_split_manifest), |
| "train_subset_seed": train_subset_seed, |
| "full_train_tiles": len(full_train_records), |
| "train_tiles_after_record_cap": len(train_records), |
| "valid_tiles": len(valid_records), |
| "max_train_records": args.max_train_records, |
| } |
| ) |
| write_split_manifest(args.artifact_dir, train_records, valid_records) |
| band_mean, band_std = get_band_stats( |
| train_records, |
| args.artifact_dir, |
| args.stats_sample_size, |
| source_stats_path=resolve_band_stats_source(args), |
| expected_bands=args.n_s1_bands + args.n_aef_bands, |
| ) |
|
|
| if args.mode == "benchmark": |
| run_benchmark_mode(args, device, train_records, valid_records, band_mean, band_std) |
| else: |
| run_train_mode(args, device, train_records, valid_records, band_mean, band_std) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
