active_learning.py

#!/usr/bin/env python3
"""Rank unlabeled frames by how much they would improve the dataset.

Label time is the scarce resource in the flywheel. Random sampling wastes it on
frames the model already handles. This scores each candidate on three axes and
returns a ranked labelling queue:

* **novelty** - 1 minus its perceptual-hash similarity to anything already
  labelled, so near-duplicates of existing walls sink to the bottom;
* **domain need** - rare metadata values (weather, lighting, material, paint
  stage, ...) score higher, filling the balance gaps the audit reports;
* **uncertainty** - optional, from a trained checkpoint (mean predictive
  entropy), so frames the model is unsure about are prioritised.

The output is a queue with per-frame reasons, so a human can label the top N and
get the most coverage per minute. numpy + Pillow for scoring; uncertainty is the
only part that needs torch.

Usage::

    python3 active_learning.py --data data --candidates unlabeled/ --top 50 \
        --candidates-metadata unlabeled/metadata.csv --out queue.csv
    python3 active_learning.py --data data --candidates unlabeled/ --checkpoint artifacts/best.pt
"""
from __future__ import annotations

import argparse
import csv
import json
from collections import Counter
from pathlib import Path

from flywheel import build_index, dhash, hamming

ROOT = Path(__file__).resolve().parent
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png"}
DEFAULT_FIELDS = ("environment", "surface_material", "lighting", "weather", "paint_stage")
DEFAULT_WEIGHTS = {"novelty": 0.5, "domain": 0.3, "uncertainty": 0.2}
HASH_BITS = 64


def labelled_counts(data_dir, fields=DEFAULT_FIELDS):
    """Per-field value counts from the labelled dataset's metadata."""
    path = Path(data_dir) / "metadata.csv"
    counts = {field: Counter() for field in fields}
    if not path.is_file():
        return counts
    with path.open(newline="", encoding="utf-8-sig") as stream:
        for row in csv.DictReader(stream):
            for field in fields:
                value = (row.get(field) or "").strip().lower()
                if value:
                    counts[field][value] += 1
    return counts


def metadata_rows(path):
    path = Path(path)
    if not path.is_file():
        return {}
    with path.open(newline="", encoding="utf-8-sig") as stream:
        return {Path(row.get("image", "")).name: row for row in csv.DictReader(stream)}


def novelty_score(fingerprint, existing_hashes):
    if not existing_hashes:
        return 1.0, None
    best = min(hamming(fingerprint, other) for other in existing_hashes)
    return best / HASH_BITS, best


def domain_score(row, counts, fields=DEFAULT_FIELDS):
    scores, reasons = [], []
    for field in fields:
        value = (row.get(field) or "").strip().lower()
        if not value:
            continue
        column = counts.get(field, Counter())
        peak = max(column.values()) if column else 1
        presence = column.get(value, 0)
        scores.append(1.0 - presence / peak)
        if presence <= peak * 0.25:
            reasons.append(f"rare {field}={value}")
    return (sum(scores) / len(scores) if scores else 0.0), reasons


def score_candidates(candidates, index, counts, rows=None, fields=DEFAULT_FIELDS, weights=None,
                     uncertainties=None, top=None):
    rows = rows or {}
    weights = {**DEFAULT_WEIGHTS, **(weights or {})}
    existing_hashes = list(index["phash"].keys())
    ranked = []
    for path in candidates:
        row = rows.get(Path(path).name, {})
        novelty, distance = novelty_score(dhash(path), existing_hashes)
        domain, reasons = domain_score(row, counts, fields)
        uncertainty = float(uncertainties.get(str(path), 0.0)) if uncertainties else None
        components = {"novelty": novelty, "domain": domain}
        active_weights = {k: weights[k] for k in ("novelty", "domain")}
        if uncertainty is not None:
            components["uncertainty"] = uncertainty
            active_weights["uncertainty"] = weights["uncertainty"]
        total = sum(active_weights.values()) or 1.0
        score = sum(components[k] * w for k, w in active_weights.items()) / total
        if novelty > 0.6:
            reasons = ["novel scene"] + reasons
        if uncertainty is not None and uncertainty > 0.3:
            reasons = ["high model uncertainty"] + reasons
        ranked.append({"image": str(path), "score": round(score, 4), "novelty": round(novelty, 4),
                       "nearest_hash_distance": distance, "domain": round(domain, 4),
                       "uncertainty": (round(uncertainty, 4) if uncertainty is not None else None),
                       "reasons": reasons})
    ranked.sort(key=lambda item: item["score"], reverse=True)
    return ranked[:top] if top else ranked


