generate.py

"""Command-line dataset builder for AngleForge.

Reads labelled source images from ``input/<label>/*.jpg`` and builds a
multi-angle image dataset, optionally pushing to Hugging Face and/or Edge
Impulse.
"""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Dict, List

from src import edge_impulse
from src.backends import select_backend
from src.builder import build_dataset
from src.config import DEFAULT_ANGLES, DatasetConfig
from src.hf_export import export_hf_dataset, push_to_hub


def _read_input_dir(input_dir: str) -> Dict[str, List[str]]:
    root = Path(input_dir)
    if not root.exists():
        raise SystemExit(f"Input directory not found: {input_dir}")
    exts = {".jpg", ".jpeg", ".png", ".bmp"}
    classes: Dict[str, List[str]] = {}
    for label_dir in sorted(p for p in root.iterdir() if p.is_dir()):
        imgs = sorted(str(p) for p in label_dir.iterdir() if p.suffix.lower() in exts)
        if imgs:
            classes[label_dir.name] = imgs
    if not classes:
        raise SystemExit(f"No labelled images found under {input_dir}/<label>/*.jpg")
    return classes


def main() -> None:
    parser = argparse.ArgumentParser(description="Build a multi-angle image dataset with AngleForge.")
    parser.add_argument("--input", default="input", help="Input dir with <label>/ subfolders of images.")
    parser.add_argument("--out", default="output", help="Dataset output directory.")
    parser.add_argument("--hf-out", default="hf_dataset", help="Hugging Face imagefolder output directory.")
    parser.add_argument("--dataset-name", default="industrial_angles")
    parser.add_argument("--angles", nargs="*", default=list(DEFAULT_ANGLES))
    parser.add_argument("--variations", type=int, default=1)
    parser.add_argument("--plain-augs", type=int, default=0)
    parser.add_argument("--image-size", type=int, default=512)
    parser.add_argument("--test-ratio", type=float, default=0.2)
    parser.add_argument("--prefer", choices=["auto", "local", "serverless"], default="auto")
    parser.add_argument("--hf-token", default="")
    parser.add_argument("--push-hf-repo", default="")
    parser.add_argument("--hf-private", action="store_true")
    parser.add_argument("--edge-impulse-api-key", default="")
    parser.add_argument("--ei-allow-duplicates", action="store_true")
    args = parser.parse_args()

    classes = _read_input_dir(args.input)
    backend = select_backend(hf_token=args.hf_token, image_size=args.image_size, prefer=args.prefer)
    print(f"Backend: {backend.source}")

    config = DatasetConfig(
        out_dir=args.out,
        dataset_name=args.dataset_name,
        image_size=args.image_size,
        angles=args.angles,
        variations_per_angle=args.variations,
        plain_augmentations_per_image=args.plain_augs,
        test_ratio=args.test_ratio,
    )
    result = build_dataset(config, backend, classes)
    print(f"Built {result.total_images} images in {result.out_dir}")

    export_hf_dataset(config, result, args.hf_out, repo_id=args.push_hf_repo or "your-username/your-dataset")
    print(f"Hugging Face imagefolder: {args.hf_out}")

    if args.push_hf_repo:
        if not args.hf_token:
            print("Skipping HF push: --hf-token is required.")
        else:
            url = push_to_hub(args.hf_out, args.push_hf_repo, args.hf_token, private=args.hf_private)
            print(f"Pushed dataset: {url}")

    if args.edge_impulse_api_key:
        res = edge_impulse.upload_dataset(
            dataset_dir=args.out,
            api_key=args.edge_impulse_api_key,
            allow_duplicates=args.ei_allow_duplicates,
        )
        print(f"Edge Impulse: {res.uploaded} uploaded, {res.failed} failed.")


if __name__ == "__main__":
    main()