Add automatic confidence-threshold calibration to both training pipelines

models/dinov3s_buildings/pipeline.py

@@ -194,22 +194,144 @@ 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,
     dataset_chips: str,
     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
 
     n_trials = int(hyperparameters.get("tune_postprocess_trials", 30))
     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)
 

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."},

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,

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",