geointel/scripts/derive_yolo_min_dimension_corpus.py

#!/usr/bin/env python3
"""Create an immutable experimental YOLO corpus with a stricter pixel floor."""

from __future__ import annotations

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


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 filter_label_lines(
    lines: list[str], *, tile_size: int, min_dimension_pixels: float
) -> tuple[list[str], list[int]]:
    kept: list[str] = []
    removed_indices: list[int] = []
    for index, line in enumerate(lines):
        stripped = line.strip()
        if not stripped:
            continue
        parts = stripped.split()
        if len(parts) != 5:
            raise ValueError(f"invalid YOLO row at zero-based index {index}")
        try:
            float(parts[0])
            width = float(parts[3])
            height = float(parts[4])
        except ValueError as exc:
            raise ValueError(
                f"invalid numeric YOLO row at zero-based index {index}"
            ) from exc
        if width <= 0 or height <= 0 or width > 1 or height > 1:
            raise ValueError(f"out-of-range YOLO row at zero-based index {index}")
        if min(width, height) * tile_size < min_dimension_pixels:
            removed_indices.append(index)
        else:
            kept.append(stripped)
    return kept, removed_indices


def aggregate_hash(paths: list[Path], root: Path) -> str:
    digest = hashlib.sha256()
    for path in sorted(paths):
        relative = path.relative_to(root).as_posix()
        digest.update(relative.encode("utf-8"))
        digest.update(b"\0")
        digest.update(sha256_file(path).encode("ascii"))
        digest.update(b"\n")
    return digest.hexdigest()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source-dir", required=True, type=Path)
    parser.add_argument("--output-dir", required=True, type=Path)
    parser.add_argument("--min-dimension-pixels", type=float, default=4.0)
    args = parser.parse_args()
    if args.min_dimension_pixels <= 0:
        parser.error("--min-dimension-pixels must be positive")

    source_dir = args.source_dir.expanduser().resolve(strict=True)
    output_dir = args.output_dir.expanduser().resolve(strict=False)
    if output_dir.exists():
        parser.error(f"output directory already exists: {output_dir}")
    source_summary_path = source_dir / "yolo_tile_dataset_summary.json"
    source_summary = json.loads(source_summary_path.read_text(encoding="utf-8"))
    tiles = source_summary.get("tiles")
    tile_size = source_summary.get("tile_size")
    if not isinstance(tiles, list):
        raise ValueError("source summary must contain a tiles list")
    if not isinstance(tile_size, int) or tile_size <= 0:
        raise ValueError("source summary must contain a positive integer tile_size")
    source_minimum = source_summary.get("min_label_px")
    if isinstance(source_minimum, (int, float)) and args.min_dimension_pixels < float(
        source_minimum
    ):
        raise ValueError("derived minimum may not weaken the source pixel floor")

    output_dir.mkdir(parents=True)
    kept_tiles: list[dict[str, Any]] = []
    excluded_tiles: list[dict[str, Any]] = []
    removed_labels: list[dict[str, Any]] = []
    output_images: list[Path] = []
    output_labels: list[Path] = []
    for tile in tiles:
        if not isinstance(tile, dict) or not tile.get("kept", True):
            continue
        source_image = Path(str(tile.get("image_path") or "")).resolve(strict=True)
        source_label = Path(str(tile.get("label_path") or "")).resolve(strict=True)
        image_relative = source_image.relative_to(source_dir)
        label_relative = source_label.relative_to(source_dir)
        lines = source_label.read_text(encoding="utf-8").splitlines()
        kept_lines, removed_indices = filter_label_lines(
            lines,
            tile_size=tile_size,
            min_dimension_pixels=args.min_dimension_pixels,
        )
        for index in removed_indices:
            removed_labels.append(
                {
                    "sample_slug": tile.get("sample_slug"),
                    "split": tile.get("split"),
                    "tile_index": tile.get("tile_index"),
                    "source_label_path": str(source_label),
                    "source_label_index": index,
                    "source_row": lines[index].strip(),
                    "reason": "dimension_below_pixel_floor",
                }
            )
        source_had_labels = any(line.strip() for line in lines)
        if source_had_labels and not kept_lines:
            excluded_tiles.append(
                {
                    "sample_slug": tile.get("sample_slug"),
                    "split": tile.get("split"),
                    "tile_index": tile.get("tile_index"),
                    "reason": "positive_tile_became_empty_after_filter",
                    "removed_label_count": len(removed_indices),
                    "source_image_path": str(source_image),
                    "source_label_path": str(source_label),
                }
            )
            continue

        target_image = output_dir / image_relative
        target_label = output_dir / label_relative
        target_image.parent.mkdir(parents=True, exist_ok=True)
        target_label.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(source_image, target_image)
        target_label.write_text(
            "".join(f"{line}\n" for line in kept_lines), encoding="utf-8"
        )
        output_images.append(target_image)
        output_labels.append(target_label)
        derived_tile = dict(tile)
        derived_tile.update(
            {
                "image_path": str(target_image),
                "label_path": str(target_label),
                "label_count": len(kept_lines),
                "is_negative": not kept_lines,
                "derived_from_image_path": str(source_image),
                "derived_from_label_path": str(source_label),
                "removed_label_count": len(removed_indices),
            }
        )
        kept_tiles.append(derived_tile)

    dataset_yaml = output_dir / "dataset.yaml"
    dataset_yaml.write_text(
        "\n".join(
            (
                f"path: {output_dir}",
                "train: images/train",
                "val: images/val",
                "names:",
                "  0: building",
                "",
            )
        ),
        encoding="utf-8",
    )
    positive_tiles = [tile for tile in kept_tiles if not tile["is_negative"]]
    negative_tiles = [tile for tile in kept_tiles if tile["is_negative"]]
    summary = {
        **{key: value for key, value in source_summary.items() if key != "tiles"},
        "schema_version": 1,
        "status": "experimental_derived_not_release_eligible",
        "claim_boundary": (
            "Deterministic min-dimension ablation only; no human-review, training, "
            "evaluation or production-release claim."
        ),
        "output_dir": str(output_dir),
        "dataset_yaml": str(dataset_yaml),
        "source_dataset_dir": str(source_dir),
        "source_summary_path": str(source_summary_path),
        "source_summary_sha256": sha256_file(source_summary_path),
        "transformation": "drop_label_if_min_dimension_pixels_below_threshold",
        "min_label_px": args.min_dimension_pixels,
        "source_min_label_px": source_minimum,
        "tile_count": len(kept_tiles),
        "positive_tile_count": len(positive_tiles),
        "negative_tile_count": len(negative_tiles),
        "train_tile_count": sum(tile.get("split") == "train" for tile in kept_tiles),
        "val_tile_count": sum(tile.get("split") == "val" for tile in kept_tiles),
        "label_count": sum(int(tile["label_count"]) for tile in kept_tiles),
        "removed_label_count": len(removed_labels),
        "excluded_tile_count": len(excluded_tiles),
        "training_release_eligible": False,
        "tiles": kept_tiles,
    }
    summary_path = output_dir / "yolo_tile_dataset_summary.json"
    summary_path.write_text(
        json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    manifest = {
        "schema_version": 1,
        "generated_at": datetime.now(UTC).isoformat(),
        "status": "complete_experimental_only",
        "source_summary_path": str(source_summary_path),
        "source_summary_sha256": sha256_file(source_summary_path),
        "output_summary_path": str(summary_path),
        "output_summary_sha256": sha256_file(summary_path),
        "dataset_yaml_sha256": sha256_file(dataset_yaml),
        "image_set_aggregate_sha256": aggregate_hash(output_images, output_dir),
        "label_set_aggregate_sha256": aggregate_hash(output_labels, output_dir),
        "min_dimension_pixels": args.min_dimension_pixels,
        "source_label_count": int(source_summary.get("label_count") or 0),
        "derived_label_count": summary["label_count"],
        "removed_label_count": len(removed_labels),
        "excluded_tiles": excluded_tiles,
        "removed_labels": removed_labels,
        "training_release_eligible": False,
    }
    manifest_path = output_dir / "derived_corpus_manifest.json"
    manifest_path.write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    print(
        json.dumps(
            {
                "status": manifest["status"],
                "derived_label_count": manifest["derived_label_count"],
                "removed_label_count": manifest["removed_label_count"],
                "excluded_tile_count": len(excluded_tiles),
                "manifest_path": str(manifest_path),
            },
            indent=2,
        )
    )
    return 0


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