From d8fc66230eb0156d0af072543456e1f6d2ecb8d8 Mon Sep 17 00:00:00 2001 From: Jens Date: Mon, 27 Jul 2026 08:01:16 +0200 Subject: [PATCH] Add split-preserving YOLO scale normalization --- scripts/retile_yolo_dataset.py | 114 +++++++++++++++++++++++++++++++++ 1 file changed, 114 insertions(+) create mode 100644 scripts/retile_yolo_dataset.py diff --git a/scripts/retile_yolo_dataset.py b/scripts/retile_yolo_dataset.py new file mode 100644 index 00000000..60a8e985 --- /dev/null +++ b/scripts/retile_yolo_dataset.py @@ -0,0 +1,114 @@ +#!/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 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("--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.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 = 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) + for top in range(0, height, args.tile_size): + for left in range(0, width, args.tile_size): + 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 x1, y1, x2, y2 in 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 + 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) + + 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_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, + "min_visible_ratio":args.min_visible_ratio, "output_tile_count":len(output_tiles), + "kept_label_count":kept_labels, "dropped_label_count":dropped_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())