Files
caption-test/scripts/make_split.py
2026-05-08 09:52:15 +03:00

88 lines
2.5 KiB
Python

from __future__ import annotations
"""Create 80/20 train/test split from GTA-UAV-LR pair JSONs.
Merges cross-area train+test (33,708 pairs), shuffles deterministically,
and saves new 80/20 split JSONs.
Usage:
python -m scripts.make_split [--ratio 0.8] [--seed 42]
"""
import argparse
import json
import logging
import random
from pathlib import Path
import coloredlogs
LOGGER = logging.getLogger("caption_test.make_split")
_RGB_ROOT = Path("/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR")
def main() -> None:
parser = argparse.ArgumentParser(description="Create 80/20 split for GTA-UAV-LR.")
parser.add_argument("--ratio", type=float, default=0.8, help="Train ratio (default 0.8).")
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
parser.add_argument(
"--output-dir", type=str, default="meta",
help="Output directory for split JSONs.",
)
args = parser.parse_args()
coloredlogs.install(
level="INFO", logger=LOGGER,
fmt="%(asctime)s %(name)s %(levelname)s %(message)s",
)
# Load both original splits.
train_path = _RGB_ROOT / "cross-area-drone2sate-train.json"
test_path = _RGB_ROOT / "cross-area-drone2sate-test.json"
LOGGER.info("📂 Loading %s", train_path.name)
with open(train_path) as f:
part1 = json.load(f)
LOGGER.info("📂 Loading %s", test_path.name)
with open(test_path) as f:
part2 = json.load(f)
all_pairs = part1 + part2
LOGGER.info("📊 Total pairs: %d", len(all_pairs))
# Shuffle deterministically.
rng = random.Random(args.seed)
rng.shuffle(all_pairs)
# Split.
n_train = int(len(all_pairs) * args.ratio)
train_pairs = all_pairs[:n_train]
test_pairs = all_pairs[n_train:]
LOGGER.info(
"✂️ Split %.0f/%.0f: train=%d (%.1f%%) test=%d (%.1f%%)",
args.ratio * 100, (1 - args.ratio) * 100,
len(train_pairs), 100 * len(train_pairs) / len(all_pairs),
len(test_pairs), 100 * len(test_pairs) / len(all_pairs),
)
# Save.
out_dir = Path(args.output_dir)
out_dir.mkdir(parents=True, exist_ok=True)
train_out = out_dir / "train_80.json"
test_out = out_dir / "test_20.json"
with train_out.open("w", encoding="utf-8") as f:
json.dump(train_pairs, f)
with test_out.open("w", encoding="utf-8") as f:
json.dump(test_pairs, f)
LOGGER.info("💾 Saved: %s (%d pairs)", train_out, len(train_pairs))
LOGGER.info("💾 Saved: %s (%d pairs)", test_out, len(test_pairs))
if __name__ == "__main__":
main()