Monitoring Class-Balance Drift Across Image Tiles

To detect class-balance drift across image tiles, compute a per-class annotation frequency vector for every acquisition batch, normalise it into a class-prior distribution, and compare each incoming batch against a trusted reference prior using the Population Stability Index (PSI). PSI scores each class by (current_prop - reference_prop) * ln(current_prop / reference_prop); a per-class PSI above 0.25 flags significant class-balance drift and should trigger targeted re-labeling of that batch’s tiles before they enter training. This measure is the batch-level companion to per-tile monitoring inside an active learning loop, where annotation effort is steered toward the data the model handles worst. The rest of this guide gives a runnable pandas implementation, a threshold table, and the failure modes that quietly break PSI monitoring.

# Why Class-Balance Drift Silently Degrades Geospatial Models

A geospatial detector learns the class priors baked into its training set. When later acquisition batches carry a different mix — a summer flight thick with cars where the winter baseline had mostly bare roofs, or a new sensor that resolves small vehicles the old one missed — the label distribution shifts even though every individual annotation is correct. That shift is label drift, and it erodes model performance in a way that overall accuracy hides: the minority classes lose recall first, and aggregate metrics stay high because the dominant class still scores well.

Per-class frequency is the cheapest drift signal you can compute. It needs no imagery, no model inference, and no reprojection — only the class property already present in each annotation. Yet a naive check (“did any class change by more than 5 percentage points?”) misfires in both directions. It ignores a rare class tripling from 1% to 3%, and it panics over a dominant class drifting from 40% to 45%. PSI fixes both by weighting every change by its log-ratio, so proportional moves in small classes dominate the score exactly where drift does the most damage. Because class labels are the same discrete categories the ROI label taxonomy already fixes, the reference vector is stable and cheap to maintain across the life of a dataset.

# Step-by-Step Implementation

Install the pinned dependencies once:

pip install pandas==2.2.2 numpy==1.26.4

# Step 1 — Tabulate Class Counts per Batch from GeoJSON

Each tile is a GeoJSON FeatureCollection whose features carry a class property. Group tiles into acquisition batches (a flight, a sensor pass, or a delivery) and count classes per batch. Store the batch identifier and the source EPSG:4326 coordinate reference system alongside the counts so a flagged batch stays traceable to its raw tiles.

from __future__ import annotations
import json
from collections import Counter
from pathlib import Path
import pandas as pd


def count_classes_in_tile(geojson_path: Path, class_key: str = "class") -> Counter[str]:
    """Return a Counter of class -> annotation count for one tile GeoJSON."""
    fc: dict = json.loads(geojson_path.read_text())
    counts: Counter[str] = Counter()
    for feature in fc.get("features", []):
        label = feature.get("properties", {}).get(class_key)
        if label is not None:
            counts[str(label)] += 1
    return counts


def tabulate_batch_counts(
    batch_tiles: dict[str, list[Path]],
    class_key: str = "class",
) -> pd.DataFrame:
    """
    Aggregate per-class counts for each batch.

    Args:
        batch_tiles: mapping of batch_id -> list of tile GeoJSON paths.
        class_key:   feature property holding the class label.

    Returns:
        DataFrame indexed by batch_id, one column per class, values = counts.
    """
    rows: dict[str, Counter[str]] = {}
    for batch_id, tiles in batch_tiles.items():
        total: Counter[str] = Counter()
        for tile in tiles:
            total.update(count_classes_in_tile(tile, class_key))
        rows[batch_id] = total
    return pd.DataFrame(rows).fillna(0).astype(int).T.sort_index(axis=1)

# Step 2 — Build the Reference Class-Prior Vector

Normalise a trusted baseline batch (or a pooled historical window) into proportions. Apply Laplace smoothing so no class prior is exactly zero — a zero reference makes the PSI log-ratio explode to infinity.

import numpy as np


def to_prior(counts: pd.Series, smoothing: float = 1.0) -> pd.Series:
    """Convert a per-class count Series into a smoothed proportion vector."""
    smoothed = counts.astype(float) + smoothing
    return smoothed / smoothed.sum()


def build_reference_prior(
    counts_df: pd.DataFrame,
    reference_batches: list[str],
    smoothing: float = 1.0,
) -> pd.Series:
    """Pool one or more baseline batches into a single reference prior."""
    pooled: pd.Series = counts_df.loc[reference_batches].sum(axis=0)
    return to_prior(pooled, smoothing)

# Step 3 — Compute PSI per Class

PSI is a sum over classes of (current - reference) * ln(current / reference). Compute the per-class contribution so you can flag which class drifted, not merely that the batch as a whole moved. Align both vectors on the full class index first so a class missing from one side still contributes.

