Spaces:

ankandrew
/

emotion2vec_plus_large

Running on Zero

Upload 5 files

5e2fc72 verified 19 days ago

9.25 kB

from __future__ import annotations

import os
os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")

import spaces
from pathlib import Path
import math

import gradio as gr
import soundfile as sf
import torch
from funasr import AutoModel as FunASRAutoModel
from transformers import AutoModelForAudioFrameClassification, AutoProcessor

from pipeline_utils import (
    SAMPLE_RATE,
    clean_segments,
    export_json,
    format_time,
    prepare_audio,
    prepare_youtube,
    timeline_html,
)

DIAR_MODEL = "nvidia/Nemotron-3-Diarization"
EMOTION_MODEL = "emotion2vec/emotion2vec_plus_large"

EMOTIONS = {
    "开心": "Happy", "happy": "Happy",
    "生气": "Angry", "angry": "Angry",
    "难过": "Sad", "sad": "Sad",
    "中立": "Neutral", "neutral": "Neutral",
    "吃惊": "Surprised", "surprised": "Surprised", "surprise": "Surprised",
    "恐惧": "Fearful", "fearful": "Fearful", "fear": "Fearful",
    "厌恶": "Disgusted", "disgusted": "Disgusted", "disgust": "Disgusted",
    "其他": "Other", "other": "Other",
    "<unk>": "Unknown", "unknown": "Unknown",
}


def english_emotion(label: str) -> str:
    parts = str(label).strip().lower().split("/")
    for part in reversed(parts):
        if part.strip() in EMOTIONS:
            return EMOTIONS[part.strip()]
    return "Unknown"


def parse_emotion(prediction, threshold: float):
    if not isinstance(prediction, list) or not prediction:
        raise ValueError("No emotion prediction returned.")

    record = prediction[0]
    labels = record.get("labels") or []
    values = record.get("scores") or []
    if len(labels) != len(values) or not labels:
        raise ValueError("Invalid emotion prediction.")

    scores = {}
    for raw_label, raw_value in zip(labels, values):
        value = float(raw_value)
        if not math.isfinite(value):
            continue
        label = english_emotion(raw_label)
        scores[label] = max(value, scores.get(label, float("-inf")))

    if not scores:
        raise ValueError("No valid emotion scores returned.")

    scores = dict(sorted(scores.items(), key=lambda pair: pair[1], reverse=True))
    top = next(iter(scores))
    score = scores[top]
    emotion = top if score >= threshold and top not in {"Other", "Unknown"} else "Uncertain"
    return emotion, round(score, 5), {name: round(value, 5) for name, value in scores.items()}


print("Loading diarization model...")
diar_processor = AutoProcessor.from_pretrained(DIAR_MODEL)
diar_model = AutoModelForAudioFrameClassification.from_pretrained(DIAR_MODEL).to("cuda").eval()

print("Loading emotion model...")
emotion_model = FunASRAutoModel(
    model=EMOTION_MODEL,
    hub="hf",
    device="cuda",
    disable_update=True,
    disable_pbar=True,
    ncpu=2,
)


def prepare(upload_path, youtube_url):
    try:
        if youtube_url and youtube_url.strip():
            wav_path, duration, _ = prepare_youtube(youtube_url)
        else:
            wav_path, duration, _ = prepare_audio(upload_path)
        return wav_path, duration
    except Exception as exc:
        raise gr.Error(str(exc)) from exc


