#!/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())