Spaces:
Running on Zero
Running on Zero
File size: 3,782 Bytes
3046f29 | 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 91 92 | """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()
|