diff --git a/models/dinov3s_buildings/pipeline.py b/models/dinov3s_buildings/pipeline.py index 7d17709..ba87c35 100644 --- a/models/dinov3s_buildings/pipeline.py +++ b/models/dinov3s_buildings/pipeline.py @@ -194,6 +194,116 @@ def evaluate_model( return metrics +THRESHOLD_GRID = [round(0.1 + 0.05 * i, 2) for i in range(17)] + + +def _burn_masks(dataset_chips: str, dataset_labels: str) -> tuple[Path, Path]: + """Burn GT polygons for all chips; returns (chips_dir, masks_dir).""" + from dinov3_hot.paths import resolve_labels_geojson + from geomltoolkits.raster.burn import burn_labels + + chips_dir = resolve_directory(dataset_chips, "*.tif*") + labels_geojson = resolve_labels_geojson(Path(resolve_directory(dataset_labels, "*.geojson"))) + masks_dir = Path(tempfile.mkdtemp()) / "masks" + burn_labels( + labels_path=str(labels_geojson), + chips_dir=str(chips_dir), + output_dir=str(masks_dir), + burn_value=255, + ) + return chips_dir, masks_dir + + +def _pooled_probs_and_targets( + net: Any, chips_dir: Path, masks_dir: Path, chip_names: list[str], split_info: dict[str, Any], device: str +) -> tuple[Any, Any]: + """Flattened mask-head probabilities and binary GT pixels for the given chips.""" + import numpy as np + import rasterio + from dinov3_hot.tune import cache_val_forwards + + cache = cache_val_forwards( + net, chips_dir, chip_names, + mean=split_info["norm_mean"], std=split_info["norm_std"], device=device, + ) + probs, targets = [], [] + for entry in cache: + probs.append(entry["mask_prob"].ravel()) + with rasterio.open(masks_dir / entry["name"]) as src: + targets.append(src.read(1).ravel() > 0) + return np.concatenate(probs), np.concatenate(targets) + + +@step +def calibrate_threshold( + trained_model: Any, + dataset_chips: str, + dataset_labels: str, + hyperparameters: dict[str, Any], + split_info: dict[str, Any], +) -> Annotated[dict[str, Any], "calibrated_threshold"]: + """Select the mask confidence_threshold by a deterministic val sweep, + decoupled from the Optuna post-process search. + + 17-point pixel-F1 sweep over 0.10-0.90 on the spatial val chips; falls + back to rate-matching on the train chips (threshold at which the + predicted positive-pixel fraction matches the labeled fraction - needs + no held-out data) when val has fewer than two chips or no positive + pixels. The result seeds tune_postprocess's defaults, which apply + verbatim whenever the Optuna search is skipped (val < 8 chips or + trials disabled) - exactly the small-dataset case where the catalog + constant used to be served unchanged. + """ + import numpy as np + import torch + + default = float( + hyperparameters.get("confidence_threshold", DEFAULT_INFERENCE_PARAMS["confidence_threshold"]) + ) + result: dict[str, Any] = {"confidence_threshold": default, "method": "default", "val_f1": None} + if not hyperparameters.get("calibrate_threshold", True): + log_metadata(metadata={"fair/threshold_calibration": result}) + return result + + chips_dir, masks_dir = _burn_masks(dataset_chips, dataset_labels) + device = "cuda" if torch.cuda.is_available() else "cpu" + val_names = split_info["val_chip_names"] + + probs = targets = None + if len(val_names) >= 2: + probs, targets = _pooled_probs_and_targets( + trained_model, chips_dir, masks_dir, val_names, split_info, device + ) + if probs is not None and targets.any(): + curve = [] + for t in THRESHOLD_GRID: + pred = probs >= t + tp = int((pred & targets).sum()) + fp = int((pred & ~targets).sum()) + fn = int((~pred & targets).sum()) + curve.append(2 * tp / max(2 * tp + fp + fn, 1)) + best = max(range(len(THRESHOLD_GRID)), key=curve.__getitem__) + result = { + "confidence_threshold": THRESHOLD_GRID[best], + "method": "val_sweep", + "val_f1": curve[best], + } + else: + probs, targets = _pooled_probs_and_targets( + trained_model, chips_dir, masks_dir, split_info["train_chip_names"], split_info, device + ) + pos_frac = float(targets.mean()) + if pos_frac > 0: + result = { + "confidence_threshold": float(np.quantile(probs, 1.0 - pos_frac)), + "method": "rate_match", + "val_f1": None, + } + + log_metadata(metadata={"fair/threshold_calibration": result}) + return result + + @step def tune_postprocess( trained_model: Any, @@ -201,8 +311,16 @@ def tune_postprocess( dataset_labels: str, hyperparameters: dict[str, Any], split_info: dict[str, Any], + calibrated_threshold: dict[str, Any] | None = None, ) -> Annotated[dict[str, Any], "recommended_inference_params"]: - """Optuna over post-process params via `dinov3_hot.tune.tune_postprocess_run`.""" + """Optuna over post-process params via `dinov3_hot.tune.tune_postprocess_run`. + + Defaults are seeded with the calibrated confidence_threshold, so the + skipped-search path (val < 8 chips or trials disabled) serves the + calibrated value instead of the catalog constant. When the search runs + it still tunes the threshold jointly within [0.3, 0.8]; constraining + that space to the calibrated value would need a dinov3_hot change. + """ from dinov3_hot.paths import resolve_labels_geojson from dinov3_hot.tune import tune_postprocess_run @@ -210,6 +328,10 @@ def tune_postprocess( chips_dir = resolve_directory(dataset_chips, "*.tif*") labels_geojson = resolve_labels_geojson(Path(resolve_directory(dataset_labels, "*.geojson"))) + defaults = dict(DEFAULT_INFERENCE_PARAMS) + if calibrated_threshold and calibrated_threshold.get("method") != "default": + defaults["confidence_threshold"] = float(calibrated_threshold["confidence_threshold"]) + result = tune_postprocess_run( trained_model, chips_dir, @@ -219,7 +341,16 @@ def tune_postprocess( std=split_info["norm_std"], n_trials=n_trials, seed=int(split_info["seed"]), - default_params=DEFAULT_INFERENCE_PARAMS, + default_params=defaults, + ) + log_metadata( + metadata={ + "fair/tune_postprocess": { + "skipped": result.get("skipped"), + "calibrated_threshold": defaults["confidence_threshold"], + "final_threshold": result["best_params"].get("confidence_threshold"), + } + } ) return result["best_params"] @@ -309,12 +440,20 @@ def training_pipeline( split_info=split_info, num_classes=num_classes, ) + calibrated = calibrate_threshold( + trained_model=trained, + dataset_chips=dataset_chips, + dataset_labels=dataset_labels, + hyperparameters=hyperparameters, + split_info=split_info, + ) tune_postprocess( trained_model=trained, dataset_chips=dataset_chips, dataset_labels=dataset_labels, hyperparameters=hyperparameters, split_info=split_info, + calibrated_threshold=calibrated, ) export_onnx(trained_model=trained, hyperparameters=hyperparameters, num_classes=num_classes) diff --git a/models/dinov3s_buildings/stac-item.json b/models/dinov3s_buildings/stac-item.json index 7a8bc4d..25c1185 100644 --- a/models/dinov3s_buildings/stac-item.json +++ b/models/dinov3s_buildings/stac-item.json @@ -118,6 +118,7 @@ "training.freeze_encoder": true, "training.norm_stats": "hot_global", "training.tune_postprocess_trials": 30, + "training.calibrate_threshold": true, "training.sample_fraction": 1.0, "inference.confidence_threshold": 0.4371, "inference.simplify_m": 0.9626, @@ -147,7 +148,8 @@ {"key": "inner_val_frac", "type": "float", "default": 0.1, "min": 0.05, "max": 0.5, "description": "Within-train holdout that upstream finetune uses for Lightning's per-epoch val_loss/val_iou and early stopping. Independent of the catalog-facing spatial val."}, {"key": "freeze_encoder", "type": "bool", "default": true, "description": "Keep the DINOv3 encoder frozen (recommended)"}, {"key": "norm_stats", "type": "str", "default": "hot_global", "values": ["hot_global", "dataset"], "description": "Input normalisation source. 'hot_global' uses mean/std from the hotosm/vhr-building-segmentation dataset (matches pretraining). 'dataset' computes per-channel mean/std from the user's finetune chips, pick this when your imagery differs substantially from HOT global imagery."}, - {"key": "confidence_threshold", "type": "float", "default": 0.4371, "min": 0.0, "max": 1.0, "description": "Sigmoid threshold for building mask"}, + {"key": "calibrate_threshold", "type": "bool", "default": true, "description": "Select the mask confidence_threshold by a deterministic 17-point pixel-F1 sweep on the val chips after training (rate-matching on the train chips when val is tiny). Seeds the tune_postprocess defaults, which apply verbatim whenever the Optuna search is skipped."}, + {"key": "confidence_threshold", "type": "float", "default": 0.4371, "min": 0.0, "max": 1.0, "description": "Sigmoid threshold for building mask; used as-is only when calibrate_threshold is false and the post-process search does not retune it"}, {"key": "simplify_m", "type": "float", "default": 0.9626, "min": 0.0, "max": 10.0, "description": "Douglas-Peucker tolerance in metres (EPSG:3857). 0 disables simplification."}, {"key": "regularize_area_threshold", "type": "float", "default": 0.4949, "min": 0.0, "max": 1.0, "description": "Minimum polygon-area / MBR-area ratio for the rectangle-substitution step. Lower = more aggressive squaring of corners."}, {"key": "regularize_overlap_tol_m2", "type": "float", "default": 3.9251, "min": 0.0, "max": 100.0, "description": "Maximum new neighbour overlap (m^2) allowed when substituting a polygon with its MBR. Keeps adjacent buildings from being merged."}, diff --git a/models/unet_segmentation/pipeline.py b/models/unet_segmentation/pipeline.py index c1baab2..2e411e6 100644 --- a/models/unet_segmentation/pipeline.py +++ b/models/unet_segmentation/pipeline.py @@ -131,9 +131,14 @@ def predict(session: Any, input_images: str, params: dict[str, Any]) -> dict[str batch, transform, crs = _preprocess_onnx_image(img_path) logits = session.run(None, {input_name: batch})[0] probs = _softmax(logits[0], axis=0) - mask = probs.argmax(axis=0) - top_prob = probs.max(axis=0) - mask = np.where(top_prob >= confidence_threshold, mask, 0) + # Threshold the most likely foreground class directly. The previous + # rule (argmax, then mask where max-prob < threshold) could only move + # the operating point above 0.5 in the binary case; thresholds below + # 0.5 were no-ops. At the 0.5 default both rules coincide. + fg_probs = probs[min_class_value:] + fg_class = fg_probs.argmax(axis=0) + min_class_value + fg_prob = fg_probs.max(axis=0) + mask = np.where(fg_prob >= confidence_threshold, fg_class, 0) features.extend(_vectorize_segmentation_mask(mask, transform, crs, min_class_value)) return _build_feature_collection(features) @@ -227,6 +232,110 @@ def _train_step( return loss.item() +THRESHOLD_GRID = [round(0.1 + 0.05 * i, 2) for i in range(17)] + + +def _collect_foreground_probs(model: Any, loader: Any, device: str) -> tuple[Any, Any]: + """Pooled sigmoid(foreground - background logit) and binary targets, flattened.""" + import torch + + probs, targets = [], [] + model.eval() + with torch.no_grad(): + for batch in loader: + images, masks = preprocess(batch) + logits = model(images.to(device)) + probs.append(torch.sigmoid(logits[:, 1] - logits[:, 0]).cpu()) + targets.append(masks > 0) + return torch.cat(probs).flatten(), torch.cat(targets).flatten() + + +def _f1_at_threshold(probs: Any, targets: Any, threshold: float) -> float: + pred = probs >= threshold + tp = (pred & targets).sum().item() + fp = (pred & ~targets).sum().item() + fn = (~pred & targets).sum().item() + return 2 * tp / max(2 * tp + fp + fn, 1) + + +@step +def calibrate_threshold( + trained_model: Any, + dataset_chips: str, + dataset_labels: str, + hyperparameters: dict[str, Any], + split_info: dict[str, Any], + num_classes: int = 2, +) -> Annotated[dict[str, Any], "recommended_inference_params"]: + """Select inference.confidence_threshold on the validation split. + + Deterministic 17-point F1 sweep over 0.10-0.90. Falls back to + rate-matching on the train sampler (threshold at which the predicted + positive-pixel fraction equals the labeled positive fraction) when the + validation split contains no positive pixels, and to the configured + default when neither split has positives. Binary models only: with more + than two classes a single foreground threshold is ill-defined, so the + default passes through unchanged. + """ + import torch + + default = float(hyperparameters.get("confidence_threshold", 0.5)) + result: dict[str, Any] = {"confidence_threshold": default, "method": "default", "val_f1": None} + if num_classes != 2 or not hyperparameters.get("calibrate_threshold", True): + log_metadata(metadata={"fair/threshold_calibration": result}) + return result + + chip_size = hyperparameters.get("chip_size", 256) + batch_size = hyperparameters.get("batch_size", 4) + samples_per_epoch = hyperparameters.get("samples_per_epoch", 50) + sample_fraction = hyperparameters.get("sample_fraction", 1.0) + device = _get_device() + model = trained_model.to(device) + + val_loader = _build_dataset( + dataset_chips, + dataset_labels, + chip_size, + length=0, + batch_size=batch_size, + split="val", + seed=split_info["seed"], + sample_fraction=sample_fraction, + ) + probs, targets = _collect_foreground_probs(model, val_loader, device) + + if bool(targets.any()): + curve = [_f1_at_threshold(probs, targets, t) for t in THRESHOLD_GRID] + best = max(range(len(THRESHOLD_GRID)), key=curve.__getitem__) + result = { + "confidence_threshold": THRESHOLD_GRID[best], + "method": "val_sweep", + "val_f1": curve[best], + } + else: + train_loader = _build_dataset( + dataset_chips, + dataset_labels, + chip_size, + length=samples_per_epoch, + batch_size=batch_size, + split="train", + seed=split_info["seed"], + sample_fraction=sample_fraction, + ) + probs, targets = _collect_foreground_probs(model, train_loader, device) + pos_frac = targets.float().mean().item() + if pos_frac > 0: + result = { + "confidence_threshold": float(torch.quantile(probs, 1.0 - pos_frac)), + "method": "rate_match", + "val_f1": None, + } + + log_metadata(metadata={"fair/threshold_calibration": result}) + return result + + @step def split_dataset( dataset_chips: str, @@ -498,6 +607,14 @@ def training_pipeline( split_info=split_info, num_classes=num_classes, ) + calibrate_threshold( + trained_model=trained_model, + dataset_chips=dataset_chips, + dataset_labels=dataset_labels, + hyperparameters=hyperparameters, + split_info=split_info, + num_classes=num_classes, + ) export_onnx( trained_model=trained_model, hyperparameters=hyperparameters, diff --git a/models/unet_segmentation/stac-item.json b/models/unet_segmentation/stac-item.json index 5d34dee..d4da646 100644 --- a/models/unet_segmentation/stac-item.json +++ b/models/unet_segmentation/stac-item.json @@ -184,6 +184,7 @@ "training.scheduler": "cosine", "training.max_grad_norm": 1.0, "training.freeze_encoder": true, + "training.calibrate_threshold": true, "inference.min_class_value": 1, "inference.confidence_threshold": 0.5 }, @@ -295,13 +296,19 @@ "default": true, "description": "Freeze pretrained encoder weights during finetuning" }, + { + "key": "calibrate_threshold", + "type": "bool", + "default": true, + "description": "Select inference.confidence_threshold by a deterministic validation-split F1 sweep after training (rate-matching fallback when validation has no positive pixels)" + }, { "key": "confidence_threshold", "type": "float", "default": 0.5, "min": 0.0, "max": 1.0, - "description": "Minimum per-pixel softmax probability to retain a predicted class at inference" + "description": "Minimum per-pixel foreground probability to retain a predicted class at inference; used as-is only when calibrate_threshold is false" }, { "key": "min_class_value", diff --git a/models/unet_segmentation/tests/test_steps.py b/models/unet_segmentation/tests/test_steps.py index 9c6a874..28bc61a 100644 --- a/models/unet_segmentation/tests/test_steps.py +++ b/models/unet_segmentation/tests/test_steps.py @@ -100,6 +100,53 @@ def test_evaluate_model( assert 0.0 <= metrics["mean_iou"] <= 1.0 +def test_calibrate_threshold( + toy_chips: Path, + toy_labels: Path, + base_hyperparameters: dict[str, Any], + pretrained_weights: str, +) -> None: + from models.unet_segmentation.pipeline import calibrate_threshold, split_dataset, train_model + + split_info = split_dataset.entrypoint( + dataset_chips=str(toy_chips), + dataset_labels=str(toy_labels), + hyperparameters=base_hyperparameters, + ) + model = train_model.entrypoint( + dataset_chips=str(toy_chips), + dataset_labels=str(toy_labels), + base_model_weights=pretrained_weights, + hyperparameters=base_hyperparameters, + split_info=split_info, + num_classes=2, + ) + params = calibrate_threshold.entrypoint( + trained_model=model, + dataset_chips=str(toy_chips), + dataset_labels=str(toy_labels), + hyperparameters=base_hyperparameters, + split_info=split_info, + num_classes=2, + ) + + assert params["method"] in {"val_sweep", "rate_match", "default"} + assert 0.0 < params["confidence_threshold"] < 1.0 + if params["method"] == "val_sweep": + assert 0.0 <= params["val_f1"] <= 1.0 + + disabled = calibrate_threshold.entrypoint( + trained_model=model, + dataset_chips=str(toy_chips), + dataset_labels=str(toy_labels), + hyperparameters={**base_hyperparameters, "calibrate_threshold": False}, + split_info=split_info, + num_classes=2, + ) + assert disabled["method"] == "default" + assert disabled["confidence_threshold"] == base_hyperparameters.get("confidence_threshold", 0.5) + + def test_export_onnx(base_hyperparameters: dict[str, Any]) -> None: import onnx from torchgeo.models import unet