@spaces.GPU(duration=120)
def analyze(wav_path, duration, min_seconds, min_score):
    if not wav_path or not Path(wav_path).is_file():
        raise gr.Error("Upload a file first.")

    waveform, sr = sf.read(wav_path, dtype="float32")
    if waveform.ndim != 1 or sr != SAMPLE_RATE:
        raise gr.Error("Could not process the audio.")

    model_input = diar_processor(waveform, sampling_rate=sr).to(
        diar_model.device, dtype=diar_model.dtype
    )
    with torch.inference_mode():
        logits = diar_model(**model_input).logits
    raw_turns = diar_processor.extract_speaker_dict(logits, model_input.attention_mask)[0]
    del logits, model_input

    turns, skipped = clean_segments(raw_turns, duration, float(min_seconds))
    run_dir = Path(wav_path).parent
    rows, selections, segments = [], [], []

    for turn in turns:
        left = max(0, round(turn["start"] * sr))
        right = min(len(waveform), round(turn["end"] * sr))
        if right <= left:
            continue

        index = len(segments) + 1
        speaker = f"Speaker {turn['speaker'] + 1}"
        clip_path = run_dir / f"clip_{index:03d}.wav"
        sf.write(clip_path, waveform[left:right], sr, subtype="PCM_16")

        emotion, score, scores = "Overlap", None, {}
        if not turn["too_much_overlap"]:
            with torch.inference_mode():
                prediction = emotion_model.generate(
                    input=str(clip_path),
                    granularity="utterance",
                    extract_embedding=False,
                )
            emotion, score, scores = parse_emotion(prediction, float(min_score))

        turn["emotion"] = emotion
        start, end = format_time(turn["start"]), format_time(turn["end"])
        rows.append([speaker, f"{start} – {end}", emotion, score])

        segment = {
            "speaker": speaker,
            "start": round(turn["start"], 3),
            "end": round(turn["end"], 3),
            "emotion": emotion,
            "score": score,
            "class_scores": scores,
            "overlap_ratio": round(turn["overlap_ratio"], 4),
        }
        segments.append(segment)
        selections.append({"path": str(clip_path), "title": f"{speaker} · {emotion} · {start}–{end}"})

    speakers = len({turn["speaker"] for turn in turns})
    summary = f"**{speakers} speakers** · **{len(segments)} segments** · {skipped} skipped"
    if not segments:
        summary = "No segments found. Try a shorter minimum duration."

    export_path = export_json(run_dir / "results.json", {
        "models": {"diarization": DIAR_MODEL, "emotion": EMOTION_MODEL},
        "settings": {"minimum_seconds": min_seconds, "minimum_score": min_score},
        "segments": segments,
    })

    first = selections[0] if selections else None
    return (
        summary,
        timeline_html(turns, max(duration, 0.01)),
        rows,
        first["path"] if first else None,
        first["title"] if first else "Select a segment",
        export_path,
        selections,
        gr.update(visible=True),
    )


def select_segment(selections, evt: gr.SelectData):
    index = evt.index[0] if isinstance(evt.index, (tuple, list)) else evt.index
    if not selections or not isinstance(index, int) or not 0 <= index < len(selections):
        return gr.skip(), gr.skip()
    item = selections[index]
    return item["path"], item["title"]


CSS = """
.gradio-container {max-width: 940px !important; margin: 0 auto;}
#page-title h1 {font-size: 1.65rem; letter-spacing: -.025em; margin-bottom: 0;}
.timeline-wrapper {padding: 12px 8px;}
.speaker-timeline-row {display: flex; align-items: center; gap: 12px; margin: 9px 0;}
.speaker-timeline-row > span {flex: 0 0 76px; font-size: .8rem;}
.timeline-track {flex: 1; height: 20px; background: var(--background-fill-secondary); border-radius: 4px; position: relative; overflow: hidden;}
.timeline-block {position: absolute; top: 0; height: 100%; border-radius: 3px; opacity: .8; filter: saturate(.4);}
.timeline-axis {display: flex; justify-content: space-between; margin: 6px 0 0 88px; font-size: .7rem; opacity: .65;}
.empty-timeline {font-size: .85rem; opacity: .6;}
"""

