File size: 3,789 Bytes
c6c2fd0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8e2ced3
 
 
 
 
 
c6c2fd0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
360a832
c6c2fd0
 
 
 
 
 
 
0aad17a
 
 
 
 
 
 
 
c6c2fd0
 
 
 
 
 
 
 
 
 
0aad17a
d8eeff7
0aad17a
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
import gradio as gr
import numpy as np
import spaces
import torch
import torch.nn as nn
from transformers import Wav2Vec2Processor
from transformers.models.wav2vec2.modeling_wav2vec2 import Wav2Vec2Model
from transformers.models.wav2vec2.modeling_wav2vec2 import Wav2Vec2PreTrainedModel

import audiofile


class ModelHead(nn.Module):
    r"""Classification head."""

    def __init__(self, config, num_labels):

        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, num_labels)

    def forward(self, features, **kwargs):

        x = features
        x = self.dropout(x)
        x = self.dense(x)
        x = torch.tanh(x)
        x = self.dropout(x)
        x = self.out_proj(x)

        return x


class AgeGenderModel(Wav2Vec2PreTrainedModel):
    r"""Speech emotion classifier."""

    def __init__(self, config):

        super().__init__(config)

        self.config = config
        self.wav2vec2 = Wav2Vec2Model(config)
        self.age = ModelHead(config, 1)
        self.gender = ModelHead(config, 3)
        self.init_weights()

    def forward(
            self,
            input_values,
    ):

        outputs = self.wav2vec2(input_values)
        hidden_states = outputs[0]
        hidden_states = torch.mean(hidden_states, dim=1)
        logits_age = self.age(hidden_states)
        logits_gender = torch.softmax(self.gender(hidden_states), dim=1)

        return hidden_states, logits_age, logits_gender



# load model from hub
device = 0 if torch.cuda.is_available() else "cpu"
model_name = "audeering/wav2vec2-large-robust-24-ft-age-gender"
processor = Wav2Vec2Processor.from_pretrained(model_name)
model = AgeGenderModel.from_pretrained(model_name)


def process_func(x: np.ndarray, sampling_rate: int) -> dict:
    r"""Predict age and gender or extract embeddings from raw audio signal."""
    # run through processor to normalize signal
    # always returns a batch, so we just get the first entry
    # then we put it on the device
    y = processor(x, sampling_rate=sampling_rate)
    y = y['input_values'][0]
    y = y.reshape(1, -1)
    y = torch.from_numpy(y).to(device)

    # run through model
    with torch.no_grad():
        y = model(y)
        y = torch.hstack([y[1], y[2]])

    # convert to numpy
    y = y.detach().cpu().numpy()

    # convert to dict
    y = {
        "age": 100 * y[0][0],
        "female": y[0][1],
        "male": y[0][2],
        "child": y[0][3],
    }

    return y


@spaces.GPU
def recognize(file):
    if file is None:
        raise gr.Error(
            "No audio file submitted! "
            "Please upload or record an audio file "
            "before submitting your request."
        )
    signal, sampling_rate = audiofile.read(file)
    age_gender = process_func(signal, sampling_rate)
    return  age_gender


outputs = gr.Label()
title = "audEERING age and gender recognition"
description = (
    "Recognize age and gender of a microphone recording or audio file. "
    "Demo uses the checkpoint [{model_name}](https://huggingface.co/{model_name})."
)
allow_flagging = "never"

# microphone = gr.Interface(
#     fn=recognize,
#     inputs=gr.Audio(sources="microphone", type="filepath"),
#     outputs=outputs,
#     title=title,
#     description=description,
#     allow_flagging=allow_flagging,
# )

file = gr.Interface(
    fn=recognize,
    inputs=gr.Audio(sources="upload", type="filepath", label="Audio file"),
    outputs=outputs,
    title=title,
    description=description,
    allow_flagging=allow_flagging,
)

# demo = gr.TabbedInterface([microphone, file], ["Microphone", "Audio file"])
# demo.queue().launch()
# demo.launch()
file.launch()