geointel/scripts/retile_yolo_dataset.py

#!/usr/bin/env python3
"""Retile a YOLO dataset while preserving sample and protected-split identity."""

from __future__ import annotations

import argparse
import hashlib
import json
import shutil
from pathlib import Path
from typing import Any


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 read_boxes(path: Path, width: int, height: int) -> list[tuple[float, float, float, float]]:
    boxes = []
    for row in path.read_text(encoding="utf-8").splitlines() if path.is_file() else []:
        parts = row.split()
        if len(parts) != 5:
            raise ValueError(f"invalid YOLO row in {path}: {row}")
        _class_id, cx, cy, box_width, box_height = map(float, parts)
        boxes.append(((cx-box_width/2)*width, (cy-box_height/2)*height,
                      (cx+box_width/2)*width, (cy+box_height/2)*height))
    return boxes


def tile_starts(length: int, tile_size: int, overlap: int) -> list[int]:
    if length < tile_size:
        return []
    stride = tile_size - overlap
    starts = list(range(0, length - tile_size + 1, stride))
    final = length - tile_size
    if starts[-1] != final:
        starts.append(final)
    return starts


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--summary", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--tile-size", type=int, default=320)
    parser.add_argument("--overlap", type=int, default=0)
    parser.add_argument("--min-visible-ratio", type=float, default=0.25)
    parser.add_argument("--force", action="store_true")
    args = parser.parse_args()
    if args.tile_size < 32 or not 0 <= args.overlap < args.tile_size or not 0 <= args.min_visible_ratio <= 1:
        raise SystemExit("invalid tile size or visible ratio")
    if args.output_dir.exists():
        if not args.force:
            raise SystemExit(f"output exists: {args.output_dir}")
        shutil.rmtree(args.output_dir)

    from PIL import Image

    summary = json.loads(args.summary.read_text(encoding="utf-8"))
    output_tiles: list[dict[str, Any]] = []
    kept_labels = dropped_labels = source_labels = uniquely_kept_labels = 0
    for source_tile in summary["tiles"]:
        with Image.open(source_tile["image_path"]) as opened:
            source = opened.convert("RGB")
        width, height = source.size
        boxes = read_boxes(Path(source_tile["label_path"]), width, height)
        source_labels += len(boxes)
        covered_source_labels: set[int] = set()
        for top in tile_starts(height, args.tile_size, args.overlap):
            for left in tile_starts(width, args.tile_size, args.overlap):
                right, bottom = min(width, left+args.tile_size), min(height, top+args.tile_size)
                if right-left < args.tile_size or bottom-top < args.tile_size:
                    continue
                local = []
                for box_index, (x1, y1, x2, y2) in enumerate(boxes):
                    cx, cy = (x1+x2)/2, (y1+y2)/2
                    if not (left <= cx < right and top <= cy < bottom):
                        continue
                    ix1, iy1, ix2, iy2 = max(x1,left), max(y1,top), min(x2,right), min(y2,bottom)
                    visible = max(0,ix2-ix1)*max(0,iy2-iy1) / max((x2-x1)*(y2-y1), 1e-9)
                    if visible < args.min_visible_ratio:
                        dropped_labels += 1
                        continue
                    local.append((ix1-left, iy1-top, ix2-left, iy2-top))
                    kept_labels += 1
                    covered_source_labels.add(box_index)
                split = source_tile["split"]
                stem = f"{Path(source_tile['image_path']).stem}_z{top}_{left}"
                image_path = args.output_dir / "images" / split / f"{stem}.png"
                label_path = args.output_dir / "labels" / split / f"{stem}.txt"
                image_path.parent.mkdir(parents=True, exist_ok=True)
                label_path.parent.mkdir(parents=True, exist_ok=True)
                source.crop((left,top,right,bottom)).save(image_path)
                rows = []
                for x1,y1,x2,y2 in local:
                    rows.append(f"0 {(x1+x2)/(2*args.tile_size):.8f} {(y1+y2)/(2*args.tile_size):.8f} "
                                f"{(x2-x1)/args.tile_size:.8f} {(y2-y1)/args.tile_size:.8f}")
                label_path.write_text("\n".join(rows) + ("\n" if rows else ""), encoding="utf-8")
                tile = dict(source_tile)
                tile.update({"image_path":str(image_path), "label_path":str(label_path),
                             "label_count":len(local), "source_tile":source_tile["image_path"],
                             "retile_window":[left,top,right,bottom]})
                output_tiles.append(tile)
        uniquely_kept_labels += len(covered_source_labels)

    args.output_dir.mkdir(parents=True, exist_ok=True)
    output_summary = dict(summary)
    output_summary.update({"output_dir":str(args.output_dir), "dataset_yaml":str(args.output_dir/"dataset.yaml"),
                           "tiles":output_tiles, "retile_size":args.tile_size,
                           "retile_overlap":args.overlap,
                           "retile_min_visible_ratio":args.min_visible_ratio})
    summary_path = args.output_dir / "yolo_tile_dataset_summary.json"
    summary_path.write_text(json.dumps(output_summary, indent=2), encoding="utf-8")
    (args.output_dir/"dataset.yaml").write_text(
        f"path: {args.output_dir}\ntrain: images/train\nval: images/val\nnames:\n  0: building\n", encoding="utf-8")
    evidence = {"schema_version":1, "status":"ok", "source_summary":str(args.summary),
                "source_summary_sha256":sha256(args.summary), "tile_size":args.tile_size,
                "overlap":args.overlap,
                "min_visible_ratio":args.min_visible_ratio, "output_tile_count":len(output_tiles),
                "kept_label_count":kept_labels, "dropped_label_count":dropped_labels,
                "source_label_count":source_labels, "unique_kept_source_label_count":uniquely_kept_labels,
                "uncovered_source_label_count":source_labels-uniquely_kept_labels,
                "summary":str(summary_path), "summary_sha256":sha256(summary_path)}
    (args.output_dir/"retile-evidence.json").write_text(json.dumps(evidence,indent=2),encoding="utf-8")
    print(json.dumps(evidence,indent=2))
    return 0


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