|
import pprint |
|
|
|
import gradio as gr |
|
import librosa |
|
import plotly.graph_objects as go |
|
import spaces |
|
import torch |
|
from loguru import logger |
|
from transformers import AutoFeatureExtractor |
|
from transformers.modeling_outputs import SequenceClassifierOutput |
|
|
|
from model import EmotionModel |
|
|
|
repo_id = "my_model" |
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
logger.info(f"device: {device}") |
|
model = EmotionModel.from_pretrained(repo_id, device_map=device) |
|
model.eval() |
|
processor = AutoFeatureExtractor.from_pretrained(repo_id) |
|
|
|
label_map = { |
|
"Angry": "😠 怒り", |
|
"Disgusted": "😒 嫌悪", |
|
"Embarrassed": "😳 戸惑い", |
|
"Fearful": "😨 恐怖", |
|
"Happy": "😊 幸せ", |
|
"Sad": "😢 悲しみ", |
|
"Surprised": "😲 驚き", |
|
"Neutral": "😐 中立", |
|
"Sexual1": "🥰 NSFW1", |
|
"Sexual2": "🍭 NSFW2", |
|
} |
|
|
|
|
|
@spaces.GPU |
|
def pipe(filename: str) -> tuple[dict[str, float], go.Figure]: |
|
audio, sr = librosa.load(filename, sr=16000) |
|
duration = librosa.get_duration(y=audio, sr=sr) |
|
logger.info(f"filename: {filename}, duration: {duration}") |
|
if duration > 30.0: |
|
return ( |
|
{f"Error: 音声ファイルの長さが長すぎます: {duration}秒": 0.0}, |
|
go.Figure(), |
|
) |
|
inputs = processor(audio, sampling_rate=sr, return_tensors="pt") |
|
inputs = {k: v.to(device) for k, v in inputs.items()} |
|
with torch.no_grad(): |
|
outputs: SequenceClassifierOutput = model(**inputs) |
|
logits = outputs.logits |
|
|
|
logits = logits[0].cpu().numpy() |
|
labels = [label_map[label] for id, label in model.config.id2label.items()] |
|
sorted_pairs = sorted(zip(logits, labels), key=lambda x: x[0]) |
|
sorted_logits, sorted_labels = zip(*sorted_pairs) |
|
logger.info(f"Result:\n{pprint.pformat(sorted_pairs)}") |
|
probabilities = outputs.logits.softmax(dim=-1) |
|
scores_dict = {label: prob.item() for label, prob in zip(labels, probabilities[0])} |
|
fig = go.Figure([go.Bar(x=sorted_logits, y=sorted_labels, orientation="h")]) |
|
return scores_dict, fig |
|
|
|
|
|
md = """ |
|
# 音声からの感情認識 ver 0.1 |
|
|
|
- 音声ファイルから感情を予測して、確率とlogits (softmax前の値) を表示します |
|
- 30秒以上の音声ファイルは受け付けません |
|
""" |
|
|
|
with gr.Blocks() as app: |
|
gr.Markdown(md) |
|
audio = gr.Audio(type="filepath") |
|
btn = gr.Button("感情を予測") |
|
with gr.Row(): |
|
result = gr.Label(label="結果") |
|
plot = gr.Plot(label="Logits") |
|
|
|
btn.click(pipe, inputs=audio, outputs=[result, plot]) |
|
|
|
app.launch() |
|
|
