""" Gradio app wrapping your diarization + separation + enhancement + transcription pipeline. """ import os import tempfile import math import json import shutil import time from datetime import timedelta from pathlib import Path from typing import List, Tuple import re import numpy as np import soundfile as sf import librosa import noisereduce as nr import gradio as gr # Lazy imports (heavy models) will be done inside the worker function # to keep the app responsive on startup. # ----------------------- # Configuration defaults # ----------------------- SAMPLE_RATE = 16000 CHUNK_DURATION = 8.0 KEYWORDS = ["red", "yellow", "green"] HF_TOKEN_E = os.environ.get("HF_TOKEN") # ----------------------- # Helper utilities # ----------------------- def time_to_samples(t: float, sr: int) -> int: return int(round(t * sr)) def save_wav(path: str, data: np.ndarray, sr: int = SAMPLE_RATE): sf.write(path, data.astype(np.float32), sr) # ----------------------- # Transcription helper # ----------------------- def transcribe_audio_array_with_whisper(audio: np.ndarray, sr: int, whisper_model) -> dict: """Whisper expects a file path; write to temp wav then transcribe.""" tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False) try: sf.write(tmp.name, audio.astype(np.float32), sr) res = whisper_model.transcribe(tmp.name, task="transcribe", fp16=False, language=None) return res except Exception as e: return {"text": "", "segments": []} finally: try: tmp.close() os.unlink(tmp.name) except Exception: pass def transcribe_file_with_whisper(wav_path: str, whisper_model) -> dict: try: res = whisper_model.transcribe(wav_path, task="transcribe", fp16=False, language=None) return res except Exception as e: return {"text": "", "segments": []} # ----------------------- # Keyword finder # ----------------------- def find_keywords_in_text(text: str, keywords: List[str]) -> List[Tuple[str, int]]: found = [] for kw in keywords: for match in re.finditer(rf"\b{re.escape(kw)}\b", text, flags=re.IGNORECASE): found.append((kw, match.start())) return found # ----------------------- # Main pipeline (wrapped for Gradio streaming) # ----------------------- def pipeline_worker(video_file_path: str, keywords: List[str]): """ Generator function that yields progress logs and finally returns (log, file_list, keyword_log, transcripts_json_path) The Gradio interface will call this function and stream the logs. """ # Prepare temporary output directory per-run run_dir = tempfile.mkdtemp(prefix="diarize_run_") out_dir = os.path.join(run_dir, "out") os.makedirs(out_dir, exist_ok=True) logs = [] def emit(message: str): nonlocal logs logs.append(message) yield "\n".join(logs), "", "", "" # 1) Convert mp4 to wav (use moviepy) yield from emit(f"Starting run — saving outputs to: {out_dir}") try: from moviepy.editor import VideoFileClip except Exception as e: yield from emit(f"ERROR: moviepy import failed: {e}") return wav_path = os.path.join(run_dir, "input_audio.wav") try: yield from emit("Extracting audio from video...") clip = VideoFileClip(video_file_path) clip.audio.write_audiofile(wav_path, codec="pcm_s16le") clip.close() yield from emit(f"Saved extracted audio: {wav_path}") except Exception as e: yield from emit(f"ERROR extracting audio: {e}") return # 2) Load audio (librosa) try: y, sr = librosa.load(wav_path, sr=SAMPLE_RATE, mono=True) duration = len(y) / sr yield from emit(f"Loaded audio: {duration:.1f}s @ {sr}Hz") except Exception as e: yield from emit(f"ERROR loading audio: {e}") return # Lazy-load heavy models yield from emit("Loading diarization & embedding models (this can take a while)...") HF_TOKEN = os.environ.get("HF_TOKEN_1") try: from pyannote.audio import Pipeline, Model # diarize_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization@2022.07", use_auth_token=HF_TOKEN) diarize_pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token=HF_TOKEN_E) embedding_model = Model.from_pretrained("pyannote/embedding", use_auth_token=HF_TOKEN_E) yield from emit("pyannote models loaded.") except Exception as e: yield from emit(f"WARNING: pyannote models failed to load: {e}\nDiarization may not work.") diarize_pipeline = None embedding_model = None # Load separation & enhancement (speechbrain) lazily try: from speechbrain.pretrained import SepformerSeparation as Sepformer from speechbrain.pretrained import SpectralMaskEnhancement as Enhancer sepformer = Sepformer.from_hparams(source="speechbrain/sepformer-whamr", savedir=os.path.join(run_dir, "tmp_speechbrain_sepformer")) enhancer = Enhancer.from_hparams(source="speechbrain/metricgan-plus-voicebank", savedir=os.path.join(run_dir, "tmp_speechbrain_enh")) yield from emit("Speechbrain sepformer + enhancer loaded.") except Exception as e: yield from emit(f"WARNING: speechbrain models failed to load: {e}\nSeparation/enhancement fallbacks will be used.") sepformer = None enhancer = None # Load whisper model lazily try: import whisper whisper_model = whisper.load_model("large-v3", device="cpu") yield from emit("Whisper loaded (large-v3) on CPU.") except Exception as e: yield from emit(f"ERROR loading Whisper model: {e}") whisper_model = None # run diarization if diarize_pipeline is None: yield from emit("Skipping diarization (pipeline unavailable). Creating single ""speaker_0"" segment covering full audio.") diarization = None speakers = ["SPEAKER_0"] segments = [ (0.0, duration, "SPEAKER_0") ] else: yield from emit("Running diarization... This may take a while.") try: diarization = diarize_pipeline({"audio": wav_path}) speakers = sorted({label for segment, track, label in diarization.itertracks(yield_label=True)}) yield from emit(f"Detected speakers: {speakers}") except Exception as e: yield from emit(f"ERROR during diarization: {e}") diarization = None speakers = ["SPEAKER_0"] # Prepare speaker buffers speaker_buffers = {sp: [] for sp in speakers} transcriptions = [] # Helper to compute embedding from numpy audio (if model available) def embedding_from_audio(audio_np: np.ndarray): if embedding_model is None: return np.zeros((1, 256)) waveform = audio_np.reshape(1, -1) try: emb = embedding_model({'waveform': waveform, 'sample_rate': SAMPLE_RATE}) return emb.data.numpy().reshape(1, -1) except Exception: return np.zeros((1, 256)) # Iterate through diarized segments (or single fallback) yield from emit("Processing diarized segments (separation/enhancement/transcription)...") if diarization is None: segments_iter = [(0.0, duration, "SPEAKER_0")] else: segments_iter = [(seg.start, seg.end, lbl) for seg, _, lbl in diarization.itertracks(yield_label=True)] for idx, (start, end, label) in enumerate(segments_iter): seg_dur = end - start a_samp = time_to_samples(start, sr) b_samp = time_to_samples(end, sr) seg_audio = y[a_samp:b_samp] yield from emit(f"Segment {idx+1}/{len(segments_iter)}: {label} [{start:.2f}-{end:.2f}] ({seg_dur:.2f}s)") # Detect overlaps (simple check) is_overlap = False if diarization is not None: overlapped_labels = [lbl for s2, _, lbl in diarization.itertracks(yield_label=True) if s2.start < end and s2.end > start and lbl != label] is_overlap = len(overlapped_labels) > 0 # Non-overlap & short => enhance and append if not is_overlap and seg_dur <= CHUNK_DURATION: # attempt enhancer try: if enhancer is not None: import torch wav_tensor = torch.tensor(seg_audio).float().unsqueeze(0) enhanced = enhancer.enhance_batch(wav_tensor).squeeze(0).numpy() else: raise Exception("enhancer unavailable") except Exception: enhanced = nr.reduce_noise(y=seg_audio, sr=sr) speaker_buffers[label].append(enhanced.flatten()) # transcribe if whisper_model is not None: try: res = transcribe_audio_array_with_whisper(enhanced, sr, whisper_model) transcript_text = res.get("text", "").strip() except Exception: transcript_text = "[Transcription failed]" else: transcript_text = "[Whisper unavailable]" transcriptions.append({ "speaker": label, "start": float(start), "end": float(end), "duration": float(seg_dur), "text": transcript_text, }) else: # Overlapped or long: chunk, separate, embed, match to prototypes samples = seg_audio n_chunks = max(1, math.ceil(len(samples) / int(CHUNK_DURATION * sr))) chunk_size = int(len(samples) / n_chunks) for i in range(n_chunks): a = i * chunk_size b = min(len(samples), (i + 1) * chunk_size) chunk = samples[a:b] if len(chunk) < 100: continue # Try sepformer separation est_sources = None try: if sepformer is not None: # speechbrain sepformer has a separate_file_chunkwise or separate_file; attempt both try: est_sources = sepformer.separate_file_chunkwise(batch_audio=chunk, sample_rate=sr) except Exception: tmpf = tempfile.NamedTemporaryFile(suffix=".wav", delete=False) sf.write(tmpf.name, chunk, sr) est = sepformer.separate_file(tmpf.name) tmpf.close() os.unlink(tmpf.name) est_sources = est except Exception: est_sources = None if est_sources is None: # fallback: attempt simple split into two channels (if mono, duplicate) — conservative fallback est_sources = [chunk, chunk] # Compute embeddings embeddings = [] for src in est_sources: try: emb = embedding_from_audio(np.asarray(src).flatten()) except Exception: emb = np.zeros((1, 256)) embeddings.append(emb) # Speaker prototypes speaker_protos = {} for sp in speakers: if len(speaker_buffers[sp]) > 0: ex = np.concatenate([np.asarray(p).flatten() for p in speaker_buffers[sp][:1]]) speaker_protos[sp] = embedding_from_audio(ex) else: speaker_protos[sp] = None for src_idx, emb in enumerate(embeddings): best_sp, best_sim = None, -1 for sp in speakers: proto = speaker_protos[sp] if proto is None: continue try: from sklearn.metrics.pairwise import cosine_similarity sim = cosine_similarity(emb, proto)[0, 0] except Exception: sim = -1 if sim > best_sim: best_sim = sim best_sp = sp assign_to = best_sp if best_sp is not None else speakers[src_idx % len(speakers)] speaker_buffers[assign_to].append(np.asarray(est_sources[src_idx]).flatten()) # Transcribe separated chunk if whisper_model is not None: try: res = transcribe_audio_array_with_whisper(np.asarray(est_sources[src_idx]).flatten(), sr, whisper_model) transcript_text = res.get("text", "").strip() except Exception: transcript_text = "[Transcription failed]" else: transcript_text = "[Whisper unavailable]" transcriptions.append({ "speaker": assign_to, "start": float(start + a / sr), "end": float(start + b / sr), "duration": float((b - a) / sr), "text": transcript_text, }) # Emit progress after each segment yield from emit(f"Processed segment {idx+1}/{len(segments_iter)}") # After processing all segments: write per-speaker concatenated wavs yield from emit("Concatenating speaker buffers and saving speaker wav files...") generated_files = [] for sp, pieces in speaker_buffers.items(): if len(pieces) == 0: continue out = np.concatenate([np.asarray(p).flatten() for p in pieces]) out_path = os.path.join(out_dir, f"{sp}.wav") save_wav(out_path, out, sr) generated_files.append(out_path) yield from emit(f"Saved speaker file: {out_path}") # Build residual noise track (simple reconstruction) yield from emit("Building residual noise track...") recon = np.zeros_like(y) cursor = 0 for sp, pieces in speaker_buffers.items(): if len(pieces) == 0: continue recon_piece = np.concatenate([np.asarray(p).flatten() for p in pieces]) length = min(len(recon_piece), len(recon) - cursor) if length <= 0: continue recon[cursor:cursor+length] += recon_piece[:length] cursor += length residual = y - recon residual_path = os.path.join(out_dir, "noise_residual.wav") save_wav(residual_path, residual, sr) generated_files.append(residual_path) yield from emit(f"Saved residual: {residual_path}") # Save timestamped transcriptions (from the `transcriptions` built earlier) transcript_file = os.path.join(out_dir, "timestamped_transcriptions.json") with open(transcript_file, "w", encoding="utf-8") as f: json.dump(transcriptions, f, indent=2, ensure_ascii=False) generated_files.append(transcript_file) yield from emit(f"Saved timestamped transcriptions: {transcript_file}") # Run a second pass: run whisper on each speaker file for segments (detailed JSON) yield from emit("Running final Whisper pass on each speaker file to produce detailed transcripts...") detailed_paths = [] for sp in speakers: sp_wav_path = os.path.join(out_dir, f"{sp}.wav") if not os.path.exists(sp_wav_path): continue if whisper_model is not None: res = transcribe_file_with_whisper(sp_wav_path, whisper_model) text = res.get("text", "").strip() segments = res.get("segments", []) else: text = "" segments = [] json_path = os.path.join(out_dir, f"{sp}_transcript.json") with open(json_path, "w", encoding="utf-8") as fj: json.dump({"speaker": sp, "text": text, "segments": segments}, fj, indent=2, ensure_ascii=False) detailed_paths.append(json_path) generated_files.append(json_path) yield from emit(f"Saved detailed JSON: {json_path}") # Keyword scanning yield from emit("Scanning transcripts for keywords...") keyword_log_lines = [] for sp in speakers: json_path = os.path.join(out_dir, f"{sp}_transcript.json") if not os.path.exists(json_path): continue with open(json_path, "r", encoding="utf-8") as f: data = json.load(f) text = data.get("text", "") segments = data.get("segments", []) if segments: for seg in segments: seg_text = seg.get("text", "") seg_start = seg.get("start", 0) seg_end = seg.get("end", 0) hits = find_keywords_in_text(seg_text, keywords) if hits: s_td = str(timedelta(seconds=float(seg_start))) e_td = str(timedelta(seconds=float(seg_end))) line = f"Speaker: {sp} [{s_td} --> {e_td}] Text: {seg_text.strip()}" keyword_log_lines.append(line) else: hits = find_keywords_in_text(text, keywords) if hits: line = f"Speaker: {sp} [No segment timestamps available] Excerpt: {text.strip()[:200]}" keyword_log_lines.append(line) if len(keyword_log_lines) == 0: keyword_log = "No keyword matches found." else: keyword_log = "\n".join(keyword_log_lines) yield from emit("Keyword scan complete.") # Final return: logs, list of generated files (as newline list), keywords, path to timestamped JSON file_list_text = "\n".join(generated_files) yield "\n".join(logs), file_list_text, keyword_log, transcript_file # # ----------------------- # # Gradio UI # # ----------------------- # def build_interface(): # with gr.Blocks() as demo: # gr.Markdown("# Voice Analysis (Diarisation and Signal Identification)\nUpload an MP4 and click Run to start analysis.") # with gr.Row(): # video_in = gr.Video(label="Input video (.mp4)") # keywords_in = gr.Textbox(value=",".join(KEYWORDS), label="Keywords (comma separated)") # run_btn = gr.Button("Run") # with gr.Row(): # # logs_out = gr.Textbox(label="Progress logs", lines=20) # # files_out = gr.Textbox(label="Generated files (saved in temp run folder)", lines=20) # keywords_out = gr.Textbox(label="Keyword matches (console-style)", lines=5) # transcript_json_out = gr.Textbox(label="Timestamped transcript JSON path") # # Loading indicator (spinner) # with gr.Row(): # status_msg = gr.Markdown("⏳ *Idle...*") # # Add a JSON viewer for transcript preview # with gr.Accordion("📜 View Detailed Transcript JSON", open=False): # transcript_view = gr.JSON(label="Transcript Data (Timestamps + Text)") # # Function to open and display transcript JSON file # def open_transcript_json(json_path): # if not os.path.exists(json_path): # return {"error": "File not found"} # try: # with open(json_path, "r", encoding="utf-8") as f: # data = json.load(f) # return data # except Exception as e: # return {"error": str(e)} # # Button to view JSON file content # view_btn = gr.Button("Open Transcript JSON") # view_btn.click(fn=open_transcript_json, inputs=transcript_json_out, outputs=transcript_view) # def run_and_stream(video_path, keywords_text, progress=gr.Progress(track_tqdm=True)): # progress(0, desc="Starting analysis...") # keys = [k.strip() for k in keywords_text.split(",") if k.strip()] # gen = pipeline_worker(video_path, keys) # for out in gen: # yield out # # Update status to "Processing..." # yield "Processing...", "", "⏳ **Processing... Please wait.**" # for out in pipeline_worker(video_path, keys): # progress(0.5, desc="Running pipeline...") # yield out, "", "⚙️ **Working...**" # # Done # progress(1, desc="Completed!") # yield "Processing done", "Processing complete", "✅ **Processing done!**" # # ----------------------- # # Attach button to function # # ----------------------- # run_btn.click( # fn=run_and_stream, # inputs=[video_in, keywords_in], # outputs=[keywords_out, transcript_json_out, status_msg] # ) # # def run_and_stream(video_path, keywords_text): # # keys = [k.strip() for k in keywords_text.split(",") if k.strip()] # # gen = pipeline_worker(video_path, keys) # # for out in gen: # # yield out # # yield "Processing done", "Output is ready" # # # run_btn.click(fn=run_and_stream, inputs=[video_in, keywords_in], outputs=[logs_out, files_out, keywords_out, transcript_json_out]) # # run_btn.click(fn=run_and_stream, inputs=[video_in, keywords_in], outputs=[keywords_out, transcript_json_out]) # return demo # ----------------------- # Gradio UI # ----------------------- def build_interface(): with gr.Blocks() as demo: gr.Markdown("# Voice Analysis (Diarisation and Signal Identification)\nUpload an MP4 and click Run to start analysis.") with gr.Row(): video_in = gr.Video(label="Input video (.mp4)") keywords_in = gr.Textbox(value=",".join(KEYWORDS), label="Keywords (comma separated)") run_btn = gr.Button("Run") with gr.Row(): logs_out = gr.Textbox(label="Progress logs", lines=20) files_out = gr.Textbox(label="Generated files (saved in temp run folder)", lines=20) with gr.Row(): keywords_out = gr.Textbox(label="Keyword matches (console-style)", lines=5) transcript_json_out = gr.Textbox(label="Timestamped transcript JSON path") # Add a JSON viewer for transcript preview with gr.Accordion("📜 View Detailed Transcript JSON", open=False): transcript_view = gr.JSON(label="Transcript Data (Timestamps + Text)") # Function to open and display transcript JSON file def open_transcript_json(json_path): if not os.path.exists(json_path): return {"error": "File not found"} try: with open(json_path, "r", encoding="utf-8") as f: data = json.load(f) return data except Exception as e: return {"error": str(e)} # Button to view JSON file content view_btn = gr.Button("Open Transcript JSON") view_btn.click(fn=open_transcript_json, inputs=transcript_json_out, outputs=transcript_view) def run_and_stream(video_path, keywords_text): keys = [k.strip() for k in keywords_text.split(",") if k.strip()] gen = pipeline_worker(video_path, keys) for out in gen: yield out run_btn.click(fn=run_and_stream, inputs=[video_in, keywords_in], outputs=[logs_out, files_out, keywords_out, transcript_json_out]) # run_btn.click(fn=run_and_stream, inputs=[video_in, keywords_in], outputs=[keywords_out, transcript_json_out]) return demo app = build_interface() if __name__ == "__main__": app.launch(server_name="0.0.0.0", server_port=7860)