Spaces:
Sleeping
Sleeping
File size: 3,317 Bytes
d4a30e5 | 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 83 84 85 86 87 88 | """Fetch the PPE dataset into work/ so train/val/test resolve for Ultralytics.
The dataset is NOT committed to git (it's ~5k images). This script reconstructs
it locally / in Colab / in CI. From repo root:
uv run --with gdown python scripts/download_data.py
DATA_URL="https://drive.google.com/uc?id=<FILE_ID>" uv run --with gdown python scripts/download_data.py
`--with gdown` installs gdown just for this run without touching the project deps.
Source options (set DATA_URL or --url):
* A Google Drive ZIP (recommended) — a single file, no download caps. Zip your
`data/` folder once, share it, and pass its `https://drive.google.com/uc?id=...` URL.
* A Google Drive FOLDER — works only if it has <50 files (gdown's folder limit),
so not suitable for the full image set. Kept as a fallback for small mirrors.
* Roboflow — the original source; `pip install roboflow` then use the versioned
download snippet from your Roboflow project for the most reproducible pull.
Expected final layout (what work/data.yaml points at):
work/train/ work/valid/ work/test/
"""
from __future__ import annotations
import argparse
import os
import sys
import zipfile
from pathlib import Path
# Documented Drive *folder* mirror (see README). Prefer a zip via DATA_URL.
DEFAULT_URL = "https://drive.google.com/drive/folders/1Vhf9QP4WG9yF11EeFUtmFKDU7sKut69f"
DEST = Path("work")
def _require_gdown():
try:
import gdown # noqa: PLC0415
except ImportError:
sys.exit(
"gdown is not installed. Re-run with uv's ephemeral install:\n"
" uv run --with gdown python scripts/download_data.py"
)
return gdown
def _extract_and_flatten(zip_path: Path) -> None:
with zipfile.ZipFile(zip_path) as zf:
zf.extractall(DEST)
# If the zip wrapped everything in a single top dir, surface train/valid/test.
for split in ("train", "valid", "test"):
if (DEST / split).exists():
continue
matches = list(DEST.glob(f"*/{split}"))
if matches:
matches[0].rename(DEST / split)
def main() -> None:
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--url", default=os.environ.get("DATA_URL", DEFAULT_URL))
args = p.parse_args()
gdown = _require_gdown()
DEST.mkdir(exist_ok=True)
if "/folders/" in args.url:
print("⚠️ Downloading a Drive FOLDER — gdown caps this at 50 files, so the")
print(" full image set will be truncated. Prefer a zip via DATA_URL.\n")
gdown.download_folder(args.url, output=str(DEST), quiet=False, use_cookies=False)
else:
zip_path = DEST / "dataset.zip"
gdown.download(args.url, str(zip_path), quiet=False, fuzzy=True)
print(f"Extracting {zip_path} → {DEST}/ ...")
_extract_and_flatten(zip_path)
zip_path.unlink(missing_ok=True)
missing = [s for s in ("train", "valid", "test") if not (DEST / s).exists()]
if missing:
print(f"\n⚠️ Done, but these splits are missing under work/: {missing}")
print(" Check the archive layout — work/data.yaml expects work/train, work/valid, work/test.")
else:
print("\n✅ Dataset ready: work/train, work/valid, work/test")
if __name__ == "__main__":
main()
|