<reponame>msanchezgonzalez/Sentiment-Analysis-On-Voice-Data<filename>run_simple_cnn.py
from speechemotionrecognition.dnn import CNN
from speechemotionrecognition.utilities import get_feature_vector_from_mfcc

import os

def convert_file(old_filepath : str) -> str:
    '''
    Converts an audio file to a mono .wav WAVE format, sample rate of 16000 Hz
    and overwrites if duplicate. Depends on having ffmpeg installed.
    '''
    new_filepath = os.path.splitext(old_filepath) + '.wav'
    os.system('ffmpeg -i {} -ar 16000 -ac 1 -c:a pcm_s16le -y {}'.format(
        old_filepath, new_filepath
    ))
    return new_filepath

def parse_labels(index : int) -> str:
    labels = ["Neutral", "Angry", "Happy", "Sad"]
    return labels[index]

def cnn_predict_example(audio_filepath) -> str:
    '''
    Simple function to test the trained CNN.

    Arguments
    ---------
        - audio_filepath: the filepath with the audio to analyze. It has to be 
        mono and have a sample rate of 16000 Hz.

    Returns
    -------
        - prediction: the prediction from the CNN (one from Neutral, Angry, 
        Happy and Sad)
    '''
    # Don't flatten the audio file ndarrays
    to_flatten = False
    # Define the model class with the specific parameters for our problem
    model = CNN(input_shape=(198, 39), # x_train[0].shape = (198, 39)
                num_classes=4) # num_labels = 4
    # Restore the model from the one that it's previously trained
    model.restore_model('./models/best_model_CNN.h5')
    prediction = model.predict_prob(
        get_feature_vector_from_mfcc(audio_filepath, flatten=to_flatten)
    )
    pred_class_names = ["Neutral", "Angry", "Happy", "Sad"]
    prediction_array = {pred_class_names[i]: prediction.tolist()[i] for i in range(len(prediction))}
    return prediction_array

def cnn_predict_example_from_mic(audio_filepath) -> str:
    '''
    Simple function to test the trained CNN.

    Arguments
    ---------
        - audio_filepath: the filepath with the audio to analyze. It has to be 
        mono and have a sample rate of 16000 Hz.

    Returns
    -------
        - prediction: the prediction from the CNN (one from Neutral, Angry, 
        Happy and Sad)
    '''
    # Don't flatten the audio file ndarrays
    to_flatten = False
    # Define the model class with the specific parameters for our problem
    model = CNN(input_shape=(198, 39), # x_train[0].shape = (198, 39)
                num_classes=4) # num_labels = 4
    # Restore the model from the one that it's previously trained
    model.restore_model('./models/best_model_CNN.h5')
    # Audio conversion
    new_audio_filepath = audio_filepath[:-4] + '_conv.wav'
    os.system('ffmpeg -i {} -c:a pcm_f32le -ar 16000 {}'.format(audio_filepath, new_audio_filepath))
    print('new_audio_filepath =', new_audio_filepath)
    prediction = model.predict_prob(
        get_feature_vector_from_mfcc(new_audio_filepath, flatten=to_flatten)
    )[0] # because this returns an array within a list
    print(prediction)
    print('prediction has been made')
    pred_class_names = ["Neutral", "Angry", "Happy", "Sad"]
    prediction_array = {pred_class_names[i]: prediction.tolist()[i] for i in range(len(prediction))}
    print(prediction_array)
    return prediction_array


if __name__ == '__main__':
    audio_filepath = './mariaenfadada.wav'
    prediction = cnn_predict_example_from_mic(audio_filepath)
    print('prediction =', prediction)
