vrushil/tiny-smart-turn-test

1import gradio as gr2import numpy as np3import librosa4import onnxruntime as ort5 6# =====================7# MODEL8# =====================9MODEL_PATH = "model.onnx"10 11sess = ort.InferenceSession(12    MODEL_PATH,13    providers=["CPUExecutionProvider"]14)15 16input_name = sess.get_inputs()[0].name17output_name = sess.get_outputs()[0].name18 19# =====================20# CONFIG21# =====================22SR = 1600023N_MELS = 8024N_FRAMES = 800      # Required for Whisper-based models25THRESHOLD = 0.3526 27# =====================28# FEATURE EXTRACTION29# =====================30def extract_features(audio: np.ndarray) -> np.ndarray:31    """32    Whisper-style log-mel spectrogram33    Output shape: (1, 80, 800)34    """35 36    mel = librosa.feature.melspectrogram(37        y=audio,38        sr=SR,39        n_fft=400,40        hop_length=160,41        n_mels=N_MELS,42        fmin=0,43        fmax=8000,44        power=2.0,45        norm="slaney",46        mel_scale="slaney",47    )48 49    mel = np.log(np.maximum(mel, 1e-10))50 51    # Pad or trim to 800 frames52    if mel.shape[1] < N_FRAMES:53        mel = np.pad(54            mel,55            ((0, 0), (0, N_FRAMES - mel.shape[1])),56            mode="constant"57        )58    else:59        mel = mel[:, :N_FRAMES]60 61    return mel[np.newaxis, :, :].astype(np.float32)62 63# =====================64# PREDICTION65# =====================66def predict(audio):67    try:68        if audio is None:69            return "❌ No audio provided"70 71        sr, y = audio72 73        # Stereo → mono74        if y.ndim == 2:75            y = np.mean(y, axis=1)76 77        # Convert to float3278        y = y.astype(np.float32)79 80        # Resample if needed81        if sr != SR:82            y = librosa.resample(y, orig_sr=sr, target_sr=SR)83 84        # Normalize audio85        y = y / (np.max(np.abs(y)) + 1e-9)86 87        # Extract features88        feats = extract_features(y)89 90        # ONNX inference91        output = sess.run([output_name], {input_name: feats})[0]92        output = np.asarray(output)93 94        # Handle different output formats95        if output.ndim == 2 and output.shape[1] == 2:96            score = float(output[0, 1])97        else:98            score = float(output.flatten()[0])99 100        # Optional: sigmoid if model outputs logits101        # score = 1 / (1 + np.exp(-score))102 103        if score >= THRESHOLD:104            return f"✅ COMPLETE (score={score:.3f})"105        else:106            return f"❌ INCOMPLETE (score={score:.3f})"107 108    except Exception as e:109        return f"❌ ERROR: {str(e)}"110 111# =====================112# GRADIO UI113# =====================114demo = gr.Interface(115    fn=predict,116    inputs=gr.Audio(117        sources=["upload", "microphone"],118        type="numpy",119        label="Upload or record audio"120    ),121    outputs=gr.Textbox(label="Prediction"),122    title="Smart Turn Detection (ONNX)",123    description="Upload or record audio to detect turn completion."124)125 126 127 128if __name__ == "__main__":129    demo.launch()130