88 lines
3.0 KiB
Python
88 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Create a checksummed luminance-normalized YOLO dataset without changing labels/splits."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
from PIL import Image
|
|
|
|
|
|
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 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("--force", action="store_true")
|
|
args = parser.parse_args()
|
|
if args.output_dir.exists():
|
|
if not args.force:
|
|
raise SystemExit(f"Output exists: {args.output_dir}")
|
|
shutil.rmtree(args.output_dir)
|
|
args.output_dir.mkdir(parents=True)
|
|
|
|
summary = json.loads(args.summary.read_text(encoding="utf-8"))
|
|
converted = 0
|
|
for tile in summary["tiles"]:
|
|
split = tile["split"]
|
|
source_image = Path(tile["image_path"])
|
|
source_label = Path(tile["label_path"])
|
|
image_target = args.output_dir / "images" / split / source_image.name
|
|
label_target = args.output_dir / "labels" / split / source_label.name
|
|
image_target.parent.mkdir(parents=True, exist_ok=True)
|
|
label_target.parent.mkdir(parents=True, exist_ok=True)
|
|
with Image.open(source_image) as image:
|
|
image.convert("L").convert("RGB").save(image_target)
|
|
shutil.copyfile(source_label, label_target)
|
|
tile["image_path"] = str(image_target)
|
|
tile["label_path"] = str(label_target)
|
|
converted += 1
|
|
|
|
dataset_yaml = args.output_dir / "dataset.yaml"
|
|
dataset_yaml.write_text(
|
|
f"path: {args.output_dir}\ntrain: images/train\nval: images/val\nnames:\n 0: building\n",
|
|
encoding="utf-8",
|
|
)
|
|
summary.update(
|
|
{
|
|
"output_dir": str(args.output_dir),
|
|
"dataset_yaml": str(dataset_yaml),
|
|
"preprocessing": "luminance_rgb_replicated",
|
|
"source_summary": str(args.summary),
|
|
"source_summary_sha256": sha256(args.summary),
|
|
}
|
|
)
|
|
output_summary = args.output_dir / "yolo_tile_dataset_summary.json"
|
|
output_summary.write_text(json.dumps(summary, indent=2), encoding="utf-8")
|
|
evidence = {
|
|
"schema_version": 1,
|
|
"status": "ok",
|
|
"preprocessing": "luminance_rgb_replicated",
|
|
"source_summary": str(args.summary),
|
|
"source_summary_sha256": sha256(args.summary),
|
|
"converted_tile_count": converted,
|
|
"output_summary": str(output_summary),
|
|
"output_summary_sha256": sha256(output_summary),
|
|
"dataset_yaml": str(dataset_yaml),
|
|
}
|
|
(args.output_dir / "grayscale-dataset-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())
|