Spaces:
Sleeping
Sleeping
File size: 2,542 Bytes
4687353 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | from __future__ import annotations
import argparse
import subprocess
import sys
from pathlib import Path
from typing import List
REPO_ROOT = Path(__file__).resolve().parents[2]
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Run chair-only Visual Genome data pipeline.")
parser.add_argument(
"--config",
type=Path,
default=REPO_ROOT / "configs" / "dataset.yaml",
help="Path to dataset config yaml.",
)
parser.add_argument(
"--dataset-name",
type=str,
default="visual_genome",
help="Dataset selector. Only visual_genome is supported in this version.",
)
parser.add_argument("--download-only", action="store_true")
parser.add_argument("--extract-only", action="store_true")
parser.add_argument("--process-only", action="store_true")
parser.add_argument("--split-only", action="store_true")
parser.add_argument("--all", action="store_true", help="Run full pipeline (default behavior).")
return parser.parse_args()
def run_step(script_path: Path, config_path: Path) -> None:
cmd: List[str] = [sys.executable, str(script_path), "--config", str(config_path)]
print(f"[run] {' '.join(cmd)}")
subprocess.run(cmd, check=True, cwd=REPO_ROOT)
def main() -> None:
args = parse_args()
if args.dataset_name != "visual_genome":
raise RuntimeError("Only --dataset-name visual_genome is supported currently.")
selected_flags = [
args.download_only,
args.extract_only,
args.process_only,
args.split_only,
args.all,
]
if sum(1 for flag in selected_flags if flag) > 1:
raise RuntimeError("Use only one stage selector at a time.")
scripts_dir = REPO_ROOT / "scripts" / "data"
steps = {
"download": scripts_dir / "download_visual_genome.py",
"extract": scripts_dir / "extract_visual_genome.py",
"process": scripts_dir / "process_visual_genome.py",
"split": scripts_dir / "build_splits.py",
}
if args.download_only:
pipeline = ["download"]
elif args.extract_only:
pipeline = ["extract"]
elif args.process_only:
pipeline = ["process"]
elif args.split_only:
pipeline = ["split"]
else:
pipeline = ["download", "extract", "process", "split"]
for step in pipeline:
run_step(script_path=steps[step], config_path=args.config)
print("[ok] Data pipeline completed.")
if __name__ == "__main__":
main()
|