Spaces:
Paused
Paused
| #!/usr/bin/env python3 | |
| """ | |
| prepare_dataset.py | |
| Build a clean, single-speaker voice dataset from a folder of video files. | |
| Pipeline (per video): | |
| 1. ffmpeg -> extract the audio track | |
| 2. Demucs -> strip music / background noise, keep vocals only | |
| 3. Whisper -> transcribe with word timestamps (faster-whisper), cut into sentences | |
| 4. Speaker ID -> check every sentence against the reference voice (ECAPA embeddings) | |
| 5. Chunking -> merge consecutive target-speaker sentences into 3-15 s clips, | |
| then re-check each finished clip | |
| 6. Quality -> reject clipped, too-quiet, or near-empty clips | |
| 7. Export -> wavs/ + metadata.csv (LJSpeech style) + report.csv (every candidate, kept or not) | |
| Resumable: videos already processed are skipped on re-run. | |
| """ | |
| import os | |
| os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1") # quiet a harmless Windows warning | |
| os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS", "1") | |
| import argparse | |
| import csv | |
| import hashlib | |
| import json | |
| import re | |
| import shutil | |
| import subprocess | |
| import sys | |
| import tempfile | |
| from pathlib import Path | |
| import numpy as np | |
| import soundfile as sf | |
| import torch | |
| import torchaudio | |
| from tqdm import tqdm | |
| VIDEO_EXTS = {".mp4", ".mov", ".mkv", ".webm", ".avi", ".m4v", ".mts", ".wav", ".mp3", ".m4a"} | |
| EMBED_SR = 16000 # speaker model + whisper both work at 16 kHz | |
| # ---------------------------------------------------------------- utilities | |
| def run(cmd): | |
| """Run a command; raise with its stderr if it fails.""" | |
| result = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True) | |
| if result.returncode != 0: | |
| raise RuntimeError(f"Command failed: {' '.join(map(str, cmd))}\n{result.stderr[-2000:]}") | |
| def check_ffmpeg(): | |
| if shutil.which("ffmpeg") is None: | |
| raise ValueError("ffmpeg not found. Install it first (see README).") | |
| def decode_audio(src, dst, sr, channels): | |
| run(["ffmpeg", "-y", "-loglevel", "error", "-i", str(src), "-vn", | |
| "-ac", str(channels), "-ar", str(sr), "-c:a", "pcm_s16le", str(dst)]) | |
| def load_mono(path): | |
| data, sr = sf.read(str(path), dtype="float32", always_2d=True) | |
| return data.mean(axis=1), sr | |
| def resample(x, sr_in, sr_out): | |
| if sr_in == sr_out: | |
| return x | |
| return torchaudio.functional.resample(torch.from_numpy(x), sr_in, sr_out).numpy() | |
| def video_id(path): | |
| slug = re.sub(r"[^A-Za-z0-9]+", "_", path.stem).strip("_")[:40] or "video" | |
| digest = hashlib.md5(str(path.resolve()).encode()).hexdigest()[:6] | |
| return f"{slug}_{digest}" | |
| def dbfs(x): | |
| rms = np.sqrt(np.mean(np.square(x)) + 1e-12) | |
| return 20 * np.log10(rms + 1e-12) | |
| # ---------------------------------------------------------------- vocal separation | |
| def separate_vocals(wav_path, workdir, device): | |
| run([sys.executable, "-m", "demucs", "--two-stems", "vocals", "-n", "htdemucs", | |
| "-d", device, "-o", str(workdir), str(wav_path)]) | |
| vocals = workdir / "htdemucs" / wav_path.stem / "vocals.wav" | |
| if not vocals.exists(): | |
| raise RuntimeError(f"Demucs did not produce {vocals}") | |
| return vocals | |
| # ---------------------------------------------------------------- speaker filter | |
| class SpeakerFilter: | |
| """Scores how closely a clip matches the reference speaker (cosine similarity, -1..1).""" | |
| def __init__(self, reference_paths, device, tmpdir): | |
| from importlib import import_module | |
| try: | |
| EncoderClassifier = import_module("speechbrain.inference.speaker").EncoderClassifier | |
| except ImportError: # speechbrain < 1.0 | |
| EncoderClassifier = import_module("speechbrain.pretrained").EncoderClassifier | |
| self.device = device | |
| kwargs = dict( | |
| source="speechbrain/spkrec-ecapa-voxceleb", | |
| savedir=str(Path.home() / ".cache" / "voiceprep" / "ecapa"), | |
| run_opts={"device": device}, | |
| ) | |
| # Windows blocks symlinks for normal users, so copy the model files instead. | |
| try: | |
| from speechbrain.utils.fetching import LocalStrategy | |
| kwargs["local_strategy"] = LocalStrategy.COPY | |
| except ImportError: # older speechbrain without LocalStrategy | |
| pass | |
| self.model = EncoderClassifier.from_hparams(**kwargs) | |
| embs = [] | |
| for i, ref in enumerate(reference_paths): | |
| tmp = Path(tmpdir) / f"ref_{i}.wav" | |
| decode_audio(ref, tmp, EMBED_SR, 1) | |
| x, _ = load_mono(tmp) | |
| embs.append(self.embed(x)) | |
| ref = torch.stack(embs).mean(dim=0) | |
| self.ref = ref / ref.norm() | |
| def embed(self, x16): | |
| wav = torch.from_numpy(x16).float().unsqueeze(0).to(self.device) | |
| emb = self.model.encode_batch(wav).squeeze().cpu() | |
| return emb / emb.norm() | |
| def score(self, x16): | |
| return float(torch.dot(self.embed(x16), self.ref)) | |
| # ---------------------------------------------------------------- chunking | |
| class Unit: | |
| """A sentence-sized piece of speech with timing, text, and a speaker verdict.""" | |
| def __init__(self, start, end, text): | |
| self.start, self.end, self.text = start, end, text | |
| self.ok = True | |
| self.sim = None | |
| def split_sentences(segments, pause=0.5): | |
| """Re-cut Whisper segments into sentence-sized units using word timestamps. | |
| A new unit starts after sentence-ending punctuation or a pause, so a quick | |
| back-and-forth between two people ends up in separate units.""" | |
| units = [] | |
| for seg in segments: | |
| words = [w for w in (seg.words or []) if w.word.strip()] | |
| if not words: | |
| if seg.text.strip(): | |
| units.append(Unit(seg.start, seg.end, seg.text.strip())) | |
| continue | |
| cur = None | |
| for w in words: | |
| if cur and (w.start - cur.end > pause): | |
| units.append(cur) | |
| cur = None | |
| if cur is None: | |
| cur = Unit(w.start, w.end, w.word.strip()) | |
| else: | |
| cur.end = w.end | |
| cur.text += w.word if w.word.startswith(" ") else " " + w.word | |
| if w.word.strip()[-1:] in ".?!": | |
| units.append(cur) | |
| cur = None | |
| if cur: | |
| units.append(cur) | |
| for u in units: | |
| u.text = re.sub(r"\s+", " ", u.text).strip() | |
| return units | |
| def build_chunks(segments, min_len, max_len, max_gap): | |
| """Merge consecutive speech units into clips between min_len and max_len seconds. | |
| Units are never split, so text stays aligned with audio. A unit marked as | |
| another speaker (ok=False) always ends the current clip, so voices never mix.""" | |
| chunks, cur = [], None | |
| for seg in segments: | |
| text = seg.text.strip() | |
| if not text: | |
| continue | |
| if not getattr(seg, "ok", True): # someone else talking: never merge across it | |
| if cur: | |
| chunks.append(cur) | |
| cur = None | |
| continue | |
| if seg.end - seg.start > max_len: # a single overlong segment can't be aligned safely | |
| if cur: | |
| chunks.append(cur) | |
| cur = None | |
| continue | |
| if cur and (seg.start - cur["end"] <= max_gap) and (seg.end - cur["start"] <= max_len): | |
| cur["end"] = seg.end | |
| cur["text"] += " " + text | |
| else: | |
| if cur: | |
| chunks.append(cur) | |
| cur = {"start": seg.start, "end": seg.end, "text": text} | |
| if cur: | |
| chunks.append(cur) | |
| return [c for c in chunks if c["end"] - c["start"] >= min_len] | |
| # ---------------------------------------------------------------- main per-video work | |
| def process_video(path, args, whisper, spk, writers, tmp_root, report=None): | |
| vid = video_id(path) | |
| work = Path(tempfile.mkdtemp(prefix=vid + "_", dir=tmp_root)) | |
| say = report or (lambda msg: tqdm.write(f" [{path.name}] {msg}")) | |
| try: | |
| say("extracting audio...") | |
| full = work / f"{vid}.wav" | |
| decode_audio(path, full, 44100, 2) | |
| say("removing music and background noise (slowest step)...") | |
| vocals_path = separate_vocals(full, work, args.device) | |
| vocals, sr = load_mono(vocals_path) | |
| vocals16 = resample(vocals, sr, EMBED_SR) | |
| out_audio = resample(vocals, sr, args.sample_rate) | |
| total = len(vocals16) / EMBED_SR | |
| say("transcribing...") | |
| segments, _ = whisper.transcribe( | |
| vocals16, language=args.language, vad_filter=True, beam_size=5, | |
| condition_on_previous_text=False, word_timestamps=True, | |
| ) | |
| units = split_sentences(list(segments)) | |
| # Check the speaker on every sentence BEFORE merging, so mixed clips can't form. | |
| say(f"checking who is speaking in {len(units)} sentences...") | |
| other = 0 | |
| for j, u in enumerate(units): | |
| dur = u.end - u.start | |
| if dur < args.min_unit: # too short to judge reliably ("Yep.") -> treat as a break | |
| u.ok = False | |
| continue | |
| s, e = u.start, u.end | |
| if dur < 1.5: # give the voice model at least 1.5 s of context | |
| mid = (s + e) / 2 | |
| s, e = max(0.0, mid - 0.75), min(total, mid + 0.75) | |
| u.sim = spk.score(vocals16[int(s * EMBED_SR):int(e * EMBED_SR)]) | |
| u.ok = u.sim >= args.sentence_threshold | |
| if not u.ok: | |
| other += 1 | |
| writers["report"].writerow([f"{vid}_s{j:04d}", path.name, f"{u.start:.2f}", | |
| f"{u.end:.2f}", f"{dur:.2f}", f"{u.sim:.3f}", "", | |
| "rejected", "other_speaker_sentence", u.text]) | |
| say(f"{other} sentences from other speakers removed, building clips...") | |
| chunks = build_chunks(units, args.min_len, args.max_len, args.max_gap) | |
| kept = 0 | |
| pad = args.pad | |
| for i, c in enumerate(chunks): | |
| start = max(0.0, c["start"] - pad) | |
| end = min(len(vocals16) / EMBED_SR, c["end"] + pad) | |
| text = re.sub(r"\s+", " ", c["text"]).strip() | |
| clip_id = f"{vid}_{i:04d}" | |
| seg16 = vocals16[int(start * EMBED_SR):int(end * EMBED_SR)] | |
| seg_out = out_audio[int(start * args.sample_rate):int(end * args.sample_rate)] | |
| sim = spk.score(seg16) | |
| level = dbfs(seg_out) | |
| clipped = float(np.mean(np.abs(seg_out) >= 0.999)) | |
| reason = "" | |
| if len(text.split()) < args.min_words: | |
| reason = "too_few_words" | |
| elif sim < args.threshold: | |
| reason = "other_speaker" | |
| elif level < args.min_dbfs: | |
| reason = "too_quiet" | |
| elif clipped > 0.001: | |
| reason = "clipped" | |
| status = "rejected" if reason else "kept" | |
| if not reason: | |
| peak = np.max(np.abs(seg_out)) + 1e-9 | |
| seg_out = seg_out * (0.89 / peak) # peak-normalize to about -1 dBFS | |
| sf.write(str(args.out / "wavs" / f"{clip_id}.wav"), seg_out, args.sample_rate, | |
| subtype="PCM_16") | |
| writers["meta"].writerow([clip_id, text, text]) | |
| kept += 1 | |
| writers["report"].writerow([clip_id, path.name, f"{start:.2f}", f"{end:.2f}", | |
| f"{end - start:.2f}", f"{sim:.3f}", f"{level:.1f}", | |
| status, reason, text]) | |
| return kept, len(chunks) | |
| finally: | |
| if args.keep_temp: | |
| print(f" temp files kept in {work}") | |
| else: | |
| shutil.rmtree(work, ignore_errors=True) | |
| def open_csv(path, header, delimiter=","): | |
| new = not path.exists() | |
| f = open(path, "a", newline="", encoding="utf-8") | |
| w = csv.writer(f, delimiter=delimiter) | |
| if new and header: | |
| w.writerow(header) | |
| return f, w | |
| def run_dataset(args, progress=None): | |
| report_progress = progress or (lambda message, value: None) | |
| check_ffmpeg() | |
| videos = sorted(f for f in args.input.rglob("*") if f.suffix.lower() in VIDEO_EXTS) | |
| if not videos: | |
| raise ValueError(f"No video/audio files found in {args.input}") | |
| (args.out / "wavs").mkdir(parents=True, exist_ok=True) | |
| progress_file = args.out / "progress.json" | |
| done = set(json.loads(progress_file.read_text())) if progress_file.exists() else set() | |
| todo = [v for v in videos if str(v.resolve()) not in done] | |
| print(f"{len(videos)} files found, {len(todo)} left to process, device={args.device}") | |
| if args.device == "cpu": | |
| print("Warning: no GPU detected. This will be very slow on a large archive.") | |
| if not todo: | |
| result = { | |
| "videos_found": len(videos), "videos_processed": 0, | |
| "videos_skipped": len(videos), "clips_kept": 0, | |
| "failures": [], "dataset": str(args.out.resolve()), | |
| } | |
| report_progress("All uploaded videos have already been processed.", 100) | |
| return result | |
| tmp_root = Path(tempfile.mkdtemp(prefix="voiceprep_")) | |
| meta_f = rep_f = None | |
| try: | |
| report_progress("Loading transcription and speaker-recognition models…", 0) | |
| from faster_whisper import WhisperModel | |
| whisper = WhisperModel(args.whisper_model, device=args.device, | |
| compute_type="float16" if args.device == "cuda" else "int8") | |
| spk = SpeakerFilter(args.reference, args.device, tmp_root) | |
| meta_f, meta_w = open_csv(args.out / "metadata.csv", None, delimiter="|") | |
| rep_f, rep_w = open_csv( | |
| args.out / "report.csv", | |
| ["clip_id", "source", "start", "end", "duration", "similarity", | |
| "dbfs", "status", "reason", "text"], | |
| ) | |
| writers = {"meta": meta_w, "report": rep_w} | |
| total_kept = 0 | |
| failures = [] | |
| for index, v in enumerate(todo, 1): | |
| base_progress = (index - 1) / max(len(todo), 1) * 100 | |
| def report_video(message, current=index, current_path=v, base=base_progress): | |
| report_progress( | |
| f"{current_path.name}: {message}", | |
| min(99, round(base + 5 / max(len(todo), 1))), | |
| ) | |
| report_progress(f"Processing video {index} of {len(todo)}: {v.name}", round(base_progress)) | |
| try: | |
| kept, cands = process_video(v, args, whisper, spk, writers, tmp_root, | |
| report=report_video) | |
| total_kept += kept | |
| tqdm.write(f"{v.name}: kept {kept}/{cands} clips") | |
| except Exception as e: # one bad file shouldn't stop the whole run | |
| tqdm.write(f"{v.name}: FAILED - {e}") | |
| failures.append({"video": v.name, "error": str(e)}) | |
| report_progress(f"{v.name} failed: {e}", round(index / max(len(todo), 1) * 100)) | |
| continue | |
| meta_f.flush() | |
| rep_f.flush() | |
| done.add(str(v.resolve())) | |
| progress_file.write_text(json.dumps(sorted(done))) | |
| result = { | |
| "videos_found": len(videos), | |
| "videos_processed": len(todo) - len(failures), | |
| "videos_skipped": len(videos) - len(todo), | |
| "clips_kept": total_kept, | |
| "failures": failures, | |
| "dataset": str(args.out.resolve()), | |
| } | |
| finally: | |
| if meta_f is not None: | |
| meta_f.close() | |
| if rep_f is not None: | |
| rep_f.close() | |
| shutil.rmtree(tmp_root, ignore_errors=True) | |
| report_progress( | |
| f"Dataset ready: {total_kept} new clips from {result['videos_processed']} videos.", | |
| 100, | |
| ) | |
| print(f"\nDone. {total_kept} new clips this run. Dataset: {args.out.resolve()}") | |
| print("Review report.csv to tune --threshold before training.") | |
| return result | |
| def main(): | |
| p = argparse.ArgumentParser(description="Build a single-speaker voice dataset from videos.") | |
| p.add_argument("input", type=Path, help="Folder of videos (searched recursively)") | |
| p.add_argument("--reference", type=Path, nargs="+", required=True, | |
| help="1-5 short clips (10-30 s each) of ONLY the target voice") | |
| p.add_argument("--out", type=Path, default=Path("dataset")) | |
| p.add_argument("--threshold", type=float, default=0.60, | |
| help="Speaker similarity cutoff for finished clips") | |
| p.add_argument("--sentence-threshold", type=float, default=0.55, | |
| help="Speaker cutoff for each sentence before merging (raise if other " | |
| "voices slip into clips)") | |
| p.add_argument("--min-unit", type=float, default=0.6, | |
| help="Sentences shorter than this (s) are too short to judge and split clips") | |
| p.add_argument("--sample-rate", type=int, default=44100, help="Output clip sample rate") | |
| p.add_argument("--min-len", type=float, default=3.0) | |
| p.add_argument("--max-len", type=float, default=15.0) | |
| p.add_argument("--max-gap", type=float, default=0.6, help="Max pause (s) when merging segments") | |
| p.add_argument("--pad", type=float, default=0.10, help="Silence padding (s) around each clip") | |
| p.add_argument("--min-words", type=int, default=2) | |
| p.add_argument("--min-dbfs", type=float, default=-40.0) | |
| p.add_argument("--whisper-model", default="large-v3", | |
| help="large-v3 (best), medium, small (faster on weak GPUs)") | |
| p.add_argument("--language", default="en") | |
| p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") | |
| p.add_argument("--keep-temp", action="store_true", help="Keep extracted/separated audio") | |
| args = p.parse_args() | |
| try: | |
| run_dataset(args) | |
| except ValueError as e: | |
| sys.exit(str(e)) | |
| if __name__ == "__main__": | |
| main() | |