def uncertainty_for_images(images, checkpoint, size, device):
    """Mean predictive entropy per image from a Painting Vision checkpoint."""
    import numpy as np
    import torch

    from models import DEFAULT_ARCH, build_model

    state = torch.load(checkpoint, map_location=device, weights_only=True)
    model = build_model(str(state.get("arch", DEFAULT_ARCH)), len(state["classes"]), pretrained_backbone=False)
    model.load_state_dict(state["model"])
    model.to(device).eval()
    mean = np.asarray(state["mean"], np.float32)[None, None, :]
    std = np.asarray(state["std"], np.float32)[None, None, :]
    scores = {}
    from PIL import Image
    import cv2
    for path in images:
        bgr = cv2.imread(str(path))
        if bgr is None:
            continue
        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
        height, width = rgb.shape[:2]
        scale = size / max(height, width)
        resized = cv2.resize(rgb, (max(1, round(width * scale)), max(1, round(height * scale))))
        canvas = np.full((size, size, 3), 114, np.uint8)
        top, left = (size - resized.shape[0]) // 2, (size - resized.shape[1]) // 2
        canvas[top:top + resized.shape[0], left:left + resized.shape[1]] = resized
        tensor = torch.from_numpy(((canvas.astype(np.float32) / 255.0 - mean) / std)
                                  .transpose(2, 0, 1).copy()).unsqueeze(0)
        with torch.inference_mode():
            probabilities = torch.softmax(model(tensor.to(device))["semantic"], dim=1)[0].cpu().numpy()
        entropy = -np.sum(probabilities * np.log(probabilities + 1e-9), axis=0)
        scores[str(path)] = float(entropy.mean() / np.log(probabilities.shape[0]))
    return scores


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--data", type=Path, required=True, help="labelled dataset root")
    parser.add_argument("--candidates", type=Path, required=True, help="directory of unlabeled frames")
    parser.add_argument("--candidates-metadata", type=Path, help="optional CSV with domain columns per candidate")
    parser.add_argument("--fields", default=",".join(DEFAULT_FIELDS))
    parser.add_argument("--top", type=int, default=50)
    parser.add_argument("--checkpoint", type=Path, help="optional checkpoint for uncertainty scoring")
    parser.add_argument("--size", type=int, default=640)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--out", type=Path, help="write the ranked queue CSV here")
    args = parser.parse_args()

    fields = tuple(f.strip() for f in args.fields.split(",") if f.strip())
    rows = metadata_rows(args.candidates_metadata)

    candidates = sorted(p for p in args.candidates.iterdir() if p.suffix.lower() in IMAGE_EXTENSIONS)
    if not candidates:
        raise SystemExit(f"no images found in {args.candidates}")
    index = build_index(args.data)
    counts = labelled_counts(args.data, fields)
    uncertainties = None
    if args.checkpoint:
        import torch
        device = torch.device(args.device if torch.cuda.is_available() or args.device == "cpu" else "cpu")
        uncertainties = uncertainty_for_images(candidates, args.checkpoint, args.size, device)

    ranked = score_candidates(candidates, index, counts, rows=rows, fields=fields,
                              uncertainties=uncertainties, top=args.top)
    summary = {"labelled_images": len(index["sha256"]), "candidates": len(candidates),
               "queue": len(ranked), "domain_gaps": {field: dict(counts[field]) for field in fields}, "ranked": ranked}
    print(json.dumps({k: v for k, v in summary.items() if k != "ranked"}, indent=2))
    for item in ranked:
        print(f"{item['score']:.3f}  {item['image']}  {', '.join(item['reasons']) or '-'}")
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        with args.out.open("w", newline="", encoding="utf-8") as stream:
            writer = csv.DictWriter(stream, fieldnames=["image", "score", "novelty", "nearest_hash_distance",
                                                        "domain", "uncertainty", "reasons"])
            writer.writeheader()
            for item in ranked:
                writer.writerow({**item, "reasons": "; ".join(item["reasons"])})


if __name__ == "__main__":
    main()