add model logic implementation

.gitignore

@@ -16,4 +16,6 @@ build/
 
 # VSCode
 .vscode/
-*.code-workspace
+*.code-workspace
+
+behaviour_model/

README.md

@@ -9,6 +9,14 @@ app_file: app.py
 pinned: false
 license: mit
 short_description: Audio and Text Emotion Recognition
+models:
+  - links-ads/kk-speech-emotion-recognition
+  - openai/whisper-large-v3
+  - facebook/wav2vec2-large-xlsr-53
+preload_from_hub:
+  - links-ads/kk-speech-emotion-recognition
+  - openai/whisper-large-v3
+  - facebook/wav2vec2-large-xlsr-53
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,13 +1,40 @@
+import torch
 import gradio as gr
-import plotly.graph_objects as go
 from src.load_html import get_description_html
-
-def process_audio(audio_file):
-    ...
+from src.audio_processor import AudioProcessor
+from src.model.behaviour_model import get_behaviour_model
+from transformers import (
+    pipeline,
+    WavLMForSequenceClassification
+)
 
 
 # Gradio interface
 def create_demo():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    segmentation_model = pipeline(
+        task="automatic-speech-recognition",
+        model="openai/whisper-large-v3-turbo",
+        tokenizer="openai/whisper-large-v3-turbo",
+        device=device
+    )
+
+    emotion_model = WavLMForSequenceClassification.from_pretrained("links-ads/kk-speech-emotion-recognition")
+    emotion_model.to(device)
+    emotion_model.eval()
+
+    behaviour_model = get_behaviour_model(
+        behaviour_model_path="behaviour_model/",
+        device=device,
+    )        
+
+    audio_processor = AudioProcessor(
+        emotion_model=emotion_model,
+        segmentation_model=segmentation_model,
+        device=device,
+        behaviour_model=behaviour_model,
+    )
+
     with gr.Blocks() as demo:
         gr.HTML(get_description_html)
         
@@ -17,7 +44,7 @@ def create_demo():
         graph_output = gr.Plot(label="Generated Graph")
 
         submit_button.click(
-            fn=process_audio, 
+            fn=audio_processor, 
             inputs=audio_input, 
             outputs=graph_output
         )

requirements.txt

@@ -1,2 +1,5 @@
 gradio==5.24.0
-plotly==6.0.1
+plotly==6.0.1
+torch==2.7.0
+librosa==0.11.0
+transformers==4.51.3

src/audio_processor.py

