Spaces:
Sleeping
Sleeping
File size: 2,160 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 | from __future__ import annotations
import argparse
from pathlib import Path
import sys
REPO_ROOT = Path(__file__).resolve().parents[2]
SRC_DIR = REPO_ROOT / "src"
if str(SRC_DIR) not in sys.path:
sys.path.insert(0, str(SRC_DIR))
from data.io_utils import ensure_dir, extract_archive, load_yaml # noqa: E402
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Extract Visual Genome raw archives.")
parser.add_argument(
"--config",
type=Path,
default=REPO_ROOT / "configs" / "dataset.yaml",
help="Path to dataset config yaml.",
)
parser.add_argument(
"--force-extract",
action="store_true",
help="Force extraction even if marker exists.",
)
return parser.parse_args()
def main() -> None:
args = parse_args()
cfg = load_yaml(args.config)
dataset_cfg = cfg["dataset"]
extraction_cfg = cfg.get("extraction", {})
if not extraction_cfg.get("enabled", True):
print("Extraction disabled in config. Nothing to do.")
return
raw_root = REPO_ROOT / dataset_cfg["root_dir"]
ensure_dir(raw_root)
marker_file = raw_root / extraction_cfg.get("marker_file", ".extract_complete")
force_extract = args.force_extract or bool(extraction_cfg.get("force_extract", False))
safe_extract = bool(extraction_cfg.get("safe_extract", True))
if marker_file.exists() and not force_extract:
print(f"[skip] Extraction marker found: {marker_file}")
return
archives = extraction_cfg.get("archives", [])
if not archives:
raise RuntimeError("No extraction archives configured.")
for archive_name in archives:
archive_path = raw_root / archive_name
if not archive_path.exists():
raise RuntimeError(f"Archive missing: {archive_path}")
print(f"[extract] {archive_name}")
extract_archive(archive_path=archive_path, output_dir=raw_root, safe_extract=safe_extract)
marker_file.write_text("ok\n", encoding="utf-8")
print(f"[ok] Extraction completed. Marker written at {marker_file}")
if __name__ == "__main__":
main()
|