def psi_per_class(
    current_counts: pd.Series,
    reference_prior: pd.Series,
    smoothing: float = 1.0,
) -> pd.DataFrame:
    """
    Population Stability Index per class between a current batch and a reference.

    Returns a DataFrame with reference_prop, current_prop, and psi per class,
    sorted by descending psi so the worst-drifting class is first.
    """
    classes = reference_prior.index.union(current_counts.index)
    ref = reference_prior.reindex(classes, fill_value=0.0)
    cur = to_prior(current_counts.reindex(classes, fill_value=0), smoothing)
    # Re-smooth the reference in case new classes appeared.
    ref = (ref + 1e-9) / (ref + 1e-9).sum()

    psi = (cur - ref) * np.log(cur / ref)
    out = pd.DataFrame(
        {"reference_prop": ref, "current_prop": cur, "psi": psi}
    )
    return out.sort_values("psi", ascending=False)

# Step 4 — Emit an Alert Table

Band each per-class PSI, compute the total PSI for the batch, and mark any class above 0.25 for targeted re-labeling. The action column feeds straight into an annotation-queue routing step.

from dataclasses import dataclass


def psi_band(value: float) -> tuple[str, str]:
    """Map a PSI value to (band, recommended action)."""
    if value < 0.10:
        return "stable", "no action"
    if value < 0.25:
        return "watch", "monitor next batch"
    return "drift", "re-label this class"


@dataclass
class BatchDriftReport:
    batch_id: str
    total_psi: float
    per_class: pd.DataFrame  # includes band and action columns


def build_alert_table(
    batch_id: str,
    current_counts: pd.Series,
    reference_prior: pd.Series,
) -> BatchDriftReport:
    table = psi_per_class(current_counts, reference_prior)
    bands = table["psi"].apply(psi_band)
    table["band"] = [b for b, _ in bands]
    table["action"] = [a for _, a in bands]
    total = float(table["psi"].sum())
    return BatchDriftReport(batch_id=batch_id, total_psi=total, per_class=table)


# --- Worked example -------------------------------------------------------
counts_df = pd.DataFrame(
    {
        "building": [400, 280],
        "road":     [250, 220],
        "vehicle":  [100, 300],
        "veg":      [250, 200],
    },
    index=["2026-04-baseline", "2026-07-flightB"],
)

reference = build_reference_prior(counts_df, ["2026-04-baseline"])
report = build_alert_table("2026-07-flightB", counts_df.loc["2026-07-flightB"], reference)

print(f"Total PSI for {report.batch_id}: {report.total_psi:.3f}")
print(report.per_class.round(3).to_string())

Running the example prints a total PSI of roughly 0.34, with vehicle carrying a per-class PSI well above 0.25 and an action of re-label this class — the batch tripled its vehicle prevalence, exactly the minority-class shift that erodes recall. Persisting each BatchDriftReport next to its dataset snapshot keeps the drift signal reproducible; the mechanics of attaching that context live in preserving metadata across dataset versions.

# PSI Thresholds and Actions

The bands below are the standard PSI conventions borrowed from population-stability monitoring; they transfer cleanly to categorical class priors because both are discrete distributions.

Read the total PSI as a batch-level smoke alarm and the per-class PSI as the diagnosis: a total of 0.34 tells you the batch drifted, but only the per-class breakdown tells you it was vehicle that moved and therefore which tiles to re-annotate. Never retrain on a batch whose dominant-class PSI sits in the drift band without first confirming the shift is real rather than a labeling-protocol change.

# Common Errors and Fixes

RuntimeWarning: divide by zero encountered in log and PSI of inf Root cause: a class has a reference proportion of exactly 0.0, so ln(current / 0) is undefined. Fix: apply Laplace smoothing (smoothing=1.0) to every count vector before normalising, or add a small epsilon and renormalise, as to_prior does above.

PSI spikes on every batch even when the label mix looks unchanged Root cause: PSI computed per tile rather than per batch — a single tile holds too few annotations for stable proportions, so the score is dominated by sampling noise. Fix: aggregate tiles into batches of at least a few hundred annotations before computing PSI.

A newly introduced class produces a misleadingly small PSI Root cause: the class exists in the current batch but not in the reference index, so it is dropped during alignment. Fix: take the union of both class indices (reference_prior.index.union(current_counts.index)) and reindex with fill_value=0 before smoothing, so new classes contribute their full drift.

Total PSI looks fine but the model still loses recall on a rare class Root cause: a large drift in one minority class is diluted when only the summed total is inspected. Fix: alert on the maximum per-class PSI, not just the total; a single class above 0.25 warrants re-labeling regardless of the aggregate.

# Related

• Detecting Distribution Drift in Spatial Datasets — the broader topic area covering spectral, covariate, and label drift monitors this class-balance check belongs to
• Preserving Metadata Across Dataset Versions — persist each batch’s drift report and reference prior so a flagged shift stays reproducible across snapshots
• Defining ROI Label Taxonomies for Aerial Imagery — a stable class taxonomy is what makes the reference prior vector meaningful batch after batch
• Active Learning & Model Feedback Loops for Geospatial Annotation — how drift alerts feed the retraining and re-labeling triage that keeps a deployed detector current

This guide covers one specialised monitor within Detecting Distribution Drift in Spatial Datasets, which is itself part of Active Learning & Model Feedback Loops for Geospatial Annotation.