Spaces:
Running on Zero
Running on Zero
| import matplotlib | |
| matplotlib.use("Agg") | |
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| import torch.nn as nn | |
| import librosa | |
| import gradio as gr | |
| import matplotlib.pyplot as plt | |
| from scipy.signal import find_peaks | |
| from scipy.ndimage import uniform_filter1d | |
| from transformers import Wav2Vec2Processor, pipeline | |
| from transformers.models.wav2vec2.modeling_wav2vec2 import Wav2Vec2Model, Wav2Vec2PreTrainedModel | |
| from silero_vad import load_silero_vad, get_speech_timestamps | |
| try: | |
| import spaces # available on ZeroGPU Spaces | |
| gpu = spaces.GPU(duration=60) | |
| except Exception: # running locally / on CPU Space | |
| def gpu(fn): | |
| return fn | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| SR = 16000 | |
| # ---------------- age + gender model (custom head from the model card) ---------------- | |
| class ModelHead(nn.Module): | |
| def __init__(self, config, n): | |
| super().__init__() | |
| self.dense = nn.Linear(config.hidden_size, config.hidden_size) | |
| self.dropout = nn.Dropout(config.final_dropout) | |
| self.out_proj = nn.Linear(config.hidden_size, n) | |
| def forward(self, x): | |
| x = self.dropout(x) | |
| x = torch.tanh(self.dense(x)) | |
| x = self.dropout(x) | |
| return self.out_proj(x) | |
| class AgeGenderModel(Wav2Vec2PreTrainedModel): | |
| def __init__(self, config): | |
| super().__init__(config) | |
| self.wav2vec2 = Wav2Vec2Model(config) | |
| self.age = ModelHead(config, 1) | |
| self.gender = ModelHead(config, 3) | |
| self.post_init() | |
| def forward(self, x): | |
| h = self.wav2vec2(x)[0].mean(1) | |
| return self.age(h), torch.softmax(self.gender(h), dim=1) | |
| AG_ID = "audeering/wav2vec2-large-robust-24-ft-age-gender" | |
| ag_proc = Wav2Vec2Processor.from_pretrained(AG_ID) | |
| ag_model = AgeGenderModel.from_pretrained(AG_ID).to(DEVICE).eval() | |
| dev_idx = 0 if DEVICE == "cuda" else -1 | |
| emo_pipe = pipeline("audio-classification", model="superb/wav2vec2-base-superb-er", device=dev_idx) | |
| lid_pipe = pipeline("audio-classification", model="facebook/mms-lid-126", device=dev_idx) | |
| vad_model = load_silero_vad() | |
| EMO = {"neu": "neutral", "hap": "happy", "ang": "angry", "sad": "sad"} | |
| def vad_segments(y): | |
| ts = get_speech_timestamps(torch.from_numpy(np.asarray(y, np.float32)), vad_model, | |
| sampling_rate=SR, return_seconds=True) | |
| return [(float(t["start"]), float(t["end"])) for t in ts] | |
| def age_bucket(a): | |
| for hi, name in [(12, "child"), (19, "teen"), (30, "young adult"), (50, "adult"), (65, "middle-aged")]: | |
| if a < hi: | |
| return name | |
| return "senior" | |
| # ---------------- GPU part: only model inference ---------------- | |
| def gpu_infer(speech): | |
| x = ag_proc(speech, sampling_rate=SR)["input_values"][0].reshape(1, -1) | |
| age_out, g = ag_model(torch.from_numpy(x).float().to(DEVICE)) | |
| age = float(age_out.item()) * 100 | |
| gp = g[0].cpu().numpy() | |
| emo = emo_pipe({"raw": speech, "sampling_rate": SR}, top_k=1)[0] | |
| lid = lid_pipe({"raw": speech, "sampling_rate": SR}, top_k=1)[0] | |
| return age, gp, emo, lid | |
| # ---------------- CPU part: audio processing ---------------- | |
| def analyze_voice(path): | |
| y, _ = librosa.load(path, sr=SR, mono=True, duration=30) | |
| if len(y) < SR // 2: | |
| raise gr.Error("Audio is too short. Please upload at least 1-2 seconds of speech.") | |
| y = y / (np.abs(y).max() + 1e-6) | |
| try: | |
| segs = vad_segments(y) | |
| except Exception: | |
| segs = [(a / SR, b / SR) for a, b in librosa.effects.split(y, top_db=30)] | |
| speech = np.concatenate([y[int(a * SR):int(b * SR)] for a, b in segs]) if segs else y | |
| speech_dur = sum(b - a for a, b in segs) or len(y) / SR | |
| if len(speech) < SR: | |
| speech = np.pad(speech, (0, SR - len(speech))) | |
| age, gp, emo, lid = gpu_infer(speech) | |
| gender = ["female", "male", "child"][int(gp.argmax())] | |
| f0, _, _ = librosa.pyin(y, fmin=65, fmax=500, sr=SR) | |
| times = librosa.times_like(f0, sr=SR) | |
| f0v = f0[~np.isnan(f0)] | |
| if len(f0v): | |
| pitch = dict(median=np.median(f0v), p10=np.percentile(f0v, 10), p90=np.percentile(f0v, 90)) | |
| else: | |
| pitch = dict(median=np.nan, p10=np.nan, p90=np.nan) | |
| rms = uniform_filter1d(librosa.feature.rms(y=speech, frame_length=400, hop_length=160)[0], 5) | |
| peaks, _ = find_peaks(rms, height=0.1 * rms.max(), prominence=0.05 * rms.max(), distance=8) | |
| rate = len(peaks) / speech_dur | |
| return dict(age=age, bucket=age_bucket(age), gender=f"{gender} ({gp.max():.0%})", | |
| emotion=f"{EMO.get(emo['label'], emo['label'])} ({emo['score']:.0%})", | |
| language=f"{lid['label']} ({lid['score']:.0%})", | |
| pitch=pitch, rate=rate, speech_dur=speech_dur, f0=f0, times=times) | |
| def compare(pa, pb): | |
| if pa is None or pb is None: | |
| raise gr.Error("Please provide both Voice A and Voice B.") | |
| A, B = analyze_voice(pa), analyze_voice(pb) | |
| rows = [ | |
| ("Gender (estimated)", A["gender"], B["gender"]), | |
| ("Age (estimated)", f"{A['age']:.0f} yrs ({A['bucket']})", f"{B['age']:.0f} yrs ({B['bucket']})"), | |
| ("Emotion", A["emotion"], B["emotion"]), | |
| ("Language (ISO 639-3)", A["language"], B["language"]), | |
| ("Median pitch (Hz)", f"{A['pitch']['median']:.0f}", f"{B['pitch']['median']:.0f}"), | |
| ("Pitch range p10-p90 (Hz)", f"{A['pitch']['p10']:.0f}-{A['pitch']['p90']:.0f}", | |
| f"{B['pitch']['p10']:.0f}-{B['pitch']['p90']:.0f}"), | |
| ("Speaking rate (syllables/s)", f"{A['rate']:.1f}", f"{B['rate']:.1f}"), | |
| ("Speech duration (s)", f"{A['speech_dur']:.1f}", f"{B['speech_dur']:.1f}"), | |
| ] | |
| df = pd.DataFrame(rows, columns=["Attribute", "Voice A", "Voice B"]) | |
| fig, ax = plt.subplots(figsize=(10, 3.5)) | |
| ax.plot(A["times"], A["f0"], label="Voice A") | |
| ax.plot(B["times"], B["f0"], label="Voice B") | |
| ax.set_xlabel("time (s)") | |
| ax.set_ylabel("F0 (Hz)") | |
| ax.legend() | |
| fig.tight_layout() | |
| return df, fig | |
| DESCRIPTION = ( | |
| "Upload or record two voices (3-30 s of clear speech each) to compare estimated gender, age, " | |
| "emotion, language, pitch and speaking rate. Results are model estimates, not facts about a person. " | |
| "Only use recordings of people who have agreed to be analyzed." | |
| ) | |
| demo = gr.Interface( | |
| fn=compare, | |
| inputs=[gr.Audio(type="filepath", label="Voice A"), gr.Audio(type="filepath", label="Voice B")], | |
| outputs=[gr.Dataframe(label="Attributes"), gr.Plot(label="Pitch contour")], | |
| title="Voice attribute comparison", | |
| description=DESCRIPTION, | |
| flagging_mode="never", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |