mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-07-04 04:52:09 +08:00
firestar5683
83 lines
3.0 KiB
Python
83 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import importlib
|
|
|
|
from pathlib import Path
|
|
|
|
import gdown
|
|
|
|
if __package__ in (None, ""):
|
|
import sys
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
from common import DEFAULT_WORKSPACE, default_raw_root, ensure_dir, resolve_workspace # type: ignore
|
|
else:
|
|
from .common import DEFAULT_WORKSPACE, default_raw_root, ensure_dir, resolve_workspace
|
|
|
|
|
|
GLARE_ROOT_URL = "https://drive.google.com/drive/folders/1gmoOSgvjR4DP7jGfGS_xAmxcMShyeThx?usp=sharing"
|
|
DEFAULT_PREFIXES = ("Images/", "Tracks/")
|
|
MANIFEST_FIELDS = ["file_id", "relative_path", "local_path"]
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Download the raw GLARE image/track files without checkpoint artifacts.")
|
|
parser.add_argument("--workspace", type=Path, default=DEFAULT_WORKSPACE, help="Training workspace root.")
|
|
parser.add_argument("--output-root", type=Path, help="Download destination. Defaults to <raw>/glare_raw.")
|
|
parser.add_argument("--prefix", dest="prefixes", nargs="+", default=list(DEFAULT_PREFIXES), help="Path prefixes to keep from the GLARE Drive tree.")
|
|
parser.add_argument("--manifest", type=Path, help="Optional CSV path for the filtered file manifest.")
|
|
parser.add_argument("--list-only", action="store_true", help="List matching files without downloading them.")
|
|
parser.add_argument("--resume", action="store_true", help="Resume partial downloads and skip completed files.")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
workspace = resolve_workspace(args.workspace)
|
|
output_root = args.output_root.resolve() if args.output_root else (default_raw_root(workspace) / "glare_raw")
|
|
manifest_path = args.manifest.resolve() if args.manifest else (output_root / "manifest.csv")
|
|
ensure_dir(output_root)
|
|
|
|
download_folder_mod = importlib.import_module("gdown.download_folder")
|
|
entries = download_folder_mod.download_folder(
|
|
url=GLARE_ROOT_URL,
|
|
skip_download=True,
|
|
quiet=True,
|
|
remaining_ok=True,
|
|
)
|
|
if entries is None:
|
|
raise RuntimeError("Failed to enumerate the GLARE Drive tree")
|
|
|
|
selected = [entry for entry in entries if any(entry.path.startswith(prefix) for prefix in args.prefixes)]
|
|
print(f"Matched {len(selected)} GLARE file(s) under prefixes: {', '.join(args.prefixes)}")
|
|
|
|
with manifest_path.open("w", encoding="utf-8", newline="") as csv_file:
|
|
writer = csv.DictWriter(csv_file, fieldnames=MANIFEST_FIELDS)
|
|
writer.writeheader()
|
|
for entry in selected:
|
|
local_path = output_root / entry.path
|
|
writer.writerow({
|
|
"file_id": entry.id,
|
|
"relative_path": entry.path,
|
|
"local_path": str(local_path),
|
|
})
|
|
if args.list_only:
|
|
continue
|
|
ensure_dir(local_path.parent)
|
|
gdown.download(
|
|
id=entry.id,
|
|
output=str(local_path),
|
|
quiet=False,
|
|
resume=args.resume,
|
|
)
|
|
|
|
print(f"GLARE manifest: {manifest_path}")
|
|
print(f"GLARE output: {output_root}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|