geointel/scripts/audit_yolo_label_relationships.py

#!/usr/bin/env python3
"""Audit same-class relationships inside YOLO label files without rewriting labels."""

from __future__ import annotations

import argparse
import hashlib
import json
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any


@dataclass(frozen=True)
class Box:
    class_id: int
    center_x: float
    center_y: float
    width: float
    height: float

    @property
    def coordinates(self) -> tuple[float, float, float, float]:
        return (
            self.center_x - self.width / 2,
            self.center_y - self.height / 2,
            self.center_x + self.width / 2,
            self.center_y + self.height / 2,
        )

    @property
    def area(self) -> float:
        return self.width * self.height

    def as_list(self) -> list[float | int]:
        return [self.class_id, self.center_x, self.center_y, self.width, self.height]


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


def resolve_path(raw_path: str, summary_path: Path) -> Path:
    candidate = Path(raw_path)
    if candidate.exists():
        return candidate
    if candidate.is_absolute() and raw_path.startswith("/app/"):
        local_candidate = Path.cwd() / raw_path.removeprefix("/app/")
        if local_candidate.exists():
            return local_candidate
    relative_candidate = summary_path.parent / raw_path
    if relative_candidate.exists():
        return relative_candidate
    return candidate


def parse_label_file(path: Path) -> list[Box]:
    boxes: list[Box] = []
    for line_number, line in enumerate(
        path.read_text(encoding="utf-8").splitlines(), start=1
    ):
        stripped = line.strip()
        if not stripped:
            continue
        parts = stripped.split()
        if len(parts) != 5:
            raise ValueError(f"invalid YOLO row at {path}:{line_number}")
        class_id = int(float(parts[0]))
        center_x, center_y, width, height = map(float, parts[1:])
        if not (
            0 <= center_x <= 1
            and 0 <= center_y <= 1
            and 0 < width <= 1
            and 0 < height <= 1
        ):
            raise ValueError(f"out-of-range YOLO row at {path}:{line_number}")
        boxes.append(Box(class_id, center_x, center_y, width, height))
    return boxes


def intersection_area(first: Box, second: Box) -> float:
    first_box = first.coordinates
    second_box = second.coordinates
    width = max(
        0.0, min(first_box[2], second_box[2]) - max(first_box[0], second_box[0])
    )
    height = max(
        0.0, min(first_box[3], second_box[3]) - max(first_box[1], second_box[1])
    )
    return width * height


def classify_pair(
    first: Box,
    second: Box,
    *,
    near_duplicate_iou: float,
    containment_threshold: float,
    max_nested_area_ratio: float,
) -> tuple[str, float] | None:
    if first.class_id != second.class_id:
        return None
    if first == second:
        return "exact_duplicate", 1.0
    intersection = intersection_area(first, second)
    if intersection <= 0:
        return None
    union = first.area + second.area - intersection
    iou = min(1.0, intersection / union)
    if iou >= near_duplicate_iou:
        return "near_duplicate", iou
    minimum_area = min(first.area, second.area)
    area_ratio = max(first.area, second.area) / minimum_area
    containment = min(1.0, intersection / minimum_area)
    if containment >= containment_threshold and area_ratio <= max_nested_area_ratio:
        return "possible_nested", containment
    return None


def audit_boxes(
    boxes: list[Box],
    *,
    near_duplicate_iou: float,
    containment_threshold: float,
    max_nested_area_ratio: float,
) -> list[dict[str, Any]]:
    relationships: list[dict[str, Any]] = []
    for first_index, first in enumerate(boxes):
        for second_index in range(first_index + 1, len(boxes)):
            second = boxes[second_index]
            classification = classify_pair(
                first,
                second,
                near_duplicate_iou=near_duplicate_iou,
                containment_threshold=containment_threshold,
                max_nested_area_ratio=max_nested_area_ratio,
            )
            if classification is None:
                continue
            relationship, score = classification
            relationships.append(
                {
                    "relationship": relationship,
                    "score": round(score, 12),
                    "first_index": first_index,
                    "second_index": second_index,
                    "first_box": first.as_list(),
                    "second_box": second.as_list(),
                }
            )
    return relationships


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("--near-duplicate-iou", type=float, default=0.90)
    parser.add_argument("--containment-threshold", type=float, default=0.98)
    parser.add_argument("--max-nested-area-ratio", type=float, default=4.0)
    args = parser.parse_args()

    if args.output.exists():
        parser.error(f"output already exists: {args.output}")
    if not 0 < args.near_duplicate_iou <= 1:
        parser.error("--near-duplicate-iou must be in (0, 1]")
    if not 0 < args.containment_threshold <= 1:
        parser.error("--containment-threshold must be in (0, 1]")
    if args.max_nested_area_ratio < 1:
        parser.error("--max-nested-area-ratio must be at least one")

    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")

    totals = {"exact_duplicate": 0, "near_duplicate": 0, "possible_nested": 0}
    reviewed_label_count = 0
    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)
        relationships = audit_boxes(
            boxes,
            near_duplicate_iou=args.near_duplicate_iou,
            containment_threshold=args.containment_threshold,
            max_nested_area_ratio=args.max_nested_area_ratio,
        )
        if not relationships:
            continue
        counts = {key: 0 for key in totals}
        for relationship in relationships:
            key = relationship["relationship"]
            counts[key] += 1
            totals[key] += 1
        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),
                "relationship_counts": counts,
                "relationships": relationships,
            }
        )
        renderable_tiles.append(tile)

    flagged_tiles.sort(
        key=lambda item: (
            -sum(item["relationship_counts"].values()),
            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 label relationship audit; possible nesting is not an automatic error decision.",
        "summary_path": str(summary_path),
        "summary_sha256": sha256_file(summary_path),
        "output_dir": summary.get("output_dir"),
        "class_names": summary.get("class_names", []),
        "thresholds": {
            "near_duplicate_iou": args.near_duplicate_iou,
            "containment_threshold": args.containment_threshold,
            "max_nested_area_ratio": args.max_nested_area_ratio,
        },
        "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),
        "relationship_totals": 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",
                    "relationship_totals",
                )
            },
            indent=2,
        )
    )
    return 0


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