eval.py

import numpy as np
import os
import pandas as pd
import torch
from sklearn.metrics import roc_curve, auc
from scipy import stats
from config import REACTION_LABELS
from data_parser import segment_stress

def calc_ccc(preds, labels):
    '''
    Concordance Correlation Coefficient
    :param preds: 1D np array
    :param labels: 1D np array
    :return:
    '''

    preds_mean, labels_mean = np.mean(preds), np.mean(labels)
    cov_mat = np.cov(preds, labels)
    covariance = cov_mat[0, 1]
    preds_var, labels_var = cov_mat[0, 0], cov_mat[1, 1]

    ccc = 2.0 * covariance / (preds_var + labels_var + (preds_mean - labels_mean) ** 2)
    return ccc

def calc_ccc_stress(preds, labels):
    '''
    Concordance Correlation Coefficient
    :param preds: 1D np array
    :param labels: 1D np array
    :return:
    '''

    preds_mean, labels_mean = np.mean(preds), np.mean(labels)
    rmse = np.sqrt(np.mean(preds - labels) ** 2)
    cov_mat = np.cov(preds, labels)
    covariance = cov_mat[0, 1]
    preds_var, labels_var = cov_mat[0, 0], cov_mat[1, 1]
    pcc = covariance / np.sqrt(preds_var * labels_var)
    ccc = 2.0 * covariance / (preds_var + labels_var + (preds_mean - labels_mean) ** 2)
    return rmse, pcc, ccc

def mean_ccc(preds, labels):
    '''

    :param preds: list of list of lists (num batches, batch_size, num_classes)
    :param labels: same
    :return: scalar
    '''
    preds = np.row_stack([np.array(p) for p in preds])
    labels = np.row_stack([np.array(l) for l in labels])
    num_classes = preds.shape[1]
    class_wise_cccs = np.array([calc_ccc(preds[:,i], labels[:,i]) for i in range(num_classes)])
    mean_ccc = np.mean(class_wise_cccs)
    return mean_ccc

def calc_pearsons(preds,labels):
    r = stats.pearsonr(preds, labels)
    return r[0]

def mean_pearsons(preds,labels):
    preds = np.row_stack([np.array(p) for p in preds])
    labels = np.row_stack([np.array(l) for l in labels])
    num_classes = preds.shape[1]
    class_wise_r = np.array([calc_pearsons(preds[:,i], labels[:,i]) for i in range(num_classes)])
    mean_r = np.mean(class_wise_r)
    return mean_r


    
def calc_auc(preds, labels):
    preds = np.concatenate(preds)
    labels = np.concatenate(labels)

    fpr, tpr, thresholds = roc_curve(labels, preds)
    return auc(fpr, tpr)

def write_reaction_predictions(full_metas, full_preds, csv_dir, filename):
    meta_arr = np.row_stack(full_metas).squeeze()
    preds_arr = np.row_stack(full_preds)
    pred_df = pd.DataFrame(columns=['File_ID']+REACTION_LABELS)
    pred_df['File_ID'] = meta_arr
    pred_df[REACTION_LABELS] = preds_arr
    pred_df.to_csv(os.path.join(csv_dir, filename), index=False)
    return None


def write_predictions(task, full_metas, full_preds, prediction_path, filename):
    assert prediction_path != ''

    if not os.path.exists(prediction_path):
        os.makedirs(prediction_path)

    if task == 'reaction':
        return write_reaction_predictions(full_metas, full_preds, prediction_path, filename)

    metas_flat = []
    for meta in full_metas:
        metas_flat.extend(meta)
    preds_flat = []
    for pred in full_preds:
        preds_flat.extend(pred if isinstance(pred, list) else (pred.squeeze() if pred.ndim>1 else pred))

    if isinstance(metas_flat[0],list):
        num_meta_cols = len(metas_flat[0])
    else:
        # np array
        num_meta_cols = metas_flat[0].shape[0]
    prediction_df = pd.DataFrame(columns = [f'meta_col_{i}' for i in range(num_meta_cols)])
    for i in range(num_meta_cols - 1):
        prediction_df[f'meta_col_{i}'] = [m[i] for m in metas_flat]
    prediction_df[f'meta_col_{num_meta_cols - 1}'] = [m[0] for m in metas_flat]
    prediction_df['prediction'] = preds_flat
    #prediction_df['label'] = labels_flat
    prediction_df.to_csv(os.path.join(prediction_path, filename), index=False)


