shift0x/smart_turn_v3

1import io2import math3import os4import logging5from typing import List, Optional6 7import numpy as np8import soundfile as sf9from scipy.signal import resample_poly10from fastapi import FastAPI, UploadFile, File, HTTPException11from fastapi.middleware.cors import CORSMiddleware12from pydantic import BaseModel13 14# Import your provided inference module (must be in the same repo)15# It should contain: predict_endpoint(audio_array: np.ndarray) -> dict16# and load the ONNX model "smart-turn-v3.0.onnx" at import time.17try:18    import inference as inf19except Exception as e:20    raise RuntimeError(21        "Failed to import inference.py. Ensure inference.py and smart-turn-v3.0.onnx "22        "are present in the repository."23    ) from e24 25logger = logging.getLogger("uvicorn")26logger.setLevel(logging.INFO)27 28TARGET_SR = 1600029 30app = FastAPI(title="Smart Turn Endpoint Detector", version="1.0.0")31 32# Allow cross-origin requests (optional)33app.add_middleware(34    CORSMiddleware,35    allow_origins=["*"],36    allow_credentials=True,37    allow_methods=["*"],38    allow_headers=["*"],39)40 41 42class ArrayInput(BaseModel):43    # Raw 1-D samples; can be any sample rate (will resample to 16 kHz)44    samples: List[float]45    sample_rate: Optional[int] = TARGET_SR  # default 16k if not provided46 47 48def _to_mono_float32(y: np.ndarray) -> np.ndarray:49    if y.ndim == 2:50        y = y.mean(axis=1)51    return y.astype(np.float32, copy=False)52 53 54def _resample_to_target_sr(y: np.ndarray, orig_sr: int, target_sr: int = TARGET_SR) -> np.ndarray:55    if orig_sr == target_sr:56        return y.astype(np.float32, copy=False)57    g = math.gcd(orig_sr, target_sr)58    up = target_sr // g59    down = orig_sr // g60    y_resampled = resample_poly(y, up=up, down=down)61    return y_resampled.astype(np.float32, copy=False)62 63 64def _load_audio_from_upload(file_bytes: bytes) -> tuple[np.ndarray, int]:65    # Decode with soundfile (supports wav, flac, ogg; not mp3)66    try:67        with sf.SoundFile(io.BytesIO(file_bytes)) as f:68            y = f.read(dtype="float32", always_2d=False)69            sr = int(f.samplerate)70    except Exception as e:71        raise HTTPException(72            status_code=400,73            detail=f"Could not read audio file. Ensure it's WAV/FLAC/OGG and not corrupted. Error: {e}",74        )75    y = _to_mono_float32(np.asarray(y))76    return y, sr77 78 79@app.get("/")80def root():81    return {82        "name": "Smart Turn Endpoint Detector",83        "status": "ok",84        "model_file": getattr(inf, "ONNX_MODEL_PATH", "smart-turn-v3.0.onnx"),85        "target_sample_rate_hz": TARGET_SR,86        "endpoints": {87            "health": "/healthz",88            "predict_file": "/predict/file",89            "predict_array": "/predict/array",90        },91    }92 93 94@app.get("/healthz")95def healthz():96    try:97        _ = inf.session98        _ = inf.feature_extractor99    except Exception as e:100        raise HTTPException(status_code=500, detail=f"Model not ready: {e}")101    return {"status": "healthy"}102 103 104@app.post("/predict/file")105async def predict_file(audio_file: UploadFile = File(...)):106    """107    Accepts an uploaded audio file (WAV/FLAC/OGG). MP3 is not supported by libsndfile.108    The audio will be converted to mono and resampled to 16 kHz before inference.109    """110    file_bytes = await audio_file.read()111    if not file_bytes:112        raise HTTPException(status_code=400, detail="Empty file upload.")113 114    y, sr = _load_audio_from_upload(file_bytes)115    y16 = _resample_to_target_sr(y, sr, TARGET_SR)116 117    result = inf.predict_endpoint(y16)118 119    return {120        "prediction": int(result["prediction"]),121        "probability": float(result["probability"]),122        "input_info": {123            "filename": audio_file.filename,124            "original_sample_rate_hz": sr,125            "original_num_samples": int(len(y)),126            "original_duration_s": float(len(y) / max(sr, 1)),127        },128        "preproc_info": {129            "sample_rate_hz": TARGET_SR,130            "num_samples": int(len(y16)),131            "duration_s": float(len(y16) / TARGET_SR),132        },133    }134 135 136@app.post("/predict/array")137def predict_array(body: ArrayInput):138    """139    Accepts raw samples in JSON:140    {141      "samples": [ ... floats ... ],142      "sample_rate": 16000   // optional, defaults to 16000143    }144    """145    if not body.samples:146        raise HTTPException(status_code=400, detail="`samples` list is empty.")147 148    y = _to_mono_float32(np.asarray(body.samples, dtype=np.float32))149    sr = int(body.sample_rate) if body.sample_rate else TARGET_SR150    y16 = _resample_to_target_sr(y, sr, TARGET_SR)151 152    result = inf.predict_endpoint(y16)153 154    return {155        "prediction": int(result["prediction"]),156        "probability": float(result["probability"]),157        "input_info": {158            "original_sample_rate_hz": sr,159            "original_num_samples": int(len(y)),160            "original_duration_s": float(len(y) / max(sr, 1)),161        },162        "preproc_info": {163            "sample_rate_hz": TARGET_SR,164            "num_samples": int(len(y16)),165            "duration_s": float(len(y16) / TARGET_SR),166        },167    }168 169 170if __name__ == "__main__":171    import uvicorn172    port = int(os.environ.get("PORT", "7860"))173    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=False)