model_regression.py

import logging
import time
import os
import pandas as pd
import numpy as np
import math
import scipy.io
import scipy.stats
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from sklearn.metrics import mean_squared_error
from scipy.optimize import curve_fit
import joblib

import seaborn as sns
import matplotlib.pyplot as plt
import copy
import argparse

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.optim.swa_utils import AveragedModel, SWALR
from torch.utils.data import DataLoader, TensorDataset
from sklearn.model_selection import KFold
from sklearn.model_selection import train_test_split

from data_processing import split_train_test

# ignore all warnings
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)


class Mlp(nn.Module):
    def __init__(self, input_features, hidden_features=256, out_features=1, drop_rate=0.2, act_layer=nn.GELU):
        super().__init__()
        self.fc1 = nn.Linear(input_features, hidden_features)
        self.bn1 = nn.BatchNorm1d(hidden_features)
        self.act1 = act_layer()
        self.drop1 = nn.Dropout(drop_rate)
        self.fc2 = nn.Linear(hidden_features, hidden_features // 2)
        self.act2 = act_layer()
        self.drop2 = nn.Dropout(drop_rate)
        self.fc3 = nn.Linear(hidden_features // 2, out_features)

    def forward(self, input_feature):
        x = self.fc1(input_feature)
        x = self.bn1(x)
        x = self.act1(x)
        x = self.drop1(x)
        x = self.fc2(x)
        x = self.act2(x)
        x = self.drop2(x)
        output = self.fc3(x)
        return output


class MAEAndRankLoss(nn.Module):
    def __init__(self, l1_w=1.0, rank_w=1.0, margin=0.0, use_margin=False):
        super(MAEAndRankLoss, self).__init__()
        self.l1_w = l1_w
        self.rank_w = rank_w
        self.margin = margin
        self.use_margin = use_margin

    def forward(self, y_pred, y_true):
        # L1 loss/MAE loss
        l_mae = F.l1_loss(y_pred, y_true, reduction='mean') * self.l1_w
        # Rank loss
        n = y_pred.size(0)
        pred_diff = y_pred.unsqueeze(1) - y_pred.unsqueeze(0)
        true_diff = y_true.unsqueeze(1) - y_true.unsqueeze(0)

        # e(ytrue_i, ytrue_j)
        masks = torch.sign(true_diff)

        if self.use_margin and self.margin > 0:
            true_diff = true_diff.abs() - self.margin
            true_diff = F.relu(true_diff)
            masks = true_diff.sign()

        l_rank = F.relu(true_diff - masks * pred_diff)
        l_rank = l_rank.sum() / (n * (n - 1))

        loss = l_mae + l_rank * self.rank_w
        return loss

def load_data(csv_file, mat_file, features, data_name, set_name):
    try:
        df = pd.read_csv(csv_file, skiprows=[], header=None)
    except Exception as e:
        logging.error(f'Read CSV file error: {e}')
        raise

    try:
        if data_name == 'lsvq_train':
            X_mat = features
        else:
            X_mat = scipy.io.loadmat(mat_file)
    except Exception as e:
        logging.error(f'Read MAT file error: {e}')
        raise

    y_data = df.values[1:, 2]
    y = np.array(list(y_data), dtype=float)

    if data_name == 'cross_dataset': # or data_name == 'lsvq_train':
        y[y > 5] = 5
    if set_name == 'test':
        print(f"Modified y_true: {y}")
    if data_name == 'lsvq_train':
        X = np.asarray(X_mat, dtype=float)
    else:
        data_name = f'{data_name}_{set_name}_features'
        X = np.asarray(X_mat[data_name], dtype=float)

    return X, y

def preprocess_data(X, y):
    X[np.isnan(X)] = 0
    X[np.isinf(X)] = 0
    imp = SimpleImputer(missing_values=np.nan, strategy='mean').fit(X)
    X = imp.transform(X)

    # scaler = StandardScaler()
    scaler = MinMaxScaler().fit(X)
    X = scaler.transform(X)
    logging.info(f'Scaler: {scaler}')

    y = y.reshape(-1, 1).squeeze()
    return X, y, imp, scaler

# define 4-parameter logistic regression
def logistic_func(X, bayta1, bayta2, bayta3, bayta4):
    logisticPart = 1 + np.exp(np.negative(np.divide(X - bayta3, np.abs(bayta4))))
    yhat = bayta2 + np.divide(bayta1 - bayta2, logisticPart)
    return yhat

def fit_logistic_regression(y_pred, y_true):
    beta = [np.max(y_true), np.min(y_true), np.mean(y_pred), 0.5]
    popt, _ = curve_fit(logistic_func, y_pred, y_true, p0=beta, maxfev=100000000)
    y_pred_logistic = logistic_func(y_pred, *popt)
    return y_pred_logistic, beta, popt

def compute_correlation_metrics(y_true, y_pred):
    y_pred_logistic, beta, popt = fit_logistic_regression(y_pred, y_true)

    plcc = scipy.stats.pearsonr(y_true, y_pred_logistic)[0]
    rmse = np.sqrt(mean_squared_error(y_true, y_pred_logistic))
    srcc = scipy.stats.spearmanr(y_true, y_pred)[0]

    try:
        krcc = scipy.stats.kendalltau(y_true, y_pred)[0]
    except Exception as e:
        logging.error(f'krcc calculation: {e}')
        krcc = scipy.stats.kendalltau(y_true, y_pred, method='asymptotic')[0]
    return y_pred_logistic, plcc, rmse, srcc, krcc

def plot_results(y_test, y_test_pred_logistic, df_pred_score, model_name, data_name, network_name, select_criteria):
    # nonlinear logistic fitted curve / logistic regression
    mos1 = y_test
    y1 = y_test_pred_logistic

    try:
        beta = [np.max(mos1), np.min(mos1), np.mean(y1), 0.5]
        popt, pcov = curve_fit(logistic_func, y1, mos1, p0=beta, maxfev=100000000)
        sigma = np.sqrt(np.diag(pcov))
    except:
        raise Exception('Fitting logistic function time-out!!')
    x_values1 = np.linspace(np.min(y1), np.max(y1), len(y1))
    plt.plot(x_values1, logistic_func(x_values1, *popt), '-', color='#c72e29', label='Fitted f(x)')

    fig1 = sns.scatterplot(x="y_test_pred_logistic", y="MOS", data=df_pred_score, markers='o', color='steelblue', label=network_name)
    plt.legend(loc='upper left')
    if data_name == 'live_vqc' or data_name == 'live_qualcomm' or data_name == 'cvd_2014' or data_name == 'lsvq_train':
        plt.ylim(0, 100)
        plt.xlim(0, 100)
    else:
        plt.ylim(1, 5)
        plt.xlim(1, 5)
    plt.title(f"Algorithm {network_name} with {model_name} on dataset {data_name}", fontsize=10)
    plt.xlabel('Predicted Score')
    plt.ylabel('MOS')
    reg_fig1 = fig1.get_figure()

    fig_path = f'../figs/{data_name}/'
    os.makedirs(fig_path, exist_ok=True)
    reg_fig1.savefig(fig_path + f"{network_name}_{model_name}_{data_name}_by{select_criteria}_kfold.png", dpi=300)
    plt.clf()
    plt.close()

def plot_and_save_losses(avg_train_losses, avg_val_losses, model_name, data_name, network_name, test_vids, i):
    plt.figure(figsize=(10, 6))

    plt.plot(avg_train_losses, label='Average Training Loss')
    plt.plot(avg_val_losses, label='Average Validation Loss')

    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title(f'Average Training and Validation Loss Across Folds - {network_name} with {model_name} (test_vids: {test_vids})', fontsize=10)

    plt.legend()
    fig_par_path = f'../log/result/{data_name}/'
    os.makedirs(fig_par_path, exist_ok=True)
    plt.savefig(f'{fig_par_path}/{network_name}_Average_Training_Loss_test{i}.png', dpi=50)
    plt.clf()
    plt.close()

def configure_logging(log_path, model_name, data_name, network_name, select_criteria):
    log_file_name = os.path.join(log_path, f"{data_name}_{network_name}_{model_name}_corr_{select_criteria}_kfold.log")
    logging.basicConfig(filename=log_file_name, filemode='w', level=logging.DEBUG, format='%(levelname)s - %(message)s')
    logging.getLogger('matplotlib').setLevel(logging.WARNING)
    logging.info(f"Evaluating algorithm {network_name} with {model_name} on dataset {data_name}")
    logging.info(f"torch cuda: {torch.cuda.is_available()}")

def load_and_preprocess_data(metadata_path, feature_path, data_name, network_name, train_features, test_features):
    if data_name == 'cross_dataset':
        data_name1 = 'youtube_ugc_all'
        data_name2 = 'cvd_2014_all'
        csv_train_file = os.path.join(metadata_path, f'mos_files/{data_name1}_MOS_train.csv')
        csv_test_file = os.path.join(metadata_path, f'mos_files/{data_name2}_MOS_test.csv')
        mat_train_file = os.path.join(f'{feature_path}split_train_test/', f'{data_name1}_{network_name}_train_features.mat')
        mat_test_file = os.path.join(f'{feature_path}split_train_test/', f'{data_name2}_{network_name}_test_features.mat')
        X_train, y_train = load_data(csv_train_file, mat_train_file, None, data_name1, 'train')
        X_test, y_test = load_data(csv_test_file, mat_test_file, None, data_name2, 'test')

    elif data_name == 'lsvq_train':
        csv_train_file = os.path.join(metadata_path, f'mos_files/{data_name}_MOS_train.csv')
        csv_test_file = os.path.join(metadata_path, f'mos_files/{data_name}_MOS_test.csv')
        X_train, y_train = load_data(csv_train_file, None, train_features, data_name, 'train')
        X_test, y_test = load_data(csv_test_file, None, test_features, data_name, 'test')

    else:
        csv_train_file = os.path.join(metadata_path, f'mos_files/{data_name}_MOS_train.csv')
        csv_test_file = os.path.join(metadata_path, f'mos_files/{data_name}_MOS_test.csv')
        mat_train_file = os.path.join(f'{feature_path}split_train_test/', f'{data_name}_{network_name}_train_features.mat')
        mat_test_file = os.path.join(f'{feature_path}split_train_test/', f'{data_name}_{network_name}_test_features.mat')
        X_train, y_train = load_data(csv_train_file, mat_train_file, None, data_name, 'train')
        X_test, y_test = load_data(csv_test_file, mat_test_file, None, data_name, 'test')

    # standard min-max normalization of traning features
    X_train, y_train, _, _ = preprocess_data(X_train, y_train)
    X_test, y_test, _, _ = preprocess_data(X_test, y_test)

    return X_train, y_train, X_test, y_test

def train_one_epoch(model, train_loader, criterion, optimizer, device):
    """Train the model for one epoch"""
    model.train()
    train_loss = 0.0
    for inputs, targets in train_loader:
        inputs, targets = inputs.to(device), targets.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets.view(-1, 1))
        loss.backward()
        optimizer.step()
        train_loss += loss.item() * inputs.size(0)
    train_loss /= len(train_loader.dataset)
    return train_loss

def evaluate(model, val_loader, criterion, device):
    """Evaluate model performance on validation sets"""
    model.eval()
    val_loss = 0.0
    y_val_pred = []
    with torch.no_grad():
        for inputs, targets in val_loader:
            inputs, targets = inputs.to(device), targets.to(device)

            outputs = model(inputs)
            y_val_pred.extend(outputs.view(-1).tolist())
            loss = criterion(outputs, targets.view(-1, 1))
            val_loss += loss.item() * inputs.size(0)
    val_loss /= len(val_loader.dataset)
    return val_loss, np.array(y_val_pred)

def update_best_model(select_criteria, best_metric, current_val, model):
    is_better = False
    if select_criteria == 'byrmse' and current_val < best_metric:
        is_better = True
    elif select_criteria == 'bykrcc' and current_val > best_metric:
        is_better = True

    if is_better:
        return current_val, copy.deepcopy(model), is_better
    return best_metric, model, is_better

def train_and_evaluate(X_train, y_train, config):
    # parameters
    n_repeats = config['n_repeats']
    n_splits = config['n_splits']
    batch_size = config['batch_size']
    epochs = config['epochs']
    hidden_features = config['hidden_features']
    drop_rate = config['drop_rate']
    loss_type = config['loss_type']
    optimizer_type = config['optimizer_type']
    select_criteria = config['select_criteria']
    initial_lr = config['initial_lr']
    weight_decay = config['weight_decay']
    patience = config['patience']
    l1_w = config['l1_w']
    rank_w = config['rank_w']
    use_swa = config.get('use_swa', False)
    logging.info(f'Parameters - Number of repeats for 80-20 hold out test: {n_repeats}, Number of splits for kfold: {n_splits}, Batch size: {batch_size}, Number of epochs: {epochs}')
    logging.info(f'Network Parameters - hidden_features: {hidden_features}, drop_rate: {drop_rate}, patience: {patience}')
    logging.info(f'Optimizer Parameters - loss_type: {loss_type}, optimizer_type: {optimizer_type}, initial_lr: {initial_lr}, weight_decay: {weight_decay}, use_swa: {use_swa}')
    logging.info(f'MAEAndRankLoss - l1_w: {l1_w}, rank_w: {rank_w}')

    kf = KFold(n_splits=n_splits, shuffle=True, random_state=42)
    best_model = None
    best_metric = float('inf') if select_criteria == 'byrmse' else float('-inf')

    # loss for every fold
    all_train_losses = []
    all_val_losses = []
    for fold, (train_idx, val_idx) in enumerate(kf.split(X_train)):
        print(f"Fold {fold + 1}/{n_splits}")

        X_train_fold, X_val_fold = X_train[train_idx], X_train[val_idx]
        y_train_fold, y_val_fold = y_train[train_idx], y_train[val_idx]

        # initialisation of model, loss function, optimiser
        model = Mlp(input_features=X_train_fold.shape[1], hidden_features=hidden_features, drop_rate=drop_rate)
        model = model.to(device) # to gpu

        if loss_type == 'MAERankLoss':
            criterion = MAEAndRankLoss()
            criterion.l1_w = l1_w
            criterion.rank_w = rank_w
        else:
            nn.MSELoss()

        if optimizer_type == 'sgd':
            optimizer = optim.SGD(model.parameters(), lr=initial_lr, momentum=0.9, weight_decay=weight_decay)
            scheduler = CosineAnnealingLR(optimizer, T_max=epochs, eta_min=1e-5)# initial eta_nim=1e-5
        else:
            optimizer = optim.Adam(model.parameters(), lr=initial_lr, weight_decay=weight_decay)  # L2 Regularisation initial: 0.01, 1e-5
            scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.95)  # step_size=10, gamma=0.1: every 10 epochs lr*0.1
        if use_swa:
            swa_model = AveragedModel(model).to(device)
            swa_scheduler = SWALR(optimizer, swa_lr=initial_lr, anneal_strategy='cos')

        # dataset loader
        train_dataset = TensorDataset(torch.FloatTensor(X_train_fold), torch.FloatTensor(y_train_fold))
        val_dataset = TensorDataset(torch.FloatTensor(X_val_fold), torch.FloatTensor(y_val_fold))
        train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
        val_loader = DataLoader(dataset=val_dataset, batch_size=batch_size, shuffle=False)

        train_losses, val_losses = [], []

        # early stopping parameters
        best_val_loss = float('inf')
        epochs_no_improve = 0
        early_stop_active = False
        swa_start = int(epochs * 0.7) if use_swa else epochs  # SWA starts after 70% of total epochs, only set SWA start if SWA is used

        for epoch in range(epochs):
            train_loss = train_one_epoch(model, train_loader, criterion, optimizer, device)
            train_losses.append(train_loss)
            scheduler.step() # update learning rate
            if use_swa and epoch >= swa_start:
                swa_model.update_parameters(model)
                swa_scheduler.step()
                early_stop_active = True
                print(f"Current learning rate with SWA: {swa_scheduler.get_last_lr()}")

            lr = optimizer.param_groups[0]['lr']
            print('Epoch %d: Learning rate: %f' % (epoch + 1, lr))

            # decide which model to evaluate: SWA model or regular model
            current_model = swa_model if use_swa and epoch >= swa_start else model
            current_model.eval()
            val_loss, y_val_pred = evaluate(current_model, val_loader, criterion, device)
            val_losses.append(val_loss)
            print(f"Epoch {epoch + 1}, Fold {fold + 1}, Training Loss: {train_loss}, Validation Loss: {val_loss}")

            y_val_pred = np.array(list(y_val_pred), dtype=float)
            _, _, rmse_val, _, krcc_val = compute_correlation_metrics(y_val_fold, y_val_pred)
            current_metric = rmse_val if select_criteria == 'byrmse' else krcc_val
            best_metric, best_model, is_better = update_best_model(select_criteria, best_metric, current_metric, current_model)
            if is_better:
                logging.info(f"Epoch {epoch + 1}, Fold {fold + 1}:")
                y_val_pred_logistic_tmp, plcc_valid_tmp, rmse_valid_tmp, srcc_valid_tmp, krcc_valid_tmp = compute_correlation_metrics(y_val_fold, y_val_pred)
                logging.info(f'Validation set - Evaluation Results - SRCC: {srcc_valid_tmp}, KRCC: {krcc_valid_tmp}, PLCC: {plcc_valid_tmp}, RMSE: {rmse_valid_tmp}')

                X_train_fold_tensor = torch.FloatTensor(X_train_fold).to(device)
                y_tra_pred_tmp = best_model(X_train_fold_tensor).detach().cpu().numpy().squeeze()
                y_tra_pred_tmp = np.array(list(y_tra_pred_tmp), dtype=float)
                y_tra_pred_logistic_tmp, plcc_train_tmp, rmse_train_tmp, srcc_train_tmp, krcc_train_tmp = compute_correlation_metrics(y_train_fold, y_tra_pred_tmp)
                logging.info(f'Train set - Evaluation Results - SRCC: {srcc_train_tmp}, KRCC: {krcc_train_tmp}, PLCC: {plcc_train_tmp}, RMSE: {rmse_train_tmp}')

            # check for loss improvement
            if early_stop_active:
                if val_loss < best_val_loss:
                    best_val_loss = val_loss
                    # save the best model if validation loss improves
                    best_model = copy.deepcopy(model)
                    epochs_no_improve = 0
                else:
                    epochs_no_improve += 1
                    if epochs_no_improve >= patience:
                        # epochs to wait for improvement before stopping
                        print(f"Early stopping triggered after {epoch + 1} epochs.")
                        break

        # saving SWA models and updating BN statistics
        if use_swa:
            train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, collate_fn=lambda x: collate_to_device(x, device))
            best_model = best_model.to(device)
            best_model.eval()
            torch.optim.swa_utils.update_bn(train_loader, best_model)
            # swa_model_path = os.path.join('save_swa_path='../model/', f'model_swa_fold{fold}.pth')
            # torch.save(swa_model.state_dict(), swa_model_path)
            # logging.info(f'SWA model saved at {swa_model_path}')

        all_train_losses.append(train_losses)
        all_val_losses.append(val_losses)
        max_length = max(len(x) for x in all_train_losses)
        all_train_losses = [x + [x[-1]] * (max_length - len(x)) for x in all_train_losses]
        max_length = max(len(x) for x in all_val_losses)
        all_val_losses = [x + [x[-1]] * (max_length - len(x)) for x in all_val_losses]

    return best_model, all_train_losses, all_val_losses