@@ -0,0 +1,190 @@
+import time
+import torch
+import librosa
+import numpy as np
+import gradio as gr
+import gradio as gr
+from .generate_graph import create_behaviour_gantt_plot
+from transformers import Wav2Vec2Processor
+
+
+SAMPLING_RATE = 16_000
+
+class AudioProcessor:
+    def __init__(
+        self,
+        emotion_model,
+        segmentation_model,
+        device,
+        behaviour_model=None,
+    ):
+        self.emotion_processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
+        self.emotion_model = emotion_model
+        self.behaviour_model = behaviour_model
+        self.device = device
+        self.audio_emotion_labels = {
+            0: "Neutralità",
+            1: "Rabbia",
+            2: "Paura",
+            3: "Gioia",
+            4: "Sorpresa",
+            5: "Tristezza",
+            6: "Disgusto",
+        }
+        self.emotion_translation = {
+            "neutrality": "Neutralità",
+            "anger": "Rabbia",
+            "fear": "Paura",
+            "joy": "Gioia",
+            "surprise": "Sorpresa",
+            "sadness": "Tristezza",
+            "disgust": "Disgusto"
+        }
+        self.behaviour_labels = {
+            0: "frustrated",
+            1: "delighted",
+            2: "dysregulated",
+        }
+        self.behaviour_translation = {
+            "frustrated": "frustazione",
+            "delighted": "incantato",
+            "dysregulated": "disregolazione",    
+        }
+        self.segmentation_model = segmentation_model
+
+        self._set_emotion_model()
+        if self.behaviour_model:
+            self._set_behaviour_model()
+
+        self.behaviour_confidence = 0.6
+
+        self.chart_generator = None
+
+    def _set_emotion_model(self):
+        self.emotion_model.to(self.device)
+        self.emotion_model.eval()
+
+    def _set_behaviour_model(self):
+        self.behaviour_model.to(self.device)
+        self.behaviour_model.eval()
+
+    def _prepare_transcribed_text(self, chunks):
+        formated_timestamps = []
+        predictions = []
+
+        for chunk in chunks:
+            start = chunk[0] / SAMPLING_RATE
+            end = chunk[1] / SAMPLING_RATE
+            formated_start = time.strftime('%H:%M:%S', time.gmtime(start))
+            formated_end = time.strftime('%H:%M:%S', time.gmtime(end))
+            formated_timestamps.append(f"**({formated_start} - {formated_end})**")
+
+            predictions.append(f"**[{chunk[2]}]**")
+
+        transcribed_texts = [chunk[3] for chunk in chunks]
+        transcribed_text = "<br/>".join(
+            [
+                f"{formated_timestamps[i]}: {transcribed_texts[i]} {predictions[i]}" for i in range(len(transcribed_texts))
+            ]
+        )
+
+        print(f"Transcribed text:\n{transcribed_text}")
+
+
+        return transcribed_text
+
+    def __call__(self, audio_path: str):
+        """
+        Predicts the emotion label for a given audio input.
+
+        Args:
+            audio (filepath): The audio input path to be processed.
+
+        Returns:
+            str: The predicted emotion label.
+
+        """
+
+        print("Segmenting audio...")
+        out = self.segmentation_model(
+            inputs=audio_path,
+            return_timestamps=True,
+        )
+
+        emotion_chunks = []
+        behaviour_chunks = []
+        timestamps = []
+        predicted_labels = []
+        all_probabilities = []
+
+        print("Analizing chunks...")
+        for chunk in out["chunks"]:                
+            # trim audio from timestamps
+            start = int(chunk["timestamp"][0] * SAMPLING_RATE)
+            end = int(chunk["timestamp"][1] * SAMPLING_RATE if chunk["timestamp"][1] else len(input_frames))
+
+            audio = input_frames[start:end]
+
+            inputs = self.emotion_processor(audio, chunk["text"], return_tensors="pt", sampling_rate=SAMPLING_RATE)
+
+            print(f"Inputs: {inputs}")
+
+            if "input_values" in inputs:
+                inputs["input_features"] = inputs.pop("input_values")
+
+            inputs['input_features'] = inputs['input_features'].to(self.device)
+            inputs['input_ids'] = inputs['input_ids'].to(self.device)
+            inputs['text_attention_mask'] = inputs['text_attention_mask'].to(self.device)
+
+            print("Predicting emotion for chunk...")
+            logits = self.emotion_model(**inputs).logits
+
+            logits = logits.detach().cpu()
+            softmax = torch.nn.Softmax(dim=1)
+            probabilities = softmax(logits).squeeze(0)
+
+            prediction = probabilities.argmax().item()
+
+            predicted_label = self.emotion_processor.config.id2label[prediction]
+            label_translation = self.emotion_translation[predicted_label]
+
+            emotion_chunks.append(
+                (
+                    start, 
+                    end,
+                    label_translation,
+                    chunk["text"], 
+                    np.round(probabilities[prediction].item(), 2)
+                )
+            )
+
+            timestamps.append((start, end))
+            predicted_labels.append(label_translation)
+            all_probabilities.append(probabilities[prediction].item())
+
+
+            inputs = self.emotion_processor(audio, return_tensors="pt", sampling_rate=SAMPLING_RATE)
+            if "input_values" in inputs:
+                inputs["input_features"] = inputs.pop("input_values")
+            
+            inputs = inputs.input_features.to(self.device)
+            print("Predicting behaviour for chunk...")
+            logits = self.behaviour_model(inputs).logits
+            probabilities = torch.nn.functional.softmax(logits.detach().cpu(), dim=-1).squeeze()
+            behaviour_chunks.append(
+                (
+                    start, 
+                    end,
+                    chunk["text"],
+                    np.round(probabilities[2].item(), 2),
+                    label_translation,
+                )
+            )
+        behaviour_gantt = create_behaviour_gantt_plot(behaviour_chunks)
+
+        # transcribed_text = self._prepare_transcribed_text(emotion_chunks)
+        
+        return (
+            behaviour_gantt,
+            # transcribed_text,
+        )