def evaluate(task, model, data_loader, loss_fn, eval_fn, use_gpu=False, predict=False, prediction_path=None, prediction_split_path=None, filename=None, batch_size=32, win_len=20):
    losses, sizes = 0, 0
    full_preds_valence = []
    full_preds_arousal = []
    full_labels_valence = []
    full_labels_arousal = []
    if predict:
        full_metas = []
    else:
        full_metas = None

    model.eval()
    with torch.no_grad():
        for batch, batch_data in enumerate(data_loader, 1):
            features, feature_lens, labels_valence,labels_arousal, metas = batch_data
            _, seq_len, features_size = features.size()

            features = pd.DataFrame(features.view(seq_len, features_size))
            stress_segmentato = segment_stress(features, win_len, 1)
            tensor_list = [torch.tensor(df.values) for df in stress_segmentato]
            stacked_tensor = torch.stack(tensor_list, dim=0)

            labels_valence = pd.DataFrame(labels_valence.view(seq_len, 1))
            labels_segmentate_valence = segment_stress(labels_valence, win_len, 1)
            labels_list_valence = [torch.tensor(df.values) for df in labels_segmentate_valence]
            last_elements_valence = [tensor[-1] for tensor in labels_list_valence]
            stacked_labels_valence = torch.stack(last_elements_valence, dim=0)

            labels_arousal = pd.DataFrame(labels_arousal.view(seq_len, 1))
            labels_segmentate_arousal = segment_stress(labels_arousal, win_len, 1)
            labels_list_arousal = [torch.tensor(df.values) for df in labels_segmentate_arousal]
            last_elements_arousal = [tensor[-1] for tensor in labels_list_arousal]
            stacked_labels_arousal = torch.stack(last_elements_arousal, dim=0)

            batches = torch.split(stacked_tensor, batch_size)
            labels_batches_valence = torch.split(stacked_labels_valence,batch_size)
            labels_batches_arousal = torch.split(stacked_labels_arousal,batch_size)

            for batch, labels_batch_valence, labels_batch_arousal in zip(batches, labels_batches_valence, labels_batches_arousal):
                bs2, sl, fs = batch.size()
                if bs2 == 1:
                    continue
                if use_gpu:
                    model.cuda()
                    batch = batch.cuda()
                    labels_batch_valence = labels_batch_valence.cuda()
                    labels_batch_arousal = labels_batch_arousal.cuda()

                preds_valence, preds_arousal = model(batch, torch.tensor(win_len))

                loss_valence = loss_fn(preds_valence.squeeze(-1), labels_batch_valence.squeeze(-1))
                loss_arousal = loss_fn(preds_arousal.squeeze(-1), labels_batch_arousal.squeeze(-1))
                loss = loss_valence + loss_arousal
                losses += loss.item() * batch_size
                sizes += batch_size

                full_labels_valence.append(labels_batch_valence.cpu().detach().squeeze().numpy().tolist())
                full_labels_arousal.append(labels_batch_arousal.cpu().detach().squeeze().numpy().tolist())
                full_preds_valence.append(preds_valence.cpu().detach().squeeze().numpy().tolist())
                full_preds_arousal.append(preds_arousal.cpu().detach().squeeze().numpy().tolist())

        if task == 'stress':
            full_labels_valence = flatten_stress_for_ccc(full_labels_valence)
            full_labels_arousal = flatten_stress_for_ccc(full_labels_arousal)
            full_preds_valence = flatten_stress_for_ccc(full_preds_valence)
            full_preds_arousal = flatten_stress_for_ccc(full_preds_arousal)
        rmse_valence, pcc_valence, score_valence = eval_fn(full_preds_valence, full_labels_valence)
        rmse_arousal, pcc_arousal, score_arousal = eval_fn(full_preds_arousal, full_labels_arousal)
        total_loss = losses / sizes
        return total_loss, rmse_valence, pcc_valence, score_valence, rmse_arousal, pcc_arousal, score_arousal


def flatten_stress_for_ccc(lst):
    '''
    Brings full_preds and full_labels of stress into the right format for the CCC function
    :param lst: list of lists of different lengths
    :return: flattened numpy array
    '''
    return np.concatenate([np.array(l) for l in lst])