File size: 2,994 Bytes
9274a8d 529f79d 9274a8d 529f79d 9274a8d | 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 89 90 | """Download the dataset from the Hugging Face Hub and make it trainable.
The inverse of push.py: fetch the repo, unpack the shards into the layout
Ultralytics expects, and rewrite data.yaml with an absolute path so the training
script works from any working directory.
Usage:
python -m src.hub.pull --repo-id <user>/cone-distance-v1 --dest data/hub
python -m src.train.train --data data/hub/data.yaml
"""
from __future__ import annotations
import argparse
import os
import tarfile
from pathlib import Path
import yaml
def extract_shards(dest: Path) -> None:
shard_dir = dest / "shards"
shards = sorted(shard_dir.glob("*.tar"))
if not shards:
raise SystemExit(f"no shards under {shard_dir}")
for shard in shards:
with tarfile.open(shard) as archive:
archive.extractall(dest)
print(f" extracted {shard.name}")
def rewrite_data_yaml(dest: Path) -> Path:
"""Point `path` at wherever the dataset actually landed."""
path = dest / "data.yaml"
config = yaml.safe_load(path.read_text())
config["path"] = str(dest.resolve())
with open(path, "w") as handle:
yaml.safe_dump(config, handle, sort_keys=False)
return path
def pull(repo_id: str, dest: Path, token: str | None, keep_shards: bool) -> None:
from huggingface_hub import snapshot_download
dest.mkdir(parents=True, exist_ok=True)
print(f"downloading {repo_id} -> {dest}")
snapshot_download(
repo_id=repo_id,
repo_type="dataset",
local_dir=str(dest),
token=token,
)
print("unpacking")
extract_shards(dest)
if not keep_shards:
for shard in (dest / "shards").glob("*.tar"):
shard.unlink()
print(" removed shards (pass --keep-shards to keep them)")
data_yaml = rewrite_data_yaml(dest)
for split in ("train", "val"):
images = list((dest / split / "images").iterdir())
print(f" {split}: {len(images)} images")
print(f"\ndata config: {data_yaml}")
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--repo-id", required=True)
parser.add_argument("--dest", type=Path, default=Path("data/hub"))
parser.add_argument("--token", default=os.environ.get("HF_TOKEN"),
help="defaults to $HF_TOKEN, then the cached "
"huggingface-cli login. Prefer the environment: a "
"token passed on the command line is visible to "
"anyone who can run ps.")
parser.add_argument("--keep-shards", action="store_true",
help="keep the .tar files after extracting (doubles disk use)")
args = parser.parse_args()
pull(repo_id=args.repo_id, dest=args.dest, token=args.token,
keep_shards=args.keep_shards)
if __name__ == "__main__":
main()
|