78 lines
2.9 KiB
Python
78 lines
2.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Create a checksummed regional YOLO view without copying protected data."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
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 select_paths(summary: dict[str, Any], manifest: dict[str, Any], region: str) -> tuple[list[str], list[str]]:
|
|
samples = {item["sample_slug"]: item for item in manifest["samples"]}
|
|
train: list[str] = []
|
|
val: list[str] = []
|
|
for tile in summary["tiles"]:
|
|
sample = samples[tile["sample_slug"]]
|
|
if sample["region"] != region or not tile.get("kept", True):
|
|
continue
|
|
if sample["split"] == "train" and tile["split"] == "train":
|
|
train.append(tile["image_path"])
|
|
elif sample["split"] == "val" and tile["split"] == "val":
|
|
val.append(tile["image_path"])
|
|
if not train or not val:
|
|
raise ValueError(f"Region {region!r} must contain train and validation images")
|
|
return sorted(train), sorted(val)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--summary", type=Path, required=True)
|
|
parser.add_argument("--corpus-manifest", type=Path, required=True)
|
|
parser.add_argument("--region", required=True)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
summary = json.loads(args.summary.read_text(encoding="utf-8"))
|
|
manifest = json.loads(args.corpus_manifest.read_text(encoding="utf-8"))
|
|
train, val = select_paths(summary, manifest, args.region)
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
train_list = args.output_dir / "train.txt"
|
|
val_list = args.output_dir / "val.txt"
|
|
train_list.write_text("\n".join(train) + "\n", encoding="utf-8")
|
|
val_list.write_text("\n".join(val) + "\n", encoding="utf-8")
|
|
dataset_yaml = args.output_dir / "dataset.yaml"
|
|
dataset_yaml.write_text(
|
|
f"path: {args.output_dir}\ntrain: {train_list}\nval: {val_list}\nnames:\n 0: building\n",
|
|
encoding="utf-8",
|
|
)
|
|
evidence = {
|
|
"schema_version": 1,
|
|
"status": "ok",
|
|
"region": args.region,
|
|
"train_image_count": len(train),
|
|
"validation_image_count": len(val),
|
|
"summary": str(args.summary),
|
|
"summary_sha256": sha256(args.summary),
|
|
"corpus_manifest": str(args.corpus_manifest),
|
|
"corpus_manifest_sha256": sha256(args.corpus_manifest),
|
|
"dataset_yaml": str(dataset_yaml),
|
|
"protected_splits_in_training": [],
|
|
}
|
|
(args.output_dir / "regional-dataset.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())
|