100 lines
3.9 KiB
Python
100 lines
3.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Create a leak-free train list for a regional YOLO expert."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
|
|
|
|
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("--corpus-manifest", type=Path, required=True)
|
|
parser.add_argument("--region", required=True)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument("--positive-repeat", type=int, default=2)
|
|
parser.add_argument("--negative-repeat", type=int, default=1)
|
|
parser.add_argument(
|
|
"--other-region-repeat",
|
|
type=int,
|
|
default=0,
|
|
help="Include each train tile outside the expert region this many times.",
|
|
)
|
|
args = parser.parse_args()
|
|
if args.positive_repeat < 1 or args.negative_repeat < 1 or args.other_region_repeat < 0:
|
|
raise SystemExit("regional repeats must be positive and other-region repeat non-negative")
|
|
|
|
summary = json.loads(args.summary.read_text(encoding="utf-8"))
|
|
manifest = json.loads(args.corpus_manifest.read_text(encoding="utf-8"))
|
|
samples = {item["sample_slug"]: item for item in manifest["samples"]}
|
|
paths: list[str] = []
|
|
selected_samples: set[str] = set()
|
|
positive_tiles = negative_tiles = other_region_tiles = 0
|
|
for tile in summary["tiles"]:
|
|
sample = samples[tile["sample_slug"]]
|
|
if sample["split"] != "train" or tile["split"] != "train":
|
|
continue
|
|
positive = int(tile.get("label_count") or 0) > 0
|
|
if sample["region"] == args.region:
|
|
repeat = args.positive_repeat if positive else args.negative_repeat
|
|
positive_tiles += int(positive)
|
|
negative_tiles += int(not positive)
|
|
else:
|
|
repeat = args.other_region_repeat
|
|
other_region_tiles += int(repeat > 0)
|
|
if repeat == 0:
|
|
continue
|
|
paths.extend([str(Path(tile["image_path"]).resolve())] * repeat)
|
|
selected_samples.add(tile["sample_slug"])
|
|
if not paths or not positive_tiles:
|
|
raise SystemExit(f"no positive train tiles found for region {args.region!r}")
|
|
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
train_list = args.output_dir / "train.txt"
|
|
train_list.write_text("\n".join(paths) + "\n", encoding="utf-8")
|
|
dataset_yaml = args.output_dir / "dataset.yaml"
|
|
dataset_yaml.write_text(
|
|
f"path: {args.output_dir}\ntrain: {train_list}\n"
|
|
f"val: {args.summary.parent / 'images' / 'val'}\nnames:\n 0: building\n",
|
|
encoding="utf-8",
|
|
)
|
|
evidence = {
|
|
"schema_version": 1,
|
|
"status": "ok",
|
|
"region": args.region,
|
|
"summary": str(args.summary),
|
|
"summary_sha256": sha256(args.summary),
|
|
"corpus_manifest": str(args.corpus_manifest),
|
|
"corpus_manifest_sha256": sha256(args.corpus_manifest),
|
|
"positive_repeat": args.positive_repeat,
|
|
"negative_repeat": args.negative_repeat,
|
|
"other_region_repeat": args.other_region_repeat,
|
|
"source_positive_tile_count": positive_tiles,
|
|
"source_negative_tile_count": negative_tiles,
|
|
"source_other_region_tile_count": other_region_tiles,
|
|
"sampled_train_entry_count": len(paths),
|
|
"selected_train_samples": sorted(selected_samples),
|
|
"protected_samples_in_training": [],
|
|
"train_list": str(train_list),
|
|
"dataset_yaml": str(dataset_yaml),
|
|
}
|
|
evidence_path = args.output_dir / "regional-expert-dataset.json"
|
|
evidence_path.write_text(json.dumps(evidence, indent=2), encoding="utf-8")
|
|
print(json.dumps(evidence, indent=2))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|