src/generate_graph.py

@@ -0,0 +1,176 @@
+import time
+import plotly.graph_objects as go
+from datetime import datetime, timedelta
+
+
+SAMPLING_RATE = 16_000
+COLOR_MAP = {
+    "Neutralità": "rgb(178, 178, 178)",
+    "Rabbia": "rgb(160, 61, 62)",
+    "Paura": "rgb(91, 57, 136)",
+    "Gioia": "rgb(255, 255, 0)",
+    "Sorpresa": "rgb(60, 175, 175)",
+    "Tristezza": "rgb(64, 106, 173)",
+    "Disgusto": "rgb(100, 153, 65)",
+}
+
+def create_behaviour_gantt_plot(behaviour_chunks, confidence_threshold=60):
+    print("Creating behaviour Gantt plot...")
+    emotion_order = [
+        "Gioia", 
+        "Sorpresa", 
+        "Disgusto", 
+        "Tristezza", 
+        "Paura", 
+        "Rabbia", 
+        "Neutralità"
+    ]
+        
+    fig = go.Figure()
+    
+    chunk_starts = [start/SAMPLING_RATE for start, _, _, _, _ in behaviour_chunks]
+    chunk_ends = [end/SAMPLING_RATE for _, end, _, _, _ in behaviour_chunks]
+    
+    # Create reference time for plotting (starting at 0)
+    # We'll use a base datetime and add seconds
+    base_time = datetime(2_000, 1, 1, 0, 0, 0) # TODO: change magic numbers
+    
+    start_times = [base_time + timedelta(seconds=t) for t in chunk_starts]
+    end_times = [base_time + timedelta(seconds=t) for t in chunk_ends]
+    
+    # Calculate midpoints for each chunk (for trend line)
+    mid_times = [base_time + timedelta(seconds=(s+e)/2) for s, e in zip(chunk_starts, chunk_ends)]
+    
+    heights = [height * 100 for _, _, _, height, _ in behaviour_chunks]
+    
+    emotions = [emotion for _, _, _, _, emotion in behaviour_chunks]
+    
+    hover_texts = []
+    for i, (start, end, label, height, emotion) in enumerate(behaviour_chunks):
+        start_fmt = time.strftime('%H:%M:%S', time.gmtime(start / SAMPLING_RATE))
+        end_fmt = time.strftime('%H:%M:%S', time.gmtime(end / SAMPLING_RATE))
+        duration_seconds = (end - start) / SAMPLING_RATE
+        duration_str = time.strftime('%H:%M:%S', time.gmtime(duration_seconds))
+        
+        hover_text = f"Inizio: {start_fmt}<br>Fine: {end_fmt}<br>Durata: {duration_str}<br>Testo: {label}<br>Attendibilità: {height*100:.2f}%<br>Emozione: {emotion}"
+        hover_texts.append(hover_text)
+    
+    fig.add_shape(
+        type="rect",
+        x0=start_times[0],
+        x1=end_times[-1],
+        y0=confidence_threshold,
+        y1=100,
+        fillcolor="rgba(188,223,241,0.8)",
+        opacity=0.8,
+        layer="below",
+        line_width=0,
+    )
+    
+    fig.add_hline(y=confidence_threshold, line_dash="dash", line_color="black", line_width=1)
+    
+    fig.add_trace(
+        go.Scatter(
+            x=mid_times,
+            y=heights,
+            mode='lines',
+            name='Disregolazione',
+            line=dict(
+                color='orange', 
+                width=2,
+                shape='spline',  # This enables smoothing
+                smoothing=1.0,   # Adjust smoothing factor
+            ),
+            text=hover_texts,
+            hoverinfo='text',
+            showlegend=False,
+        )
+    )
+    
+    emotion_data = {}
+    
+    for i, height in enumerate(heights):
+        if height >= confidence_threshold:
+            emotion = emotions[i]
+            if emotion not in emotion_data:
+                emotion_data[emotion] = {
+                    'times': [],
+                    'heights': [],
+                    'hover_texts': []
+                }
+            
+            emotion_data[emotion]['times'].append(mid_times[i])
+            emotion_data[emotion]['heights'].append(height)
+            emotion_data[emotion]['hover_texts'].append(hover_texts[i])
+    
+    for emotion in emotion_order:
+        color = COLOR_MAP.get(emotion, '#000000')
+        
+        if emotion in emotion_data:
+            data = emotion_data[emotion]
+            fig.add_trace(
+                go.Scatter(
+                    x=data['times'],
+                    y=data['heights'],
+                    mode='markers',
+                    name=emotion.capitalize(),
+                    marker=dict(
+                        size=15,
+                        color=color,
+                        symbol='circle'
+                    ),
+                    text=data['hover_texts'],
+                    hoverinfo='text',
+                    showlegend=True,
+                )
+            )
+        else:
+            fig.add_trace(
+                go.Scatter(
+                    x=[None],
+                    y=[None],
+                    mode='markers',
+                    name=emotion.capitalize(),
+                    marker=dict(
+                        size=15,
+                        color=color,
+                        symbol='circle'
+                    ),
+                    showlegend=True,
+                )
+            )
+    
+    fig.update_layout(
+        title='Distribuzione della disregolazione',
+        xaxis_title='Tempo',
+        yaxis_title='Attendibilità',
+        xaxis=dict(
+            type='date',
+            tickformat='%H:%M:%S',
+            showline=True,
+            zeroline=False,
+            side='bottom',
+            showgrid=False,
+        ),
+        yaxis=dict(
+            range=[0, 100],
+            tickvals=[0, 20, 40, 60, 80, 100],
+            ticktext=['0%', '20%', '40%', '60%', '80%', '100%'],
+            tickmode='array',
+            showgrid=False,
+        ),
+        legend_title=None,
+        legend=dict(
+            yanchor="top"
+        ),
+        hoverlabel=dict(
+            font_size=12,
+            font_family="Arial"
+        ),
+        paper_bgcolor='white',
+        plot_bgcolor='white',
+    )
+    
+    fig.update_traces(hovertemplate=None)
+    
+    return fig

