#!/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())