scripts/download_data.py

"""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()