def collate_to_device(batch, device):
    data, targets = zip(*batch)
    return torch.stack(data).to(device), torch.stack(targets).to(device)

def model_test(best_model, X, y, device):
    test_dataset = TensorDataset(torch.FloatTensor(X), torch.FloatTensor(y))
    test_loader = DataLoader(dataset=test_dataset, batch_size=1, shuffle=False)

    best_model.eval()
    y_pred = []
    with torch.no_grad():
        for inputs, _ in test_loader:
            inputs = inputs.to(device)

            outputs = best_model(inputs)
            y_pred.extend(outputs.view(-1).tolist())

    return y_pred

def main(config):
    model_name = config['model_name']
    data_name = config['data_name']
    network_name = config['network_name']

    metadata_path = config['metadata_path']
    feature_path = config['feature_path']
    log_path = config['log_path']
    save_path = config['save_path']
    score_path = config['score_path']
    result_path = config['result_path']

    # parameters
    select_criteria = config['select_criteria']
    n_repeats = config['n_repeats']

    # logging and result
    os.makedirs(log_path, exist_ok=True)
    os.makedirs(save_path, exist_ok=True)
    os.makedirs(score_path, exist_ok=True)
    os.makedirs(result_path, exist_ok=True)
    result_file = f'{result_path}{data_name}_{network_name}_{select_criteria}.mat'
    pred_score_filename = os.path.join(score_path, f"{data_name}_{network_name}_{select_criteria}.csv")
    file_path = os.path.join(save_path, f"{data_name}_{network_name}_{select_criteria}_trained_median_model_param.pth")
    configure_logging(log_path, model_name, data_name, network_name, select_criteria)

    '''======================== Main Body ==========================='''
    PLCC_all_repeats_test = []
    SRCC_all_repeats_test = []
    KRCC_all_repeats_test = []
    RMSE_all_repeats_test = []
    PLCC_all_repeats_train = []
    SRCC_all_repeats_train = []
    KRCC_all_repeats_train = []
    RMSE_all_repeats_train = []
    all_repeats_test_vids = []
    all_repeats_df_test_pred = []
    best_model_list = []

    for i in range(1, n_repeats + 1):
        print(f"{i}th repeated 80-20 hold out test")
        logging.info(f"{i}th repeated 80-20 hold out test")
        t0 = time.time()

        # train test split
        test_size = 0.2
        random_state = math.ceil(8.8 * i)
        # NR: original
        if data_name == 'lsvq_train':
            test_data_name = 'lsvq_test' #lsvq_test, lsvq_test_1080p
            train_features, test_features, test_vids = split_train_test.process_lsvq(data_name, test_data_name, metadata_path, feature_path, network_name)
        elif data_name == 'cross_dataset':
            train_data_name = 'youtube_ugc_all'
            test_data_name = 'cvd_2014_all'
            _, _, test_vids = split_train_test.process_cross_dataset(train_data_name, test_data_name, metadata_path, feature_path, network_name)
        else:
            _, _, test_vids = split_train_test.process_other(data_name, test_size, random_state, metadata_path, feature_path, network_name)

        '''======================== read files =============================== '''
        if data_name == 'lsvq_train':
            X_train, y_train, X_test, y_test = load_and_preprocess_data(metadata_path, feature_path, data_name, network_name, train_features, test_features)
        else:
            X_train, y_train, X_test, y_test = load_and_preprocess_data(metadata_path, feature_path, data_name, network_name, None, None)

        '''======================== regression model =============================== '''
        best_model, all_train_losses, all_val_losses = train_and_evaluate(X_train, y_train, config)

        # average loss plots
        avg_train_losses = np.mean(all_train_losses, axis=0)
        avg_val_losses = np.mean(all_val_losses, axis=0)
        test_vids = test_vids.tolist()
        plot_and_save_losses(avg_train_losses, avg_val_losses, model_name, data_name, network_name, len(test_vids), i)

        # predict best model on the train dataset
        y_train_pred = model_test(best_model, X_train, y_train, device)
        y_train_pred = np.array(list(y_train_pred), dtype=float)
        y_train_pred_logistic, plcc_train, rmse_train, srcc_train, krcc_train = compute_correlation_metrics(y_train, y_train_pred)

        # test best model on the test dataset
        y_test_pred = model_test(best_model, X_test, y_test, device)
        y_test_pred = np.array(list(y_test_pred), dtype=float)
        y_test_pred_logistic, plcc_test, rmse_test, srcc_test, krcc_test = compute_correlation_metrics(y_test, y_test_pred)

        # save the predict score results
        test_pred_score = {'MOS': y_test, 'y_test_pred': y_test_pred, 'y_test_pred_logistic': y_test_pred_logistic}
        df_test_pred = pd.DataFrame(test_pred_score)

        # logging logistic predicted scores
        logging.info("============================================================================================================")
        SRCC_all_repeats_test.append(srcc_test)
        KRCC_all_repeats_test.append(krcc_test)
        PLCC_all_repeats_test.append(plcc_test)
        RMSE_all_repeats_test.append(rmse_test)
        SRCC_all_repeats_train.append(srcc_train)
        KRCC_all_repeats_train.append(krcc_train)
        PLCC_all_repeats_train.append(plcc_train)
        RMSE_all_repeats_train.append(rmse_train)
        all_repeats_test_vids.append(test_vids)
        all_repeats_df_test_pred.append(df_test_pred)
        best_model_list.append(copy.deepcopy(best_model))

        # logging.info results for each iteration
        logging.info('Best results in Mlp model within one split')
        logging.info(f'MODEL: {best_model}')
        logging.info('======================================================')
        logging.info(f'Train set - Evaluation Results')
        logging.info(f'SRCC_train: {srcc_train}')
        logging.info(f'KRCC_train: {krcc_train}')
        logging.info(f'PLCC_train: {plcc_train}')
        logging.info(f'RMSE_train: {rmse_train}')
        logging.info('======================================================')
        logging.info(f'Test set - Evaluation Results')
        logging.info(f'SRCC_test: {srcc_test}')
        logging.info(f'KRCC_test: {krcc_test}')
        logging.info(f'PLCC_test: {plcc_test}')
        logging.info(f'RMSE_test: {rmse_test}')
        logging.info('======================================================')
        logging.info(' -- {} seconds elapsed...\n\n'.format(time.time() - t0))

    logging.info('')
    SRCC_all_repeats_test = np.nan_to_num(SRCC_all_repeats_test)
    KRCC_all_repeats_test = np.nan_to_num(KRCC_all_repeats_test)
    PLCC_all_repeats_test = np.nan_to_num(PLCC_all_repeats_test)
    RMSE_all_repeats_test = np.nan_to_num(RMSE_all_repeats_test)
    SRCC_all_repeats_train = np.nan_to_num(SRCC_all_repeats_train)
    KRCC_all_repeats_train = np.nan_to_num(KRCC_all_repeats_train)
    PLCC_all_repeats_train = np.nan_to_num(PLCC_all_repeats_train)
    RMSE_all_repeats_train = np.nan_to_num(RMSE_all_repeats_train)
    logging.info('======================================================')
    logging.info('Average training results among all repeated 80-20 holdouts:')
    logging.info('SRCC: %f (std: %f)', np.median(SRCC_all_repeats_train), np.std(SRCC_all_repeats_train))
    logging.info('KRCC: %f (std: %f)', np.median(KRCC_all_repeats_train), np.std(KRCC_all_repeats_train))
    logging.info('PLCC: %f (std: %f)', np.median(PLCC_all_repeats_train), np.std(PLCC_all_repeats_train))
    logging.info('RMSE: %f (std: %f)', np.median(RMSE_all_repeats_train), np.std(RMSE_all_repeats_train))
    logging.info('======================================================')
    logging.info('Average testing results among all repeated 80-20 holdouts:')
    logging.info('SRCC: %f (std: %f)', np.median(SRCC_all_repeats_test), np.std(SRCC_all_repeats_test))
    logging.info('KRCC: %f (std: %f)', np.median(KRCC_all_repeats_test), np.std(KRCC_all_repeats_test))
    logging.info('PLCC: %f (std: %f)', np.median(PLCC_all_repeats_test), np.std(PLCC_all_repeats_test))
    logging.info('RMSE: %f (std: %f)', np.median(RMSE_all_repeats_test), np.std(RMSE_all_repeats_test))
    logging.info('======================================================')
    logging.info('\n')

    # find the median model and the index of the median
    print('======================================================')
    if select_criteria == 'byrmse':
        median_metrics = np.median(RMSE_all_repeats_test)
        indices = np.where(RMSE_all_repeats_test == median_metrics)[0]
        select_criteria = select_criteria.replace('by', '').upper()
        print(RMSE_all_repeats_test)
        logging.info(f'all {select_criteria}: {RMSE_all_repeats_test}')
    elif select_criteria == 'bykrcc':
        median_metrics = np.median(KRCC_all_repeats_test)
        indices = np.where(KRCC_all_repeats_test == median_metrics)[0]
        select_criteria = select_criteria.replace('by', '').upper()
        print(KRCC_all_repeats_test)
        logging.info(f'all {select_criteria}: {KRCC_all_repeats_test}')

    median_test_vids = [all_repeats_test_vids[i] for i in indices]
    test_vids = [arr.tolist() for arr in median_test_vids] if len(median_test_vids) > 1 else (median_test_vids[0] if median_test_vids else [])

    # select the model with the first index where the median is located
    # Note: If there are multiple iterations with the same median RMSE, the first index is selected here
    median_model = None
    if len(indices) > 0:
        median_index = indices[0]  # select the first index
        median_model = best_model_list[median_index]
        median_model_df_test_pred = all_repeats_df_test_pred[median_index]

        median_model_df_test_pred.to_csv(pred_score_filename, index=False)
        plot_results(y_test, y_test_pred_logistic, median_model_df_test_pred, model_name, data_name, network_name, select_criteria)

    print(f'Median Metrics: {median_metrics}')
    print(f'Indices: {indices}')
    # print(f'Test Videos: {test_vids}')
    print(f'Best model: {median_model}')

    logging.info(f'median test {select_criteria}: {median_metrics}')
    logging.info(f"Indices of median metrics: {indices}")
    # logging.info(f'Best training and test dataset: {test_vids}')
    logging.info(f'Best model predict score: {median_model_df_test_pred}')
    logging.info(f'Best model: {median_model}')

    # ================================================================================
    # save mats
    scipy.io.savemat(result_file, mdict={'SRCC_train': np.asarray(SRCC_all_repeats_train, dtype=float), \
                            'KRCC_train': np.asarray(KRCC_all_repeats_train, dtype=float), \
                            'PLCC_train': np.asarray(PLCC_all_repeats_train, dtype=float), \
                            'RMSE_train': np.asarray(RMSE_all_repeats_train, dtype=float), \
                            'SRCC_test': np.asarray(SRCC_all_repeats_test, dtype=float), \
                            'KRCC_test': np.asarray(KRCC_all_repeats_test, dtype=float), \
                            'PLCC_test': np.asarray(PLCC_all_repeats_test, dtype=float), \
                            'RMSE_test': np.asarray(RMSE_all_repeats_test, dtype=float), \
                            f'Median_{select_criteria}': median_metrics, \
                            'Test_Videos_list': all_repeats_test_vids, \
                            'Test_videos_Median_model': test_vids, \
                            })

    # save model
    torch.save(median_model.state_dict(), file_path)
    print(f"Model state_dict saved to {file_path}")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # input parameters
    parser.add_argument('--model_name', type=str, default='Mlp')
    parser.add_argument('--data_name', type=str, default='lsvq_train', help='konvid_1k, youtube_ugc, live_vqc, cvd_2014, lsvq_train, cross_dataset')
    parser.add_argument('--network_name', type=str, default='relaxvqa', help='relaxvqa, {frag_name}_{network_name}_{layer_name}')

    parser.add_argument('--metadata_path', type=str, default='../metadata/')
    parser.add_argument('--feature_path', type=str, default='../features/')
    parser.add_argument('--log_path', type=str, default='../log/')
    parser.add_argument('--save_path', type=str, default='../model/')
    parser.add_argument('--score_path', type=str, default='../log/predict_score/')
    parser.add_argument('--result_path', type=str, default='../log/result/')
    # training parameters
    parser.add_argument('--select_criteria', type=str, default='byrmse', help='byrmse, bykrcc')
    parser.add_argument('--n_repeats', type=int, default=21, help='Number of repeats for 80-20 hold out test')
    parser.add_argument('--n_splits', type=int, default=10, help='Number of splits for k-fold validation')
    parser.add_argument('--batch_size', type=int, default=256, help='Batch size for training')
    parser.add_argument('--epochs', type=int, default=20, help='Epochs for training') # 120(small), 20(big)
    parser.add_argument('--hidden_features', type=int, default=256, help='Hidden features')
    parser.add_argument('--drop_rate', type=float, default=0.1, help='Dropout rate.')
    # misc
    parser.add_argument('--loss_type', type=str, default='MAERankLoss', help='MSEloss or MAERankLoss')
    parser.add_argument('--optimizer_type', type=str, default='sgd', help='adam or sgd')
    parser.add_argument('--initial_lr', type=float, default=1e-1, help='Initial learning rate: 1e-2')
    parser.add_argument('--weight_decay', type=float, default=0.005, help='Weight decay (L2 loss): 1e-4')
    parser.add_argument('--patience', type=int, default=5, help='Early stopping patience.')
    parser.add_argument('--use_swa', type=bool, default=True, help='Use Stochastic Weight Averaging')
    parser.add_argument('--l1_w', type=float, default=0.6, help='MAE loss weight')
    parser.add_argument('--rank_w', type=float, default=1.0, help='Rank loss weight')

    args = parser.parse_args()
    config = vars(args)  # args to dict
    print(config)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(device)
    if device.type == "cuda":
        torch.cuda.set_device(0)

    main(config)