geointel/scripts/audit_yolo_label_outliers.py

#!/usr/bin/env python3
"""Audit individual YOLO rows for geometric training risks without mutation."""

from __future__ import annotations

import argparse
import json
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

try:
    from scripts.audit_yolo_label_relationships import (
        Box,
        parse_label_file,
        resolve_path,
        sha256_file,
    )
except ModuleNotFoundError:  # Standalone operator-tool copy beside the auditor.
    from audit_yolo_label_relationships import (
        Box,
        parse_label_file,
        resolve_path,
        sha256_file,
    )


def classify_box(
    box: Box,
    *,
    tile_size: int,
    min_dimension_pixels: float,
    extreme_aspect_ratio: float,
    edge_tolerance_pixels: float,
) -> list[dict[str, Any]]:
    width_px = box.width * tile_size
    height_px = box.height * tile_size
    aspect_ratio = max(width_px / height_px, height_px / width_px)
    left, top, right, bottom = box.coordinates
    edge_sides = [
        side
        for side, distance in (
            ("left", left * tile_size),
            ("top", top * tile_size),
            ("right", (1 - right) * tile_size),
            ("bottom", (1 - bottom) * tile_size),
        )
        if distance <= edge_tolerance_pixels
    ]
    metrics = {
        "width_px": round(width_px, 6),
        "height_px": round(height_px, 6),
        "area_px2": round(width_px * height_px, 6),
        "aspect_ratio": round(aspect_ratio, 6),
    }
    outliers: list[dict[str, Any]] = []
    if min(width_px, height_px) < min_dimension_pixels:
        outliers.append({"category": "small_dimension", **metrics})
    if aspect_ratio >= extreme_aspect_ratio:
        outliers.append({"category": "extreme_aspect_ratio", **metrics})
    if edge_sides:
        outliers.append({"category": "tile_edge", "edge_sides": edge_sides, **metrics})
    return outliers


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--summary-path", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--tile-size", type=int)
    parser.add_argument("--min-dimension-pixels", type=float, default=4.0)
    parser.add_argument("--extreme-aspect-ratio", type=float, default=8.0)
    parser.add_argument("--edge-tolerance-pixels", type=float, default=0.5)
    parser.add_argument(
        "--category",
        action="append",
        choices=("small_dimension", "extreme_aspect_ratio", "tile_edge"),
        default=[],
        help="Limit output to one risk category; repeat to select multiple.",
    )
    args = parser.parse_args()

    if args.output.exists():
        parser.error(f"output already exists: {args.output}")
    if args.min_dimension_pixels <= 0:
        parser.error("--min-dimension-pixels must be positive")
    if args.extreme_aspect_ratio < 1:
        parser.error("--extreme-aspect-ratio must be at least one")
    if args.edge_tolerance_pixels < 0:
        parser.error("--edge-tolerance-pixels must be non-negative")

    summary_path = args.summary_path.expanduser().resolve(strict=True)
    summary = json.loads(summary_path.read_text(encoding="utf-8"))
    tiles = summary.get("tiles")
    if not isinstance(tiles, list):
        raise ValueError("dataset summary must contain a tiles list")
    tile_size = args.tile_size or summary.get("tile_size")
    if not isinstance(tile_size, int) or tile_size <= 0:
        raise ValueError("a positive integer tile size is required")

    selected_categories = set(
        args.category or ("small_dimension", "extreme_aspect_ratio", "tile_edge")
    )
    category_totals = {
        "small_dimension": 0,
        "extreme_aspect_ratio": 0,
        "tile_edge": 0,
    }
    reviewed_label_count = 0
    unique_flagged_rows: set[tuple[str, int]] = set()
    flagged_tiles: list[dict[str, Any]] = []
    renderable_tiles: list[dict[str, Any]] = []

    for tile in tiles:
        if not isinstance(tile, dict) or not tile.get("kept", True):
            continue
        label_path = resolve_path(str(tile.get("label_path") or ""), summary_path)
        boxes = parse_label_file(label_path)
        reviewed_label_count += len(boxes)
        outliers: list[dict[str, Any]] = []
        for index, box in enumerate(boxes):
            classifications = classify_box(
                box,
                tile_size=tile_size,
                min_dimension_pixels=args.min_dimension_pixels,
                extreme_aspect_ratio=args.extreme_aspect_ratio,
                edge_tolerance_pixels=args.edge_tolerance_pixels,
            )
            for classification in classifications:
                if classification["category"] not in selected_categories:
                    continue
                category_totals[classification["category"]] += 1
                unique_flagged_rows.add((str(label_path), index))
                outliers.append(
                    {
                        "index": index,
                        "box": box.as_list(),
                        **classification,
                    }
                )
        if not outliers:
            continue
        flagged_tiles.append(
            {
                "sample_slug": str(tile.get("sample_slug") or "unknown"),
                "split": str(tile.get("split") or "unknown"),
                "tile_index": int(tile.get("tile_index") or 0),
                "label_path": str(label_path),
                "label_count": len(boxes),
                "outlier_count": len(outliers),
                "outliers": outliers,
            }
        )
        renderable_tiles.append(tile)

    flagged_tiles.sort(
        key=lambda item: (
            -item["outlier_count"],
            item["sample_slug"],
            item["split"],
            item["tile_index"],
        )
    )
    payload = {
        "schema_version": 1,
        "generated_at": datetime.now(UTC).isoformat(),
        "status": "attention" if flagged_tiles else "ok",
        "claim_boundary": (
            "Read-only geometric risk triage; a flagged label is not an automatic "
            "ground-truth error or rewrite instruction."
        ),
        "summary_path": str(summary_path),
        "summary_sha256": sha256_file(summary_path),
        "output_dir": summary.get("output_dir"),
        "class_names": summary.get("class_names", []),
        "tile_size": tile_size,
        "thresholds": {
            "min_dimension_pixels": args.min_dimension_pixels,
            "extreme_aspect_ratio": args.extreme_aspect_ratio,
            "edge_tolerance_pixels": args.edge_tolerance_pixels,
        },
        "selected_categories": sorted(selected_categories),
        "reviewed_tile_count": sum(
            1 for tile in tiles if isinstance(tile, dict) and tile.get("kept", True)
        ),
        "reviewed_label_count": reviewed_label_count,
        "flagged_tile_count": len(flagged_tiles),
        "unique_flagged_label_count": len(unique_flagged_rows),
        "category_totals": category_totals,
        "flagged_tiles": flagged_tiles,
        "tiles": renderable_tiles,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    print(
        json.dumps(
            {
                key: payload[key]
                for key in (
                    "status",
                    "reviewed_tile_count",
                    "reviewed_label_count",
                    "flagged_tile_count",
                    "unique_flagged_label_count",
                    "category_totals",
                )
            },
            indent=2,
        )
    )
    return 0


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