Merge branch 'main' into pr/2

app.py

@@ -50,7 +50,8 @@ def process_with_live_transcript(
     sample_stem,
     stt_model,
     last_sample_stem,
-):
+    current_sample_rate,
+):  
     """Generator that runs the offline pipeline in real time (chunked): enhanced audio and
     both transcripts stream from the first chunk so playback and transcription start immediately."""
     progress_state = {}
@@ -60,10 +61,11 @@ def process_with_live_transcript(
         try:
             result_holder["result"] = run_offline_pipeline_streaming(
                 input_array,
+                current_sample_rate,
                 enhancement_level,
                 sample_stem,
                 stt_model,
-                progress_state,
+                progress_state
             )
         except Exception as e:
             result_holder["error"] = e
@@ -153,6 +155,7 @@ with gr.Blocks() as demo:
     input_array = gr.State()
     enhanced_array = gr.State()
     precomputed_noisy_transcript = gr.State("")
+    current_sample_rate = gr.State(16000)  # default sample rate for dataset samples; updated on local file load if different
 
     gr.HTML(
         '<a href="https://ai-coustics.com/" target="_blank">'
@@ -232,7 +235,7 @@ with gr.Blocks() as demo:
                 choices=ALL_FILES, value="en_00412_i_h_36", label="Sample"
             )
             audio_file_from_dataset = gr.Audio(
-                type="filepath", interactive=False, buttons=["download"], autoplay=True
+                type="filepath", interactive=False, buttons=["download"], autoplay=False
             )
 
         with gr.Tab("Upload local file") as upload_tab:
@@ -257,7 +260,7 @@ with gr.Blocks() as demo:
             type="numpy",
             interactive=False,
             buttons=["download"],
-            autoplay=True,
+            autoplay=False,
         )
 
         with gr.Row(equal_height=True, elem_classes="results-row"):
@@ -288,8 +291,8 @@ with gr.Blocks() as demo:
     def load_dataset_sample_on_tab_visit(dropdown_value):
         """Load the selected sample when visiting the Dataset tab; use default if dropdown is empty."""
         sample_id = dropdown_value or DEFAULT_DATASET_SAMPLE
-        audio_path, arr, stem = load_file_from_dataset(sample_id)
-        return sample_id, audio_path, arr, stem
+        audio_path, arr, stem, sample_rate = load_file_from_dataset(sample_id)
+        return sample_id, audio_path, arr, stem, sample_rate
 
     stream_tab.select(
         lambda: (
@@ -332,7 +335,7 @@ with gr.Blocks() as demo:
     ).then(
         load_dataset_sample_on_tab_visit,
         inputs=[dataset_dropdown],
-        outputs=[dataset_dropdown, audio_file_from_dataset, input_array, sample_stem],
+        outputs=[dataset_dropdown, audio_file_from_dataset, input_array, sample_stem, current_sample_rate],
     )
 
     stt_model.change(
@@ -386,7 +389,7 @@ with gr.Blocks() as demo:
     ).then(
         load_file_from_dataset,
         inputs=dataset_dropdown,
-        outputs=[audio_file_from_dataset, input_array, sample_stem],
+        outputs=[audio_file_from_dataset, input_array, sample_stem, current_sample_rate],
     )
 
     # Uploading a local file triggers loading the audio file and hiding results until enhancement
@@ -395,7 +398,7 @@ with gr.Blocks() as demo:
         outputs=results_card,
     ).then(
         load_local_file,
-        inputs=[audio_file_upload, normalize],
+        inputs=[audio_file_upload, normalize, current_sample_rate],
         outputs=[input_array, sample_stem, audio_preview],
     )
  
