geointel/scripts/train_building_proposal_classifier.py

#!/usr/bin/env python3
"""Train and export the binary building-proposal filter."""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
from collections import Counter
from pathlib import Path, PurePosixPath
from typing import Any, Mapping


SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

from training_dataset_eligibility import (  # noqa: E402
    TrainingEligibilityError,
    assert_frozen_manifest_training_eligible,
)
from training_release_manifest import (  # noqa: E402
    TrainingReleaseError,
    assert_yolo_summary_bound_to_embedded_training_release,
)


PROPOSAL_DATASET_PROVENANCE_NAME = "proposal-dataset-provenance.json"
PROPOSAL_DATASET_SCHEMA_VERSION = 1
_CROP_IMAGE_SUFFIXES = {".bmp", ".jpeg", ".jpg", ".png", ".tif", ".tiff", ".webp"}


class ProposalDatasetProvenanceError(ValueError):
    """Raised when proposal crops no longer match their governed provenance."""


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _canonical_json_bytes(payload: Mapping[str, Any]) -> bytes:
    return json.dumps(payload, ensure_ascii=False, separators=(",", ":"), sort_keys=True).encode("utf-8")


def _payload_sha256(payload: Mapping[str, Any]) -> str:
    normalized = dict(payload)
    normalized.pop("manifest_sha256", None)
    return hashlib.sha256(_canonical_json_bytes(normalized)).hexdigest()


def _is_sha256(value: Any) -> bool:
    return isinstance(value, str) and len(value) == 64 and all(character in "0123456789abcdef" for character in value)


def _require_mapping(value: Any, *, field: str) -> Mapping[str, Any]:
    if not isinstance(value, Mapping):
        raise ProposalDatasetProvenanceError(f"proposal dataset provenance {field} must be an object")
    return value


def _require_sha256(value: Any, *, field: str) -> str:
    if not _is_sha256(value):
        raise ProposalDatasetProvenanceError(f"proposal dataset provenance {field} must be a lowercase SHA-256")
    return str(value)


