src/hf_export.py

"""Hugging Face export for AngleForge image datasets.

Produces a standard ``imagefolder`` layout (``train/<label>/*.jpg``) plus a
``metadata.csv`` and a dataset card, and can push it to the Hub.
"""

from __future__ import annotations

import csv
import json
import shutil
from pathlib import Path
from typing import Dict, List

from .builder import BuildResult
from .config import ANGLE_PRESETS, DatasetConfig


def _dataset_card(config: DatasetConfig, result: BuildResult, repo_id: str) -> str:
    labels_table = "\n".join(
        f"| `{label}` | {count} |" for label, count in sorted(result.label_counts.items())
    )
    angles_list = "\n".join(
        f"- `{a}` — {ANGLE_PRESETS.get(a, (a, ''))[0]}" for a in config.angles
    )
    return f"""---
license: cc-by-4.0
pretty_name: {config.dataset_name} multi-angle robotic-arm image dataset
task_categories:
- image-classification
tags:
- image
- robotics
- robotic-arm
- synthetic-data
- multi-view
- qwen-image-edit
- edge-impulse
size_categories:
- n<1K
---

# {config.dataset_name} — Multi-Angle Robotic-Arm Image Dataset

Synthetic multi-viewpoint image dataset generated from real-world photos with
**Qwen Image Edit** ({result.backend_source}). Each source image is re-rendered
from several camera/gripper viewpoints to simulate a robotic arm inspecting an
object from multiple angles.

## Classes

| Label | Images |
|---|---|
{labels_table}

## Viewpoints (angles)

{angles_list}

## Layout

```text
train/<label>/<label>.<id>.jpg
test/<label>/<label>.<id>.jpg
metadata.csv
```

## Loading

```python
from datasets import load_dataset

ds = load_dataset("imagefolder", data_dir="{repo_id.split('/')[-1]}")
# or, once pushed to the Hub:
ds = load_dataset("{repo_id}")
print(ds)
```

## Edge Impulse

Filenames use the `label.<id>.jpg` convention, so they upload directly:

```bash
edge-impulse-uploader --category training train/**/*.jpg
```

## Notes

Synthetic multi-view images are a bootstrap for robotic-arm perception and
inspection models. Validate with real captures from the arm's own camera
before deployment.
"""


def export_hf_dataset(
    config: DatasetConfig,
    result: BuildResult,
    hf_dir: str,
    repo_id: str = "your-username/your-dataset",
) -> str:
    """Assemble an imagefolder dataset from a completed build. Returns its path."""
    source_dir = Path(result.out_dir)
    hf_path = Path(hf_dir).resolve()
    if hf_path.exists():
        shutil.rmtree(hf_path)
    hf_path.mkdir(parents=True, exist_ok=True)

    # Copy the imagefolder tree.
    src_imagefolder = source_dir / "hf_imagefolder"
    rows: List[Dict[str, str]] = []
    for split in ("train", "test"):
        split_src = src_imagefolder / split
        if not split_src.exists():
            continue
        for label_dir in sorted(p for p in split_src.iterdir() if p.is_dir()):
            for img in sorted(label_dir.glob("*.jpg")):
                rel = Path(split) / label_dir.name / img.name
                dst = hf_path / rel
                dst.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy2(img, dst)
                rows.append({"file_name": str(rel), "label": label_dir.name, "split": split})

    with (hf_path / "metadata.csv").open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=["file_name", "label", "split"])
        writer.writeheader()
        writer.writerows(rows)

    for name in ("dataset_summary.json",):
        src = source_dir / name
        if src.exists():
            shutil.copy2(src, hf_path / name)

    (hf_path / "README.md").write_text(_dataset_card(config, result, repo_id), encoding="utf-8")
    return str(hf_path)


def push_to_hub(hf_dir: str, repo_id: str, token: str, private: bool = False) -> str:
    """Upload the HF image dataset folder to the Hub. Returns the dataset URL."""
    from huggingface_hub import HfApi

    api = HfApi(token=token)
    api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True, private=private)
    api.upload_folder(folder_path=hf_dir, repo_id=repo_id, repo_type="dataset")
    return f"https://huggingface.co/datasets/{repo_id}"