with gr.Blocks(
    title="Voice Emotions",
    delete_cache=(3600, 7200),
) as demo:
    gr.Markdown("# Voice Emotions", elem_id="page-title")
    with gr.Row():
        upload = gr.File(label="Audio / Video", type="filepath")
        youtube_url = gr.Textbox(label="YouTube", placeholder="Video URL", lines=1)

    with gr.Row():
        min_seconds = gr.Slider(0.5, 10, value=2, step=0.5, label="Min. segment (s)")
        min_score = gr.Slider(0, 1, value=0.4, step=0.05, label="Min. score")

    button = gr.Button("Analyze", variant="primary")
    wav_state = gr.State()
    duration_state = gr.State()
    segments_state = gr.State([])

    with gr.Column(visible=False) as results:
        summary = gr.Markdown()
        timeline = gr.HTML()
        with gr.Row():
            with gr.Column(scale=3):
                table = gr.Dataframe(
                    headers=["Speaker", "Time", "Emotion", "Score"],
                    datatype=["str", "str", "str", "number"],
                    type="array",
                    interactive=False,
                    label="Segments",
                )
            with gr.Column(scale=2):
                player = gr.Audio(label="Clip", type="filepath", interactive=False)
                clip_title = gr.Markdown()
                download = gr.File(label="JSON", interactive=False)

    button.click(
        prepare,
        inputs=[upload, youtube_url],
        outputs=[wav_state, duration_state],
        concurrency_limit=2,
    ).success(
        analyze,
        inputs=[wav_state, duration_state, min_seconds, min_score],
        outputs=[summary, timeline, table, player, clip_title, download, segments_state, results],
        concurrency_limit=1,
    )
    table.select(select_segment, inputs=segments_state, outputs=[player, clip_title])

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(
        show_error=True, theme=gr.themes.Base(primary_hue="slate"), css=CSS
    )

	from __future__ import annotations

	import os
	os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")

	import spaces
	from pathlib import Path
	import math

	import gradio as gr
	import soundfile as sf
	import torch
	from funasr import AutoModel as FunASRAutoModel
	from transformers import AutoModelForAudioFrameClassification, AutoProcessor

	from pipeline_utils import (
	SAMPLE_RATE,
	clean_segments,
	export_json,
	format_time,
	prepare_audio,
	prepare_youtube,
	timeline_html,
	)

	DIAR_MODEL = "nvidia/Nemotron-3-Diarization"
	EMOTION_MODEL = "emotion2vec/emotion2vec_plus_large"

	EMOTIONS = {
	"开心": "Happy", "happy": "Happy",
	"生气": "Angry", "angry": "Angry",
	"难过": "Sad", "sad": "Sad",
	"中立": "Neutral", "neutral": "Neutral",
	"吃惊": "Surprised", "surprised": "Surprised", "surprise": "Surprised",
	"恐惧": "Fearful", "fearful": "Fearful", "fear": "Fearful",
	"厌恶": "Disgusted", "disgusted": "Disgusted", "disgust": "Disgusted",
	"其他": "Other", "other": "Other",
	"<unk>": "Unknown", "unknown": "Unknown",
	}


	def english_emotion(label: str) -> str:
	parts = str(label).strip().lower().split("/")
	for part in reversed(parts):
	if part.strip() in EMOTIONS:
	return EMOTIONS[part.strip()]
	return "Unknown"


	def parse_emotion(prediction, threshold: float):
	if not isinstance(prediction, list) or not prediction:
	raise ValueError("No emotion prediction returned.")

	record = prediction[0]
	labels = record.get("labels") or []
	values = record.get("scores") or []
	if len(labels) != len(values) or not labels:
	raise ValueError("Invalid emotion prediction.")

	scores = {}
	for raw_label, raw_value in zip(labels, values):
	value = float(raw_value)
	if not math.isfinite(value):
	continue
	label = english_emotion(raw_label)
	scores[label] = max(value, scores.get(label, float("-inf")))

	if not scores:
	raise ValueError("No valid emotion scores returned.")

	scores = dict(sorted(scores.items(), key=lambda pair: pair[1], reverse=True))
	top = next(iter(scores))
	score = scores[top]
	emotion = top if score >= threshold and top not in {"Other", "Unknown"} else "Uncertain"
	return emotion, round(score, 5), {name: round(value, 5) for name, value in scores.items()}


	print("Loading diarization model...")
	diar_processor = AutoProcessor.from_pretrained(DIAR_MODEL)
	diar_model = AutoModelForAudioFrameClassification.from_pretrained(DIAR_MODEL).to("cuda").eval()

	print("Loading emotion model...")
	emotion_model = FunASRAutoModel(
	model=EMOTION_MODEL,
	hub="hf",
	device="cuda",
	disable_update=True,
	disable_pbar=True,
	ncpu=2,
	)


	def prepare(upload_path, youtube_url):
	try:
	if youtube_url and youtube_url.strip():
	wav_path, duration, _ = prepare_youtube(youtube_url)
	else:
	wav_path, duration, _ = prepare_audio(upload_path)
	return wav_path, duration
	except Exception as exc:
	raise gr.Error(str(exc)) from exc


	@spaces.GPU(duration=120)
	def analyze(wav_path, duration, min_seconds, min_score):
	if not wav_path or not Path(wav_path).is_file():
	raise gr.Error("Upload a file first.")

	waveform, sr = sf.read(wav_path, dtype="float32")
	if waveform.ndim != 1 or sr != SAMPLE_RATE:
	raise gr.Error("Could not process the audio.")

	model_input = diar_processor(waveform, sampling_rate=sr).to(
	diar_model.device, dtype=diar_model.dtype
	)
	with torch.inference_mode():
	logits = diar_model(**model_input).logits
	raw_turns = diar_processor.extract_speaker_dict(logits, model_input.attention_mask)[0]
	del logits, model_input

	turns, skipped = clean_segments(raw_turns, duration, float(min_seconds))
	run_dir = Path(wav_path).parent
	rows, selections, segments = [], [], []

	for turn in turns:
	left = max(0, round(turn["start"] * sr))
	right = min(len(waveform), round(turn["end"] * sr))
	if right <= left:
	continue

	index = len(segments) + 1
	speaker = f"Speaker {turn['speaker'] + 1}"
	clip_path = run_dir / f"clip_{index:03d}.wav"
	sf.write(clip_path, waveform[left:right], sr, subtype="PCM_16")

	emotion, score, scores = "Overlap", None, {}
	if not turn["too_much_overlap"]:
	with torch.inference_mode():
	prediction = emotion_model.generate(
	input=str(clip_path),
	granularity="utterance",
	extract_embedding=False,
	)
	emotion, score, scores = parse_emotion(prediction, float(min_score))

	turn["emotion"] = emotion
	start, end = format_time(turn["start"]), format_time(turn["end"])
	rows.append([speaker, f"{start} – {end}", emotion, score])

	segment = {
	"speaker": speaker,
	"start": round(turn["start"], 3),
	"end": round(turn["end"], 3),
	"emotion": emotion,
	"score": score,
	"class_scores": scores,
	"overlap_ratio": round(turn["overlap_ratio"], 4),
	}
	segments.append(segment)
	selections.append({"path": str(clip_path), "title": f"{speaker} · {emotion} · {start}–{end}"})

	speakers = len({turn["speaker"] for turn in turns})
	summary = f"{speakers} speakers · {len(segments)} segments · {skipped} skipped"
	if not segments:
	summary = "No segments found. Try a shorter minimum duration."

	export_path = export_json(run_dir / "results.json", {
	"models": {"diarization": DIAR_MODEL, "emotion": EMOTION_MODEL},
	"settings": {"minimum_seconds": min_seconds, "minimum_score": min_score},
	"segments": segments,
	})

	first = selections[0] if selections else None
	return (
	summary,
	timeline_html(turns, max(duration, 0.01)),
	rows,
	first["path"] if first else None,
	first["title"] if first else "Select a segment",
	export_path,
	selections,
	gr.update(visible=True),
	)


	def select_segment(selections, evt: gr.SelectData):
	index = evt.index[0] if isinstance(evt.index, (tuple, list)) else evt.index
	if not selections or not isinstance(index, int) or not 0 <= index < len(selections):
	return gr.skip(), gr.skip()
	item = selections[index]
	return item["path"], item["title"]


	CSS = """
	.gradio-container {max-width: 940px !important; margin: 0 auto;}
	#page-title h1 {font-size: 1.65rem; letter-spacing: -.025em; margin-bottom: 0;}
	.timeline-wrapper {padding: 12px 8px;}
	.speaker-timeline-row {display: flex; align-items: center; gap: 12px; margin: 9px 0;}
	.speaker-timeline-row > span {flex: 0 0 76px; font-size: .8rem;}
	.timeline-track {flex: 1; height: 20px; background: var(--background-fill-secondary); border-radius: 4px; position: relative; overflow: hidden;}
	.timeline-block {position: absolute; top: 0; height: 100%; border-radius: 3px; opacity: .8; filter: saturate(.4);}
	.timeline-axis {display: flex; justify-content: space-between; margin: 6px 0 0 88px; font-size: .7rem; opacity: .65;}
	.empty-timeline {font-size: .85rem; opacity: .6;}
	"""

	with gr.Blocks(
	title="Voice Emotions",
	delete_cache=(3600, 7200),
	) as demo:
	gr.Markdown("# Voice Emotions", elem_id="page-title")
	with gr.Row():
	upload = gr.File(label="Audio / Video", type="filepath")
	youtube_url = gr.Textbox(label="YouTube", placeholder="Video URL", lines=1)

	with gr.Row():
	min_seconds = gr.Slider(0.5, 10, value=2, step=0.5, label="Min. segment (s)")
	min_score = gr.Slider(0, 1, value=0.4, step=0.05, label="Min. score")

	button = gr.Button("Analyze", variant="primary")
	wav_state = gr.State()
	duration_state = gr.State()
	segments_state = gr.State([])

	with gr.Column(visible=False) as results:
	summary = gr.Markdown()
	timeline = gr.HTML()
	with gr.Row():
	with gr.Column(scale=3):
	table = gr.Dataframe(
	headers=["Speaker", "Time", "Emotion", "Score"],
	datatype=["str", "str", "str", "number"],
	type="array",
	interactive=False,
	label="Segments",
	)
	with gr.Column(scale=2):
	player = gr.Audio(label="Clip", type="filepath", interactive=False)
	clip_title = gr.Markdown()
	download = gr.File(label="JSON", interactive=False)

	button.click(
	prepare,
	inputs=[upload, youtube_url],
	outputs=[wav_state, duration_state],
	concurrency_limit=2,
	).success(
	analyze,
	inputs=[wav_state, duration_state, min_seconds, min_score],
	outputs=[summary, timeline, table, player, clip_title, download, segments_state, results],
	concurrency_limit=1,
	)
	table.select(select_segment, inputs=segments_state, outputs=[player, clip_title])

	if __name__ == "__main__":
	demo.queue(default_concurrency_limit=1).launch(
	show_error=True, theme=gr.themes.Base(primary_hue="slate"), css=CSS
	)