src/model/behaviour_model.py

@@ -0,0 +1,28 @@
+import os
+import argparse
+import torch
+from .make_model import make_model
+
+
+hparams_dict = {
+    'HF_MODEL_PATH': 'facebook/wav2vec2-large-xlsr-53',
+    'DATASET': 'recanvo',
+    'MAX_DURATION': 4,
+    'SAMPLING_RATE': 16_000,
+    'OUTPUT_HIDDEN_STATES': True,
+    'CLASSIFIER_NAME': 'multilevel',
+    'CLASSIFIER_PROJ_SIZE': 256,
+    'NUM_LABELS': 3,
+    'LABEL_WEIGHTS': [1.0],
+    'LOSS': 'cross-entropy',
+    'GPU_ID': 0,
+    'RETURN_RAW_ARRAY': False,
+}
+hparams = argparse.Namespace(**hparams_dict)
+
+def get_behaviour_model(behaviour_model_path, device):
+    state_dict = torch.load(os.path.join(behaviour_model_path, 'pytorch_model.bin'), map_location=device)
+    model = make_model(hparams)
+    model.load_state_dict(state_dict)
+
+    return model

src/model/custom_model.py

@@ -0,0 +1,38 @@
+import torch
+import torch.nn as nn
+
+from typing import Optional, Union, Tuple
+from transformers.modeling_outputs import SequenceClassifierOutput
+from .wav2vec2_wrapper import Wav2VecWrapper
+from .multilevel_classifier import MultiLevelDownstreamModel
+
+class CustomModelForAudioClassification(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.output_hidden_states == True, "The upstream model must return all hidden states"
+        self.config = config
+        self.encoder = Wav2VecWrapper(config)
+
+        self.classifier = MultiLevelDownstreamModel(config, use_conv_output=True)
+        
+    def forward(
+        self,
+        input_features: Optional[torch.LongTensor],
+        length: Optional[torch.LongTensor] = None,
+        encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
+        if encoder_outputs is None:
+            encoder_output = self.encoder(
+                input_features,
+                length=length,
+            )
+
+        logits = self.classifier(**encoder_output)
+
+        loss = None
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=encoder_output['encoder_hidden_states']
+        )

src/model/make_model.py

@@ -0,0 +1,20 @@
+from transformers import (
+    AutoConfig
+)
+from .custom_model import CustomModelForAudioClassification
+
+def make_model(hparams):
+    """ Returns a model instance based on the provided hyperparameters. """
+    hparams = vars(hparams)
+    config = AutoConfig.from_pretrained(hparams['HF_MODEL_PATH'])
+    config.max_duration = hparams['MAX_DURATION']
+    config.sampling_rate = hparams['SAMPLING_RATE']
+    config.output_hidden_states = hparams['OUTPUT_HIDDEN_STATES']
+    config.classifier_name = hparams['CLASSIFIER_NAME']
+    config.classifier_proj_size = hparams['CLASSIFIER_PROJ_SIZE']
+    config.num_labels = hparams['NUM_LABELS']
+    config.label_weights = hparams['LABEL_WEIGHTS']
+    config.lossname = hparams['LOSS']
+    model = CustomModelForAudioClassification(config)
+    
+    return model

src/model/multilevel_classifier.py

@@ -0,0 +1,76 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+
+class MultiLevelDownstreamModel(nn.Module):
+    def __init__(
+            self,
+            model_config,
+            use_conv_output: Optional[bool] = True,
+        ):
+        super().__init__()
+        assert model_config.output_hidden_states == True, "The upstream model must return all hidden states"
+
+        self.model_config = model_config
+        self.use_conv_output = use_conv_output
+
+
+        self.model_seq = nn.Sequential(
+            nn.Conv1d(self.model_config.hidden_size, self.model_config.classifier_proj_size, 1, padding=0),
+            nn.ReLU(),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(self.model_config.classifier_proj_size, self.model_config.classifier_proj_size, 1, padding=0),
+            nn.ReLU(),
+            nn.Dropout(p=0.1),
+            nn.Conv1d(self.model_config.classifier_proj_size, self.model_config.classifier_proj_size, 1, padding=0)
+        )
+
+        if self.use_conv_output:
+            num_layers = self.model_config.num_hidden_layers + 1  # transformer layers + input embeddings
+            self.weights = nn.Parameter(torch.ones(num_layers)/num_layers)
+        else:
+            num_layers = self.model_config.num_hidden_layers
+            self.weights = nn.Parameter(torch.zeros(num_layers))
+        
+        self.out_layer = nn.Sequential(
+            nn.Linear(self.model_config.classifier_proj_size, self.model_config.classifier_proj_size),
+            nn.ReLU(),
+            nn.Linear(self.model_config.classifier_proj_size, self.model_config.num_labels),
+        )
+
+    def forward(self, encoder_hidden_states, length=None):
+        if self.use_conv_output:
+            stacked_feature = torch.stack(encoder_hidden_states, dim=0)
+        else:
+            stacked_feature = torch.stack(encoder_hidden_states, dim=0)[1:] # exclude the convolution output
+        
+        _, *origin_shape = stacked_feature.shape
+        
+        if self.use_conv_output:
+            stacked_feature = stacked_feature.view(self.model_config.num_hidden_layers + 1, -1)
+        else:
+            stacked_feature = stacked_feature.view(self.model_config.config.num_hidden_layers, -1)
+        
+        norm_weights = F.softmax(self.weights, dim=-1)
+        
+        
+        weighted_feature = (norm_weights.unsqueeze(-1) * stacked_feature).sum(dim=0)
+        features = weighted_feature.view(*origin_shape)
+        
+        features = features.transpose(1, 2)
+        features = self.model_seq(features)
+        features = features.transpose(1, 2)
+        
+        if length is not None:
+            length = length.cuda()
+            masks = torch.arange(features.size(1)).expand(length.size(0), -1).cuda() < length.unsqueeze(1)
+            masks = masks.float()
+            features = (features * masks.unsqueeze(-1)).sum(1) / length.unsqueeze(1)
+        else:
+            features = torch.mean(features, dim=1)
+        
+        predicted = self.out_layer(features)
+        return predicted 

src/model/wav2vec2_wrapper.py

@@ -0,0 +1,93 @@
+import torch
+import transformers.models.wav2vec2.modeling_wav2vec2 as w2v2
+from torch import nn
+from transformers import Wav2Vec2Model
+
+class Wav2Vec2EncoderLayer(nn.Module):
+    def __init__(
+        self,
+        config,
+        i
+    ):
+        super().__init__()
+        self.attention = w2v2.Wav2Vec2Attention(
+            embed_dim=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            dropout=config.attention_dropout,
+            is_decoder=False,
+        )
+        self.dropout = nn.Dropout(config.hidden_dropout)
+        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.feed_forward = w2v2.Wav2Vec2FeedForward(config)
+        self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.config = config
+        self.i = i
+
+    def forward(self, hidden_states, attention_mask=None, output_attentions=False):
+        attn_residual = hidden_states
+        
+        hidden_states, attn_weights, _ = self.attention(
+            hidden_states, attention_mask=attention_mask, output_attentions=output_attentions
+        )
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = attn_residual + hidden_states
+        hidden_states = self.layer_norm(hidden_states)
+        hidden_states = hidden_states + self.feed_forward(hidden_states)            
+        hidden_states = self.final_layer_norm(hidden_states)
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+        return outputs
+
+class Wav2VecWrapper(nn.Module):
+    def __init__(
+        self,
+        config,
+    ):
+        super(Wav2VecWrapper, self).__init__()
+        self.config = config
+
+        self.backbone_model = Wav2Vec2Model.from_pretrained(
+            config._name_or_path,
+            output_hidden_states=config.output_hidden_states,
+        )
+        state_dict = self.backbone_model.state_dict()
+
+        self.model_config = self.backbone_model.config
+        self.backbone_model.encoder.layers = nn.ModuleList([Wav2Vec2EncoderLayer(self.model_config, i) for i in range(self.model_config.num_hidden_layers)])
+        
+    def forward(self,
+                input_features: torch.Tensor,
+                length: torch.Tensor = None,
+            ):
+        with torch.no_grad():
+            hidden_states = self.backbone_model.feature_extractor(input_features)
+            hidden_states = hidden_states.transpose(1, 2)
+            hidden_states, _ = self.backbone_model.feature_projection(hidden_states) 
+        
+        if length is not None:
+            length = self.get_feat_extract_output_lengths(length.detach().cpu())
+
+        hidden_states = self.backbone_model.encoder(
+            hidden_states,
+            output_hidden_states=self.config.output_hidden_states
+        ).hidden_states
+        
+        return {'encoder_hidden_states': hidden_states, 'length': length}
+        
+    def get_feat_extract_output_lengths(self, input_length):
+        def _conv_out_length(input_length, kernel_size, stride):
+            return (input_length - kernel_size) // stride + 1
+        for kernel_size, stride in zip(self.backbone_model.config.conv_kernel, self.backbone_model.config.conv_stride):
+            input_length = _conv_out_length(input_length, kernel_size, stride)
+        return input_length
+
+def prepare_mask(length, shape, dtype):
+    mask = torch.zeros(
+        shape, dtype=dtype
+    )
+    mask[(torch.arange(mask.shape[0]), length.cpu() - 1)] = 1
+    mask = mask.flip([-1]).cumsum(-1).flip([-1]).bool()
+    return mask