"""Gradio Hugging Face Space: wake-word dataset creator. Generates a keyword-spotting dataset using Google Cloud TTS when an API key is supplied, and automatically falls back to free Piper TTS otherwise. Optionally pushes the result to a Hugging Face dataset repo and/or uploads directly to an Edge Impulse project. """ from __future__ import annotations import os import shutil import tempfile from pathlib import Path from typing import List, Optional import gradio as gr from src import edge_impulse from src.backends import select_backend from src.builder import build_dataset from src.config import ( DEFAULT_UNKNOWN_PHRASES, DEFAULT_WAKE_PHRASES, DatasetConfig, ) from src.hf_export import export_hf_dataset, push_to_hub # API keys can also be provided as Space secrets. ENV_GCP_KEY = os.environ.get("GCP_TTS_API_KEY", "") ENV_HF_TOKEN = os.environ.get("HF_TOKEN", "") ENV_EI_KEY = os.environ.get("EDGE_IMPULSE_API_KEY", "") def _split_lines(text: str, fallback: List[str]) -> List[str]: items = [line.strip() for line in (text or "").splitlines() if line.strip()] return items or list(fallback) def create_dataset( dataset_name: str, wake_label: str, wake_phrases_text: str, unknown_phrases_text: str, gcp_api_key: str, base_repeats: int, augmentations: int, background_noise: int, max_voices: int, test_ratio: float, hf_repo_id: str, hf_token: str, hf_private: bool, do_push_hf: bool, ei_api_key: str, do_upload_ei: bool, ei_allow_duplicates: bool, progress=gr.Progress(track_tqdm=False), ): logs: List[str] = [] def log(message: str) -> str: logs.append(message) return "\n".join(logs) work_root = Path(tempfile.mkdtemp(prefix="wakeword_")) dataset_dir = work_root / "dataset" hf_dir = work_root / "hf_dataset" gcp_api_key = (gcp_api_key or "").strip() or ENV_GCP_KEY try: progress(0.05, desc="Selecting TTS backend") log("Selecting TTS backend...") backend = select_backend( gcp_api_key=gcp_api_key, language_prefixes=["en", "nl", "de", "fr", "es"], max_gcp_voices_per_locale=3, max_piper_voices=int(max_voices), sample_rate_hz=16000, ) engine = ( "Google Cloud TTS" if backend.source == "google_cloud_tts" else "Piper TTS (free fallback)" ) yield log(f"Using backend: {engine}"), None, None config = DatasetConfig( out_dir=str(dataset_dir), dataset_name=dataset_name or "hey_android", wake_label=wake_label or "hey_android", wake_phrases=_split_lines(wake_phrases_text, DEFAULT_WAKE_PHRASES), unknown_phrases=_split_lines(unknown_phrases_text, DEFAULT_UNKNOWN_PHRASES), base_repeats_per_phrase_per_voice=int(base_repeats), augmentations_per_speech_clip=int(augmentations), background_noise_samples=int(background_noise), max_piper_voices=int(max_voices), test_ratio=float(test_ratio), ) progress(0.15, desc="Generating audio") result = build_dataset(config, backend, progress=lambda m: logs.append(m)) yield log( f"Generated {result.total_samples} samples " f"(base={result.generated_base}, augmented={result.generated_augmented}, " f"failed={result.failed})." ), None, None progress(0.7, desc="Preparing Hugging Face layout") export_hf_dataset(config, result, str(hf_dir), repo_id=hf_repo_id or "your-username/your-dataset") log("Hugging Face dataset folder prepared.") # Zip for download. zip_base = work_root / f"{config.dataset_name}_hf_dataset" zip_path = shutil.make_archive(str(zip_base), "zip", str(hf_dir)) yield log(f"Created download archive: {Path(zip_path).name}"), zip_path, None # Optional: push to Hugging Face Hub. token = (hf_token or "").strip() or ENV_HF_TOKEN if do_push_hf: if not token: log("Skipping HF push: no token provided.") elif not hf_repo_id or "/" not in hf_repo_id: log("Skipping HF push: provide a repo id like 'username/dataset-name'.") else: progress(0.85, desc="Pushing to Hugging Face") log(f"Pushing to Hugging Face dataset '{hf_repo_id}'...") url = push_to_hub(str(hf_dir), hf_repo_id, token, private=bool(hf_private)) log(f"Pushed: {url}") yield "\n".join(logs), zip_path, None # Optional: upload to Edge Impulse. ei_key = (ei_api_key or "").strip() or ENV_EI_KEY if do_upload_ei: if not ei_key: log("Skipping Edge Impulse upload: no API key provided.") else: progress(0.92, desc="Uploading to Edge Impulse") log("Uploading dataset to your Edge Impulse project...") ei_result = edge_impulse.upload_dataset( dataset_dir=str(dataset_dir), api_key=ei_key, allow_duplicates=bool(ei_allow_duplicates), progress=lambda m: logs.append(m), ) log( f"Edge Impulse: {ei_result.uploaded} uploaded, {ei_result.failed} failed." ) if ei_result.errors: log("Edge Impulse errors:\n" + "\n".join(ei_result.errors[:5])) progress(1.0, desc="Done") summary = ( f"### Done\n" f"- Backend: **{engine}**\n" f"- Total samples: **{result.total_samples}**\n" + "\n".join(f"- `{k}`: {v}" for k, v in sorted(result.label_counts.items())) ) yield "\n".join(logs), zip_path, summary except Exception as exc: # noqa: BLE001 - surface errors to the UI log(f"ERROR: {exc}") yield "\n".join(logs), None, f"### Failed\n\n```\n{exc}\n```" with gr.Blocks(title="WakeForge — GCP & Piper TTS Wake Word Dataset Creator") as demo: gr.Markdown( """ # 🔨 WakeForge ### GCP & Piper TTS Wake Word Dataset Creator Generate a keyword-spotting dataset for **Hugging Face** and **Edge Impulse**. - Provide a **Google Cloud TTS API key** to use Google voices. - **No key? It automatically falls back to free Piper TTS.** - Optionally **push to a Hugging Face dataset** and/or **upload to your Edge Impulse project**. """ ) with gr.Row(): with gr.Column(): gr.Markdown("### Dataset") dataset_name = gr.Textbox(label="Dataset name", value="hey_android") wake_label = gr.Textbox(label="Wake label", value="hey_android") wake_phrases_text = gr.Textbox( label="Wake phrases (one per line)", value="\n".join(DEFAULT_WAKE_PHRASES), lines=6, ) unknown_phrases_text = gr.Textbox( label="Unknown / near-miss phrases (one per line)", value="\n".join(DEFAULT_UNKNOWN_PHRASES), lines=8, ) gr.Markdown("### Size") base_repeats = gr.Slider(1, 5, value=1, step=1, label="Base clips per phrase per voice") augmentations = gr.Slider(0, 20, value=8, step=1, label="Augmentations per clip") background_noise = gr.Slider(0, 500, value=200, step=10, label="Background noise clips") max_voices = gr.Slider(1, 7, value=7, step=1, label="Max voices") test_ratio = gr.Slider(0.05, 0.5, value=0.2, step=0.05, label="Test split ratio") with gr.Column(): gr.Markdown("### Google Cloud TTS (optional)") gcp_api_key = gr.Textbox( label="GCP TTS API key", type="password", placeholder="Leave blank to use free Piper TTS", ) gr.Markdown("### Push to Hugging Face (optional)") do_push_hf = gr.Checkbox(label="Push dataset to Hugging Face Hub", value=False) hf_repo_id = gr.Textbox(label="HF dataset repo id", placeholder="username/dataset-name") hf_token = gr.Textbox(label="HF write token", type="password", placeholder="hf_...") hf_private = gr.Checkbox(label="Private dataset", value=False) gr.Markdown("### Upload to Edge Impulse (optional)") do_upload_ei = gr.Checkbox(label="Upload dataset to Edge Impulse project", value=False) ei_api_key = gr.Textbox( label="Edge Impulse API key", type="password", placeholder="ei_... (Project → Dashboard → Keys)", ) ei_allow_duplicates = gr.Checkbox(label="Allow duplicate samples", value=False) generate_btn = gr.Button("Generate dataset", variant="primary") summary_md = gr.Markdown() download = gr.File(label="Download dataset (zip)") logs_box = gr.Textbox(label="Logs", lines=16, max_lines=30) generate_btn.click( fn=create_dataset, inputs=[ dataset_name, wake_label, wake_phrases_text, unknown_phrases_text, gcp_api_key, base_repeats, augmentations, background_noise, max_voices, test_ratio, hf_repo_id, hf_token, hf_private, do_push_hf, ei_api_key, do_upload_ei, ei_allow_duplicates, ], outputs=[logs_box, download, summary_md], ) if __name__ == "__main__": demo.queue().launch()