Jens
faeb58ef6d
GeoIntel release gates / Compile, test, contracts and builds (push) Successful in 1m49s
GeoIntel release gates / Python and npm vulnerability policy (push) Successful in 21s
GeoIntel release gates / Production AI image, SBOM and container scan (push) Successful in 5m39s
GeoIntel release gates / Deploy exact gated revision to Unraid (push) Failing after 58m43s
135 lines
6.5 KiB
Python
135 lines
6.5 KiB
Python
#!/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())
|