274 lines
9.2 KiB
Python
274 lines
9.2 KiB
Python
#!/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())
|