@@ -408,7 +411,7 @@ with gr.Blocks() as demo:
     # Enhancement button: run pipeline with live transcript progress (dataset + local file modes).
     enhance_btn.click(
         process_with_live_transcript,
-        inputs=[input_array, enhancement_level, sample_stem, stt_model, last_sample_stem],
+        inputs=[input_array, enhancement_level, sample_stem, stt_model, last_sample_stem, current_sample_rate],
         outputs=[
             results_card,
             enhanced_audio,

offline_pipeline.py

@@ -1,210 +1,15 @@
 import os
-from concurrent.futures import ThreadPoolExecutor
-from typing import Optional
+from random import sample
 
 import gradio as gr
-import librosa
+import soundfile as sf
 from sdk import SDKWrapper
-from utils import spec_image, compute_wer, transcribe_audio, to_gradio_audio, normalize_lufs
+from utils import spec_image, compute_wer, to_gradio_audio, normalize_lufs
 from hf_dataset_utils import get_audio, get_transcript
-from constants import DEFAULT_SR, APP_TMP_DIR, STREAMER_CLASSES
+from constants import APP_TMP_DIR, STREAMER_CLASSES
 import numpy as np
 
 
-def retrieve_audio_information(
-    original_array: np.ndarray,
-    enhanced_array: np.ndarray,
-    sample_id: str,
-    stt_model: str,
-    noisy_transcript: str,
-    progress_state: Optional[dict] = None,
-) -> tuple[str, str, str, str, str, str]:
-    """Build spectrograms, transcribe enhanced audio, and compute WER. Caller must supply
-    noisy_transcript (transcription of original_array) so STT on the original is never run here.
-    If progress_state is provided, progress_state['enhanced'] is updated with partial transcript as enhanced STT streams."""
-    if original_array is None or enhanced_array is None:
-        raise ValueError("Audio arrays are not available.")
-    noisy_spec_path = f"{APP_TMP_DIR}/{sample_id}_noisy_spectrogram.png"
-    enhanced_spec_path = f"{APP_TMP_DIR}/{sample_id}_enhanced_spectrogram.png"
-    spec_image(original_array).save(noisy_spec_path)
-    spec_image(enhanced_array).save(enhanced_spec_path)
-    on_enhanced = (lambda t: progress_state.__setitem__("enhanced", t)) if progress_state is not None else None
-    enhanced_transcript = transcribe_audio(
-        enhanced_array,
-        DEFAULT_SR,
-        stt_model,
-        stream_name=f"{sample_id}_enhanced",
-        on_update=on_enhanced,
-    )
-    try:
-        original_transcript = get_transcript(sample_id)
-        wer_enhanced = compute_wer(original_transcript, enhanced_transcript)
-        wer_noisy = compute_wer(original_transcript, noisy_transcript)
-        enhanced_transcript += f" (WER: {wer_enhanced * 100:.2f}%)"
-        noisy_transcript += f" (WER: {wer_noisy * 100:.2f}%)"
-    except Exception:
-        original_transcript = "Unavailable"
-    return enhanced_spec_path, noisy_spec_path, enhanced_transcript, original_transcript, noisy_transcript, sample_id
-
-
-def denoise_and_transcribe_noisy_parallel(
-    sample_16k: np.ndarray,
-    enhancement_level: float,
-    sample_id: str,
-    stt_model: str,
-    progress_state: Optional[dict] = None,
-) -> tuple[np.ndarray | None, tuple[int, np.ndarray] | None, str]:
-    """Run denoising and noisy transcription in parallel. Returns (enhanced_array, gradio_audio, noisy_transcript).
-    If progress_state is provided, progress_state['noisy'] is updated with partial transcript as noisy STT streams."""
-    if sample_16k is None:
-        raise ValueError("No audio to enhance. Please upload a file first.")
-    on_noisy = (lambda t: progress_state.__setitem__("noisy", t)) if progress_state is not None else None
-    with ThreadPoolExecutor(max_workers=2) as executor:
-        future_denoise = executor.submit(
-            _denoise_audio_impl, sample_16k, float(enhancement_level) / 100.0
-        )
-        future_noisy = executor.submit(
-            transcribe_audio,
-            sample_16k,
-            DEFAULT_SR,
-            stt_model,
-            stream_name=f"{sample_id}_noisy",
-            on_update=on_noisy,
-        )
-        enhanced_array, gradio_enhanced_audio = future_denoise.result()
-        noisy_transcript = future_noisy.result()
-    return enhanced_array, gradio_enhanced_audio, noisy_transcript
-
-
-def _denoise_audio_impl(
-    sample_16k: np.ndarray, enhancement_level_frac: float
-) -> tuple[np.ndarray, tuple[int, np.ndarray]]:
-    """Run the enhancer on a single array and return both the array and Gradio-ready audio.
-
-    Used by denoise_and_transcribe_noisy_parallel so the denoise step can run in a
-    thread without Gradio UI calls. Expects enhancement_level_frac in [0, 1] (not percent).
-    Returns (enhanced_array, (sample_rate, numpy_array)) for Gradio Audio.
-    """
-    sdk = SDKWrapper()
-    sdk.init_processor(
-        sample_rate=DEFAULT_SR,
-        enhancement_level=enhancement_level_frac,
-    )
-    enhanced_array = sdk.process_sync(sample_16k)
-    gradio_enhanced_audio = to_gradio_audio(enhanced_array, DEFAULT_SR)
-    return enhanced_array, gradio_enhanced_audio
-
-
-def denoise_audio(
-    sample_16k: np.ndarray,
-    enhancement_level: float = 50.0,
-) -> tuple[np.ndarray | None , tuple[int, np.ndarray]| None]:
-    """Enhance-only entry point: run the SDK on the given audio and return enhanced array + Gradio audio.
-
-    The main app uses denoise_and_transcribe_noisy_parallel (denoise + noisy STT in parallel)
-    instead. This function remains for backward compatibility, scripts, or any caller that
-    needs only enhancement without transcription (e.g. backup flows, tests).
-    """
-    if sample_16k is None:
-        raise ValueError("No audio to enhance. Please upload a file first.")
-    try:
-        sdk = SDKWrapper()
-        sdk.init_processor(sample_rate=DEFAULT_SR, enhancement_level=float(enhancement_level) / 100.0)
-        enhanced_array = sdk.process_sync(sample_16k)
-    except Exception as e:
-        gr.Warning(f"{e}")
-        raise e
-    gradio_enhanced_audio = to_gradio_audio(enhanced_array, DEFAULT_SR)
-    return enhanced_array, gradio_enhanced_audio
-
-
-def run_offline_pipeline_ordered(
-    sample_16k: np.ndarray,
-    enhancement_level: float,
-    sample_id: str,
-    stt_model: str,
-    progress_state: dict,
-) -> tuple[str, str, str, tuple[int, np.ndarray], str, np.ndarray, str]:
-    """Run pipeline in UI order: 1) Denoise only. 2) When done, set progress_state['enhanced_spec_path']
-    and progress_state['enhanced_audio']. 3) Run noisy STT and enhanced STT in parallel (both stream via
-    progress_state['noisy'] and progress_state['enhanced']). 4) Return final transcripts with WER.
-
-    Returns: (enhanced_spec_path, enhanced_transcript, noisy_transcript_with_wer, enhanced_audio,
-              last_stem, enhanced_array, precomputed_noisy).
-    """
-    if sample_16k is None:
-        raise ValueError("No audio to enhance. Please upload a file first.")
-    # 1) Denoise only
-    enhanced_array, gradio_enhanced_audio = _denoise_audio_impl(
-        sample_16k, float(enhancement_level) / 100.0
-    )
-    # 2) As soon as enhanced audio is ready: build enhanced spectrogram and expose to UI
-    enhanced_spec_path = f"{APP_TMP_DIR}/{sample_id}_enhanced_spectrogram.png"
-    spec_image(enhanced_array).save(enhanced_spec_path)
-    progress_state["enhanced_spec_path"] = enhanced_spec_path
-    progress_state["enhanced_audio"] = gradio_enhanced_audio
-    # 3) Noisy and enhanced transcription starting both at the same time (parallel, both with on_update).
-    # Sync so the UI shows both boxes updating together: only write to noisy/enhanced once both have sent at least one update.
-    progress_state["noisy_pending"] = ""
-    progress_state["enhanced_pending"] = ""
-    progress_state["noisy_has_sent"] = False
-    progress_state["enhanced_has_sent"] = False
-
-    def _flush_both():
-        if progress_state.get("noisy_has_sent") and progress_state.get("enhanced_has_sent"):
-            progress_state["noisy"] = progress_state.get("noisy_pending", "")
-            progress_state["enhanced"] = progress_state.get("enhanced_pending", "")
-
-    def on_noisy(t: str):
-        progress_state["noisy_pending"] = t
-        progress_state["noisy_has_sent"] = True
-        _flush_both()
-
-    def on_enhanced(t: str):
-        progress_state["enhanced_pending"] = t
-        progress_state["enhanced_has_sent"] = True
-        _flush_both()
-
-    with ThreadPoolExecutor(max_workers=2) as executor:
-        future_noisy = executor.submit(
-            transcribe_audio,
-            sample_16k,
-            DEFAULT_SR,
-            stt_model,
-            stream_name=f"{sample_id}_noisy",
-            on_update=on_noisy,
-        )
-        future_enhanced = executor.submit(
-            transcribe_audio,
-            enhanced_array,
-            DEFAULT_SR,
-            stt_model,
-            stream_name=f"{sample_id}_enhanced",
-            on_update=on_enhanced,
-        )
-        noisy_transcript = future_noisy.result()
-        enhanced_transcript = future_enhanced.result()
-    # 4) WER and final strings
-    precomputed_noisy = noisy_transcript
-    try:
-        original_transcript = get_transcript(sample_id)
-        wer_enhanced = compute_wer(original_transcript, enhanced_transcript)
-        wer_noisy = compute_wer(original_transcript, noisy_transcript)
-        enhanced_transcript += f" (WER: {wer_enhanced * 100:.2f}%)"
-        noisy_transcript += f" (WER: {wer_noisy * 100:.2f}%)"
-    except Exception:
-        original_transcript = "Unavailable"
-    return (
-        enhanced_spec_path,
-        enhanced_transcript,
-        noisy_transcript,
-        gradio_enhanced_audio,
-        sample_id,
-        enhanced_array,
-        precomputed_noisy,
-    )
-
-
 def _close_stt_stream(streamer) -> None:
     """Signal end-of-stream; streamer type may be Soniox (close_stream) or Deepgram (close)."""
     if hasattr(streamer, "close_stream"):
@@ -212,9 +17,9 @@ def _close_stt_stream(streamer) -> None:
     else:
         streamer.close()
 
-
 def run_offline_pipeline_streaming(
-    sample_16k: np.ndarray,
+    sample: np.ndarray,
+    sample_rate: int,
     enhancement_level: float,
     sample_id: str,
     stt_model: str,
@@ -224,13 +29,13 @@ def run_offline_pipeline_streaming(
     via progress_state['noisy'] and progress_state['enhanced']. Enhanced audio is returned
     only at the end; the app plays it automatically when processing is complete.
     Returns same tuple as run_offline_pipeline_ordered."""
-    if sample_16k is None:
+    if sample is None:
         raise ValueError("No audio to enhance. Please upload a file first.")
-    sample_16k = np.asarray(sample_16k, dtype=np.float32).flatten()
+    sample = np.asarray(sample, dtype=np.float32).flatten()
 
     sdk = SDKWrapper()
     sdk.init_processor(
-        sample_rate=DEFAULT_SR,
+        sample_rate=sample_rate,
         enhancement_level=float(enhancement_level) / 100.0,
     )
     chunk_size = sdk.num_frames
@@ -259,14 +64,14 @@ def run_offline_pipeline_streaming(
     if stt_model not in STREAMER_CLASSES:
         raise ValueError(f"Unknown STT model: {stt_model}")
     StreamerClass = STREAMER_CLASSES[stt_model]
-    streamer_noisy = StreamerClass(DEFAULT_SR, f"{sample_id}_noisy", on_update=on_noisy)
-    streamer_enhanced = StreamerClass(DEFAULT_SR, f"{sample_id}_enhanced", on_update=on_enhanced)
+    streamer_noisy = StreamerClass(sample_rate, f"{sample_id}_noisy", on_update=on_noisy)
+    streamer_enhanced = StreamerClass(sample_rate, f"{sample_id}_enhanced", on_update=on_enhanced)
 
     accumulated_enhanced: list[np.ndarray] = []
-    n = len(sample_16k)
+    n = len(sample)
 
     for i in range(0, n, chunk_size):
-        raw_chunk = sample_16k[i : i + chunk_size]
+        raw_chunk = sample[i : i + chunk_size]
         if raw_chunk.size < chunk_size:
             raw_chunk = np.pad(
                 raw_chunk,
@@ -292,7 +97,7 @@ def run_offline_pipeline_streaming(
         enhanced_transcript = streamer_enhanced.render_tokens(streamer_enhanced.final_tokens, [])
 
     enhanced_array = np.concatenate(accumulated_enhanced).astype(np.float32)
-    gradio_enhanced_audio = to_gradio_audio(enhanced_array, DEFAULT_SR)
+    gradio_enhanced_audio = to_gradio_audio(enhanced_array, sample_rate)
 
     enhanced_spec_path = f"{APP_TMP_DIR}/{sample_id}_enhanced_spectrogram.png"
     spec_image(enhanced_array).save(enhanced_spec_path)
@@ -318,32 +123,31 @@ def run_offline_pipeline_streaming(
         precomputed_noisy,
     )
 
-
 def load_local_file(
     sample_path: str,
     normalize: bool = True,
-    ) -> tuple[np.ndarray | None, str, tuple | None]:
+    ) -> tuple[np.ndarray | None, str, tuple | None, int]:
     if not sample_path or not os.path.exists(sample_path):
         return None, "", None
     if os.path.getsize(sample_path) > 5 * 1024 * 1024:
         gr.Warning("File size exceeds 5 MB limit. Please upload a smaller file.")
         raise ValueError("Uploaded file exceeds the 5 MB size limit.")
     new_sample_stem = os.path.splitext(os.path.basename(sample_path))[0]
-    y_16k, _ = librosa.load(sample_path, sr=DEFAULT_SR, dtype="float32", mono=True)
+    y, sample_rate = sf.read(sample_path, dtype="float32", always_2d=False)
     if normalize:
-        y_16k = normalize_lufs(y_16k, DEFAULT_SR)
-    gradio_audio = to_gradio_audio(y_16k, DEFAULT_SR)
-    return y_16k, new_sample_stem, gradio_audio
+        y = normalize_lufs(y, sample_rate)
+    gradio_audio = to_gradio_audio(y, sample_rate)
+    return y, new_sample_stem, gradio_audio, sample_rate
 
-def load_file_from_dataset(sample_id: str) -> tuple[tuple | None, np.ndarray | None, str]:
+def load_file_from_dataset(sample_id: str) -> tuple[tuple | None, np.ndarray | None, str, int | None]:
     if not sample_id:
         gr.Warning("Please select a sample from the dropdown.")
-        return None, None, ""
+        return None, None, "", None
     new_sample_stem = sample_id
     try:
-        y_16k, sr = get_audio(sample_id, prefix="mix")
-        y_16k_for_gradio = to_gradio_audio(y_16k, sr)
+        y, sample_rate = get_audio(sample_id, prefix="mix")
+        y_for_gradio = to_gradio_audio(y, sample_rate)
     except Exception as e:  # Convert to 16-bit PCM for gradio audio component
         gr.Warning(f"{e}")
         raise e
-    return y_16k_for_gradio, y_16k, new_sample_stem
+    return y_for_gradio, y, new_sample_stem, sample_rate

stt_streamers/deepgram_streamer.py

@@ -78,12 +78,11 @@ class DeepgramStreamer:
         """
         Returns parameters for the Deepgram V1 URL query string.
         """
-        assert fs_hz == 16000, "Only 16 kHz audio is supported."
 
         return {
             "model": "nova-3",  # Recommended general model
             "encoding": "linear16",  # Corresponds to pcm_s16le
-            "sample_rate": 16000,
+            "sample_rate": fs_hz,
             "channels": 1,
             "smart_format": "true",  # handling punctuation/formatting
             "interim_results": "true",  # required for non-final updates

stt_streamers/soniox_streamer.py

@@ -13,7 +13,6 @@ class SonioxStreamer:
         api_key = os.environ.get("SONIOX_API_KEY")
         if not api_key:
             raise RuntimeError("Missing SONIOX_API_KEY.")
-        assert fs_hz == 16000, "Only 16 kHz audio is supported."
 
         self.stream_name = stream_name
         self.api_name = "Soniox RT"
@@ -69,7 +68,7 @@ class SonioxStreamer:
             "api_key": api_key,
             "model": "stt-rt-v3",
             "audio_format": "pcm_s16le",
-            "sample_rate": 16000,
+            "sample_rate": fs_hz,
             "num_channels": 1,
             "language_hints": ["en", "de"],
             "language_hints_strict": True,

utils.py

@@ -1,11 +1,10 @@
-from typing import Optional, Callable
-import resampy
+from typing import Optional
 import numpy as np
 import librosa
 from PIL import Image
 import io
 import matplotlib.pyplot as plt
-from constants import DEFAULT_SR, TARGET_LOUDNESS, TARGET_TP, STREAMER_CLASSES
+from constants import DEFAULT_SR, TARGET_LOUDNESS, TARGET_TP
 import warnings
 import pyloudnorm as pyln
 
@@ -140,43 +139,3 @@ def normalize_lufs(x: np.ndarray, sr: int) -> np.ndarray:
     except Exception as e:
         warnings.warn(f"LUFS normalization failed, returning input unchanged: {e}")
         return x.astype("float32")
-
-
-
-
-def transcribe_audio(
-    audio_array: np.ndarray,
-    sr: int,
-    streamer_type: str = "Deepgram",
-    stream_name: str = "RAW",
-    on_update: Optional[Callable[[str], None]] = None,
-):
-    """
-    Transcribe an audio array using the specified STT streamer.
-
-    Args:
-        audio_array (np.ndarray): Audio data array
-        sr (int): Sample rate of the audio array
-        streamer_type (str): "Soniox" or "Deepgram"
-        stream_name (str): Optional label for streamer instance ("RAW", "ENHANCED", etc.)
-        on_update: Optional callback(text: str) called with partial transcript as results stream in.
-
-    Returns:
-        str: Final transcript text
-    """
-    if sr != DEFAULT_SR:
-        audio_array = resampy.resample(audio_array, sr, DEFAULT_SR)
-        sr = DEFAULT_SR
-
-    if streamer_type not in STREAMER_CLASSES:
-        raise ValueError(
-            f"Invalid streamer_type '{streamer_type}'. "
-            f"Choose from: {', '.join(STREAMER_CLASSES.keys())}"
-        )
-
-    StreamerClass = STREAMER_CLASSES[streamer_type]
-    streamer = StreamerClass(sr, stream_name, on_update=on_update)
-
-    transcript = streamer.stream_array(audio_array)
-
-    return transcript