def _require_nonempty_path(value: Any, *, field: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise ProposalDatasetProvenanceError(f"proposal dataset provenance {field} must be a non-empty path")
    return value


def _safe_relative_crop_path(value: Any, *, dataset_root: Path) -> tuple[str, Path]:
    if not isinstance(value, str) or not value or "\\" in value:
        raise ProposalDatasetProvenanceError("proposal crop relative_path must be a non-empty POSIX relative path")
    relative = PurePosixPath(value)
    if relative.is_absolute() or ".." in relative.parts or "." in relative.parts:
        raise ProposalDatasetProvenanceError(f"proposal crop has an unsafe relative path: {value!r}")
    path = (dataset_root / Path(*relative.parts)).resolve(strict=False)
    try:
        path.relative_to(dataset_root)
    except ValueError as exc:  # pragma: no cover - resolved-path defence
        raise ProposalDatasetProvenanceError(f"proposal crop escapes dataset directory: {value!r}") from exc
    return relative.as_posix(), path


def load_governed_corpus_manifest(corpus_manifest_path: Path, *, fixture_mode: bool) -> dict[str, Any]:
    """Assert frozen evidence and current live Dataset eligibility before PyTorch."""

    try:
        manifest = assert_frozen_manifest_training_eligible(
            corpus_manifest_path,
            fixture_mode=fixture_mode,
            verify_live=True,
        )
    except TrainingEligibilityError as exc:
        raise ProposalDatasetProvenanceError(str(exc)) from exc
    if not isinstance(manifest, dict):  # pragma: no cover - defensive contract boundary
        raise ProposalDatasetProvenanceError("governed corpus manifest must be a JSON object")
    return manifest


def validate_proposal_dataset_provenance(
    dataset_dir: Path,
    corpus_manifest_path: Path,
    *,
    fixture_mode: bool,
) -> dict[str, Any]:
    """Fail closed unless every trainable crop still matches its source evidence.

    ImageFolder discovers files from the directory tree.  Merely verifying a
    manifest is insufficient: any unbound image placed under train/ or val/
    would otherwise enter PyTorch training.  This validator compares the
    complete tree with the immutable crop manifest and re-binds it to the
    supplied frozen corpus and summary.
    """

    root = dataset_dir.expanduser().resolve(strict=False)
    if not root.is_dir():
        raise ProposalDatasetProvenanceError(f"proposal dataset directory is unavailable: {root}")
    manifest_path = root / PROPOSAL_DATASET_PROVENANCE_NAME
    try:
        payload = json.loads(manifest_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise ProposalDatasetProvenanceError(f"proposal dataset provenance is unreadable: {manifest_path}") from exc
    if not isinstance(payload, dict):
        raise ProposalDatasetProvenanceError("proposal dataset provenance must be a JSON object")
    if payload.get("schema_version") != PROPOSAL_DATASET_SCHEMA_VERSION:
        raise ProposalDatasetProvenanceError("proposal dataset provenance schema_version is unsupported")
    if payload.get("status") != "ok" or payload.get("immutable") is not True:
        raise ProposalDatasetProvenanceError("proposal dataset provenance is not an immutable successful release")
    if payload.get("governed_corpus_live_recheck") is not True:
        raise ProposalDatasetProvenanceError("proposal dataset provenance lacks the governed live-source recheck")
    if bool(payload.get("fixture_mode")) != fixture_mode:
        raise ProposalDatasetProvenanceError("proposal dataset provenance fixture-mode binding mismatches this invocation")
    if payload.get("manifest_sha256") != _payload_sha256(payload):
        raise ProposalDatasetProvenanceError("proposal dataset provenance manifest checksum mismatches its content")

    expected_corpus_sha256 = sha256(corpus_manifest_path)
    source = _require_mapping(payload.get("source"), field="source")
    recorded_corpus = _require_mapping(source.get("corpus_manifest"), field="source.corpus_manifest")
    if _require_sha256(recorded_corpus.get("sha256"), field="source.corpus_manifest.sha256") != expected_corpus_sha256:
        raise ProposalDatasetProvenanceError("proposal crops are not bound to the supplied governed corpus manifest")

    corpus_freeze_path = corpus_manifest_path.parent / "corpus-freeze.json"
    recorded_freeze = _require_mapping(source.get("corpus_freeze"), field="source.corpus_freeze")
    if not corpus_freeze_path.is_file() or _require_sha256(
        recorded_freeze.get("sha256"), field="source.corpus_freeze.sha256"
    ) != sha256(corpus_freeze_path):
        raise ProposalDatasetProvenanceError("proposal crops are not bound to the current frozen corpus sidecar")

    summary = _require_mapping(source.get("summary"), field="source.summary")
    summary_path_value = _require_nonempty_path(summary.get("path"), field="source.summary.path")
    summary_path = Path(summary_path_value).expanduser().resolve(strict=False)
    if not summary_path.is_file():
        raise ProposalDatasetProvenanceError(f"proposal source summary is unavailable: {summary_path}")
    if _require_sha256(summary.get("sha256"), field="source.summary.sha256") != sha256(summary_path):
        raise ProposalDatasetProvenanceError("proposal source summary checksum no longer matches the crop provenance")
    try:
        summary_payload = json.loads(summary_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise ProposalDatasetProvenanceError(f"proposal source summary is unreadable: {summary_path}") from exc
    if not isinstance(summary_payload, Mapping) or summary_payload.get("source_manifest_sha256") != expected_corpus_sha256:
        raise ProposalDatasetProvenanceError("proposal source summary is not bound to the supplied governed corpus manifest")
    if summary.get("source_manifest_sha256") != expected_corpus_sha256:
        raise ProposalDatasetProvenanceError("proposal crop provenance summary binding is invalid")
    try:
        source_release = assert_yolo_summary_bound_to_embedded_training_release(
            summary_path=summary_path,
            corpus_manifest=corpus_manifest_path,
            fixture_mode=fixture_mode,
        )
    except TrainingReleaseError as exc:
        raise ProposalDatasetProvenanceError(
            "proposal source summary has no eligible immutable training release"
        ) from exc
    recorded_release = _require_mapping(source.get("training_release"), field="source.training_release")
    if (
        recorded_release.get("dataset_yaml_path") != source_release["dataset_yaml"]["path"]
        or recorded_release.get("dataset_yaml_sha256") != source_release["dataset_yaml"]["sha256"]
        or recorded_release.get("corpus_manifest_sha256") != source_release["corpus"]["manifest_sha256"]
    ):
        raise ProposalDatasetProvenanceError("proposal crop provenance release binding is invalid")

    proposal_model = _require_mapping(source.get("proposal_model"), field="source.proposal_model")
    _require_nonempty_path(proposal_model.get("path"), field="source.proposal_model.path")
    _require_sha256(proposal_model.get("sha256"), field="source.proposal_model.sha256")
    entries = payload.get("crops")
    if not isinstance(entries, list) or not entries:
        raise ProposalDatasetProvenanceError("proposal dataset provenance has no crop entries")

    expected_paths: set[str] = set()
    calculated_counts: Counter[str] = Counter()
    normalized_entries: list[dict[str, Any]] = []
    for entry in entries:
        item = _require_mapping(entry, field="crops[]")
        relative_path, crop_path = _safe_relative_crop_path(item.get("relative_path"), dataset_root=root)
        if relative_path in expected_paths:
            raise ProposalDatasetProvenanceError(f"proposal crop provenance has duplicate path: {relative_path}")
        expected_paths.add(relative_path)
        split = item.get("split")
        label = item.get("label")
        if split not in {"train", "val"} or label not in {"negative", "positive"}:
            raise ProposalDatasetProvenanceError(f"proposal crop has invalid split/class: {relative_path}")
        if not isinstance(item.get("sample_slug"), str) or not item["sample_slug"].strip():
            raise ProposalDatasetProvenanceError(f"proposal crop has no sample identity: {relative_path}")
        if not crop_path.is_file():
            raise ProposalDatasetProvenanceError(f"proposal crop is missing: {relative_path}")
        if _require_sha256(item.get("sha256"), field=f"crops[{relative_path}].sha256") != sha256(crop_path):
            raise ProposalDatasetProvenanceError(f"proposal crop checksum mismatches provenance: {relative_path}")
        if item.get("size_bytes") != crop_path.stat().st_size:
            raise ProposalDatasetProvenanceError(f"proposal crop byte size mismatches provenance: {relative_path}")
        _require_nonempty_path(item.get("source_image_path"), field=f"crops[{relative_path}].source_image_path")
        _require_sha256(item.get("source_image_sha256"), field=f"crops[{relative_path}].source_image_sha256")
        _require_nonempty_path(item.get("source_label_path"), field=f"crops[{relative_path}].source_label_path")
        _require_sha256(item.get("source_label_sha256"), field=f"crops[{relative_path}].source_label_sha256")
        calculated_counts[f"{split}/{label}"] += 1
        normalized_entries.append(dict(item))

    canonical_entries = sorted(normalized_entries, key=lambda item: str(item["relative_path"]))
    expected_crops_sha256 = hashlib.sha256(_canonical_json_bytes({"crops": canonical_entries})).hexdigest()
    if payload.get("crops_sha256") != expected_crops_sha256:
        raise ProposalDatasetProvenanceError("proposal crop collection checksum mismatches provenance")
    if payload.get("crop_count") != len(entries):
        raise ProposalDatasetProvenanceError("proposal crop count mismatches provenance")
    if payload.get("counts") != dict(sorted(calculated_counts.items())):
        raise ProposalDatasetProvenanceError("proposal crop class counts mismatch provenance")
    if any(calculated_counts[f"{split}/{label}"] < 1 for split in ("train", "val") for label in ("negative", "positive")):
        raise ProposalDatasetProvenanceError("proposal dataset has an empty train/validation class")

    actual_paths = {
        file_path.relative_to(root).as_posix()
        for file_path in root.rglob("*")
        if file_path.is_file() and file_path.suffix.lower() in _CROP_IMAGE_SUFFIXES
    }
    if actual_paths != expected_paths:
        unexpected = sorted(actual_paths - expected_paths)
        missing = sorted(expected_paths - actual_paths)
        raise ProposalDatasetProvenanceError(
            "proposal dataset files are not exactly the immutable crop manifest "
            f"(unexpected={unexpected}, missing={missing})"
        )
    return payload


def binary_metrics(scores: list[float], labels: list[int], threshold: float = 0.5) -> dict[str, float | int]:
    tp = sum(score >= threshold and label == 1 for score, label in zip(scores, labels, strict=True))
    fp = sum(score >= threshold and label == 0 for score, label in zip(scores, labels, strict=True))
    fn = sum(score < threshold and label == 1 for score, label in zip(scores, labels, strict=True))
    precision = tp / (tp + fp) if tp + fp else 1.0
    recall = tp / (tp + fn) if tp + fn else 1.0
    return {"tp": tp, "fp": fp, "fn": fn, "precision": precision, "recall": recall,
            "f1": 2 * precision * recall / (precision + recall) if precision + recall else 0.0}


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset-dir", type=Path, required=True)
    parser.add_argument(
        "--corpus-manifest",
        type=Path,
        required=True,
        help="The exact frozen governed corpus manifest that produced the proposal crops.",
    )
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--epochs", type=int, default=12)
    parser.add_argument("--batch", type=int, default=64)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--export-existing-best", action="store_true")
    parser.add_argument(
        "--fixture-mode",
        action="store_true",
        help="Permit only an explicitly fixture-only corpus; operational runs always resolve live dataset evidence.",
    )
    args = parser.parse_args()
    if args.output_dir.exists() and not args.export_existing_best:
        parser.error(f"output already exists: {args.output_dir}")

    # All governed-source checks run before importing torch/torchvision or
    # touching CUDA.  A source revocation therefore cannot start PyTorch work.
    try:
        governed_corpus = load_governed_corpus_manifest(
            args.corpus_manifest,
            fixture_mode=args.fixture_mode,
        )
        proposal_provenance = validate_proposal_dataset_provenance(
            args.dataset_dir,
            args.corpus_manifest,
            fixture_mode=args.fixture_mode,
        )
    except ProposalDatasetProvenanceError as exc:
        raise SystemExit(str(exc)) from exc
    proposal_provenance_path = args.dataset_dir / PROPOSAL_DATASET_PROVENANCE_NAME
    proposal_provenance_sha256 = sha256(proposal_provenance_path)

    import torch
    from torch import nn
    from torch.utils.data import DataLoader
    from torchvision.datasets import ImageFolder
    from torchvision.models import ResNet18_Weights, resnet18

    weights = ResNet18_Weights.DEFAULT
    transform = weights.transforms()
    train_ds = ImageFolder(args.dataset_dir / "train", transform=transform)
    val_ds = ImageFolder(args.dataset_dir / "val", transform=transform)
    if train_ds.class_to_idx != {"negative": 0, "positive": 1}:
        raise RuntimeError(f"unexpected class order: {train_ds.class_to_idx}")
    device = torch.device(args.device)
    model = resnet18(weights=weights)
    model.fc = nn.Linear(model.fc.in_features, 1)
    model.to(device)
    if args.export_existing_best:
        state_path = args.output_dir / "best-state.pt"
        if not state_path.is_file():
            raise RuntimeError(f"missing existing best state: {state_path}")
        model.load_state_dict(torch.load(state_path, map_location=device, weights_only=True))
        model.eval()
        torch.jit.script(model).save(str(args.output_dir / "proposal-classifier.torchscript.pt"))
        print(
            json.dumps(
                {
                    "status": "exported_existing_best",
                    "model": str(state_path),
                    "proposal_dataset_provenance": str(proposal_provenance_path),
                    "proposal_dataset_provenance_sha256": proposal_provenance_sha256,
                    "governed_corpus_manifest_sha256": sha256(args.corpus_manifest),
                    "governed_corpus_live_recheck": True,
                }
            )
        )
        return 0
    train_loader = DataLoader(train_ds, batch_size=args.batch, shuffle=True, num_workers=0)
    val_loader = DataLoader(val_ds, batch_size=args.batch, shuffle=False, num_workers=0)
    positives = sum(label == 1 for _path, label in train_ds.samples)
    negatives = len(train_ds) - positives
    loss_fn = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([negatives / positives], device=device))
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
    args.output_dir.mkdir(parents=True)
    history = []
    best_f1 = -1.0
    for epoch in range(1, args.epochs + 1):
        model.train()
        train_loss = 0.0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.float().to(device)
            optimizer.zero_grad(set_to_none=True)
            logits = model(images).flatten()
            loss = loss_fn(logits, labels)
            loss.backward()
            optimizer.step()
            train_loss += float(loss) * len(images)
        model.eval()
        scores: list[float] = []
        labels_out: list[int] = []
        with torch.inference_mode():
            for images, labels in val_loader:
                scores.extend(torch.sigmoid(model(images.to(device)).flatten()).cpu().tolist())
                labels_out.extend(labels.tolist())
        metric = binary_metrics(scores, labels_out)
        row = {"epoch": epoch, "train_loss": train_loss / len(train_ds), **metric}
        history.append(row)
        print(json.dumps(row), flush=True)
        if float(metric["f1"]) > best_f1:
            best_f1 = float(metric["f1"])
            torch.save(model.state_dict(), args.output_dir / "best-state.pt")
    model.load_state_dict(torch.load(args.output_dir / "best-state.pt", map_location=device, weights_only=True))
    model.eval()
    scripted = torch.jit.script(model)
    model_path = args.output_dir / "proposal-classifier.torchscript.pt"
    scripted.save(str(model_path))
    report = {
        "schema_version": 1,
        "status": "ok",
        "classes": train_ds.class_to_idx,
        "train_count": len(train_ds),
        "validation_count": len(val_ds),
        "best_validation_f1": best_f1,
        "history": history,
        "model": str(model_path),
        "model_sha256": sha256(model_path),
        "proposal_dataset_provenance": str(proposal_provenance_path),
        "proposal_dataset_provenance_sha256": proposal_provenance_sha256,
        "proposal_dataset_manifest_sha256": proposal_provenance["manifest_sha256"],
        "corpus_manifest": str(args.corpus_manifest),
        "corpus_manifest_sha256": sha256(args.corpus_manifest),
        "fixture_mode": bool(args.fixture_mode),
        "governed_corpus_live_recheck": True,
        "governed_corpus_sample_count": len(governed_corpus.get("samples", [])),
    }
    (args.output_dir / "training-report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())