scripts/upload_hf_dataset.py

"""Upload Golden-0-to-25 + Golden-25plus raw NIfTI datasets to Hugging Face.

Uploads manifests first, then NIfTI files folder-by-folder with progress.
Resumable: HF Hub skips files already uploaded (content-addressed by SHA).
"""
from __future__ import annotations
import os, sys, time
from pathlib import Path
from huggingface_hub import HfApi, CommitOperationAdd

TOKEN = "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"
REPO_ID = "bilalahmad176176/BrainAge-Golden-Raw"
REPO_TYPE = "dataset"

SPLITS = [
    Path("/home/MRI-DataSet/Golden-0-to-25"),
    Path("/home/MRI-DataSet/Golden-25plus"),
]

BATCH_SIZE = 80  # files per commit (keeps commit payloads manageable)


def collect_files(root: Path) -> list[tuple[Path, str]]:
    """Return (local_path, repo_path) pairs for all files under root."""
    pairs = []
    for p in sorted(root.rglob("*")):
        if p.is_file():
            rel = p.relative_to(root.parent)
            pairs.append((p, str(rel)))
    return pairs


def upload_split(api: HfApi, root: Path):
    pairs = collect_files(root)
    total = len(pairs)
    print(f"\n{'='*60}")
    print(f"Uploading {root.name}: {total} files")
    print(f"{'='*60}")

    for i in range(0, total, BATCH_SIZE):
        batch = pairs[i : i + BATCH_SIZE]
        ops = []
        for local, repo_path in batch:
            ops.append(CommitOperationAdd(
                path_in_repo=repo_path,
                path_or_fileobj=str(local),
            ))
        n = min(i + BATCH_SIZE, total)
        msg = f"Add {root.name} files {i+1}–{n} of {total}"
        print(f"  [{n}/{total}] committing batch … ", end="", flush=True)
        t0 = time.time()
        try:
            api.create_commit(
                repo_id=REPO_ID,
                repo_type=REPO_TYPE,
                operations=ops,
                commit_message=msg,
            )
            print(f"done ({time.time()-t0:.0f}s)")
        except Exception as e:
            print(f"ERROR: {e}")
            print("  (will retry this batch once)")
            time.sleep(5)
            try:
                api.create_commit(
                    repo_id=REPO_ID,
                    repo_type=REPO_TYPE,
                    operations=ops,
                    commit_message=msg + " (retry)",
                )
                print(f"  retry succeeded")
            except Exception as e2:
                print(f"  retry also failed: {e2}")
                print(f"  skipping batch, re-run script to resume.")


def upload_readme(api: HfApi):
    readme = """---
license: cc-by-nc-4.0
task_categories:
  - image-classification
  - other
task_ids:
  - brain-age-prediction
language:
  - en
pretty_name: BrainAge Golden Raw MRI Dataset
size_categories:
  - 1K<n<10K
tags:
  - neuroimaging
  - mri
  - brain-age
  - t1w
  - nifti
  - pediatric
  - adult
---

# BrainAge Golden Raw MRI Dataset

Curated collection of **6,152 healthy-brain T1-weighted MRI scans** spanning
ages 0–86 years, assembled from 12 public neuroimaging datasets.

## Splits

| Split | Subjects | Age range | Size |
|-------|----------|-----------|------|
| `Golden-0-to-25/` | 4,782 | 0 – 25 y | ~42 GB |
| `Golden-25plus/` | 1,370 | 25 – 86 y | ~13 GB |

## Source datasets

BCP, Calgary, ds002726, ds000248, PTBP, IXI, MPI-Leipzig,
AOMIC-ID1000, NKI-Rockland, ABIDE-I, ABIDE-II, ADHD-200.

## File format

Each scan is a `.nii.gz` NIfTI file (native space, T1w).
Manifests (`manifest.csv`) list subject IDs, dataset of origin,
chronological age, sex, split, and file paths.

## Intended use

Training and evaluating brain-age prediction models on healthy controls.

## Citation

If you use this dataset, please cite the original source studies
(listed in each manifest row under the `dataset` column).
"""
    api.upload_file(
        path_or_fileobj=readme.encode(),
        path_in_repo="README.md",
        repo_id=REPO_ID,
        repo_type=REPO_TYPE,
        commit_message="Add dataset card",
    )
    print("README.md uploaded.")


def main():
    api = HfApi(token=TOKEN)
    print(f"Authenticated as: {api.whoami()['name']}")
    print(f"Target repo: https://huggingface.co/datasets/{REPO_ID}")

    upload_readme(api)

    for split_dir in SPLITS:
        upload_split(api, split_dir)

    print(f"\n{'='*60}")
    print(f"DONE — https://huggingface.co/datasets/{REPO_ID}")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()