Upload 5 files

app.py

@@ -0,0 +1,87 @@
+from spaces import GPU
+import gradio as gr
+import torch
+import os
+import time
+from torchvision import models
+from joblib import load
+from extractor.visualise_vit_layer import VitGenerator
+from relax_vqa import get_deep_feature, process_video_feature, process_patches, get_frame_patches, flow_to_rgb, merge_fragments, concatenate_features
+from extractor.vf_extract import process_video_residual
+from model_regression import Mlp, preprocess_data
+from demo_test_gpu import evaluate_video_quality, load_model
+
+
+@GPU
+def run_relax_vqa(video_path, is_finetune, framerate, video_type):
+    if not os.path.exists(video_path):
+        return "❌ No video uploaded or the uploaded file has expired. Please upload again."
+
+    config = {
+        'is_finetune': is_finetune,
+        'framerate': framerate,
+        'video_type': video_type,
+        'save_path': 'model/',
+        'train_data_name': 'lsvq_train',
+        'select_criteria': 'byrmse',
+        'video_path': video_path,
+        'video_name': os.path.splitext(os.path.basename(video_path))[0]
+    }
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    resnet50 = models.resnet50(pretrained=True).to(device)
+    vit = VitGenerator('vit_base', 16, device, evaluate=True, random=False, verbose=False)
+    model_mlp = load_model(config, device)
+
+    try:
+        score, runtime = evaluate_video_quality(config, resnet50, vit, model_mlp, device)
+        return f"Predicted Quality Score: {score:.4f} (in {runtime:.2f}s)"
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+    finally:
+        if "gradio" in video_path and os.path.exists(video_path):
+            os.remove(video_path)
+
+
+def toggle_dataset_visibility(is_finetune):
+    return gr.update(visible=is_finetune)
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("## 🎬 ReLaX-VQA Online Demo")
+    gr.Markdown(
+        "Upload a short video and get the predicted perceptual quality score using the ReLaX-VQA model. "
+        "You can try our demo video from the "
+        "<a href='https://huggingface.co/spaces/xinyiW915/ReLaX-VQA/blob/main/ugc_original_videos/5636101558_540p.mp4' target='_blank'>demo video</a> "
+        "(fps = 24, dataset = konvid_1k).<br><br>"
+        "⚙️ This demo is currently running on <strong>Hugging Face ZeroGPU Space</strong>: Dynamic resources (NVIDIA A100)."
+    )
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            video_input = gr.Video(label="Upload a Video (e.g. mp4)")
+            framerate_slider = gr.Slider(label="Source Video Framerate (fps)", minimum=1, maximum=60, step=1, value=24)
+            is_finetune_checkbox = gr.Checkbox(label="Use Finetuning?", value=False)
+            dataset_dropdown = gr.Dropdown(
+                label="Source Video Dataset for Finetuning",
+                choices=["konvid_1k", "youtube_ugc", "live_vqc", "cvd_2014"],
+                value="konvid_1k",
+                visible=False
+            )
+            run_button = gr.Button("Run Prediction")
+        with gr.Column(scale=1):
+            output_box = gr.Textbox(label="Predicted Quality Score", lines=5)
+
+    is_finetune_checkbox.change(
+        fn=toggle_dataset_visibility,
+        inputs=is_finetune_checkbox,
+        outputs=dataset_dropdown
+    )
+
+    run_button.click(
+        fn=run_relax_vqa,
+        inputs=[video_input, is_finetune_checkbox, framerate_slider, dataset_dropdown],
+        outputs=output_box
+    )
+
+demo.launch()

demo_test_gpu.py

@@ -0,0 +1,249 @@
+import argparse
+import time
+import math
+import os
+import shutil
+from joblib import load
+import cv2
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, Dataset
+from thop import profile
+from torchvision import models, transforms
+
+from extractor.visualise_vit_layer import VitGenerator
+from relax_vqa import get_deep_feature, process_video_feature, process_patches, get_frame_patches, flow_to_rgb, merge_fragments, concatenate_features
+from extractor.vf_extract import process_video_residual
+from model_regression import Mlp, preprocess_data
+
+def fix_state_dict(state_dict):
+    new_state_dict = {}
+    for k, v in state_dict.items():
+        if k.startswith('module.'):
+            name = k[7:]
+        elif k == 'n_averaged':
+            continue
+        else:
+            name = k
+        new_state_dict[name] = v
+    return new_state_dict
+
+def preprocess_data(X, y=None, imp=None, scaler=None):
+    if not isinstance(X, torch.Tensor):
+        X = torch.tensor(X, device='cuda' if torch.cuda.is_available() else 'cpu')
+    X = torch.where(torch.isnan(X) | torch.isinf(X), torch.tensor(0.0, device=X.device), X)
+
+    if imp is not None or scaler is not None:
+        X_np = X.cpu().numpy()
+        if imp is not None:
+            X_np = imp.transform(X_np)
+        if scaler is not None:
+            X_np = scaler.transform(X_np)
+        X = torch.from_numpy(X_np).to(X.device)
+
+    if y is not None and y.size > 0:
+        if not isinstance(y, torch.Tensor):
+            y = torch.tensor(y, device=X.device)
+        y = y.reshape(-1).squeeze()
+    else:
+        y = None
+
+    return X, y, imp, scaler
+
+def load_model(config, device, input_features=35203):
+    network_name = 'relaxvqa'
+    # input_features = X_test_processed.shape[1]
+    model = Mlp(input_features=input_features, out_features=1, drop_rate=0.2, act_layer=nn.GELU).to(device)
+    if config['is_finetune']:
+        model_path = os.path.join(config['save_path'], f"fine_tune_model/{config['video_type']}_{network_name}_{config['select_criteria']}_fine_tuned_model.pth")
+    else:
+        model_path = os.path.join(config['save_path'], f"{config['train_data_name']}_{network_name}_{config['select_criteria']}_trained_median_model_param_onLSVQ_TEST.pth")
+    print("Loading model from:", model_path)
+    state_dict = torch.load(model_path, map_location=device)
+    fixed_state_dict = fix_state_dict(state_dict)
+    try:
+        model.load_state_dict(fixed_state_dict)
+    except RuntimeError as e:
+        print(e)
+    return model
+
+def evaluate_video_quality(config, resnet50, vit, model_mlp, device):
+    is_finetune = config['is_finetune']
+    save_path = config['save_path']
+    video_type = config['video_type']
+    video_name = config['video_name']
+    framerate = config['framerate']
+    sampled_fragment_path = os.path.join("../video_sampled_frame/sampled_frame/", "test_sampled_fragment")
+
+    video_path = config.get("video_path")
+    if video_path is None:
+        if video_type == 'youtube_ugc':
+            video_path = f'./ugc_original_videos/{video_name}.mkv'
+        else:
+            video_path = f'./ugc_original_videos/{video_name}.mp4'
+    target_size = 224
+    patch_size = 16
+    top_n = int((target_size / patch_size) * (target_size / patch_size))
+
+    # sampled video frames
+    start_time = time.time()
+    frames, frames_next = process_video_residual(video_type, video_name, framerate, video_path, sampled_fragment_path)
+
+    # get ResNet50 layer-stack features and ViT pooling features
+    all_frame_activations_resnet = []
+    all_frame_activations_vit = []
+    # get fragments ResNet50 features and ViT features
+    all_frame_activations_sampled_resnet = []
+    all_frame_activations_merged_resnet = []
+    all_frame_activations_sampled_vit = []
+    all_frame_activations_merged_vit = []
+
+    batch_size = 64  # Define the number of frames to process in parallel
+    for i in range(0, len(frames_next), batch_size):
+        batch_frames = frames[i:i + batch_size]
+        batch_rgb_frames = [cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) for frame in batch_frames]
+        batch_frames_next = frames_next[i:i + batch_size]
+        batch_tensors = torch.stack([transforms.ToTensor()(frame) for frame in batch_frames]).to(device)
+        batch_rgb_tensors = torch.stack([transforms.ToTensor()(frame_rgb) for frame_rgb in batch_rgb_frames]).to(device)
+        batch_tensors_next = torch.stack([transforms.ToTensor()(frame_next) for frame_next in batch_frames_next]).to(device)
+
+        # compute residuals
+        residuals = torch.abs(batch_tensors_next - batch_tensors)
+
+        # calculate optical flows
+        batch_gray_frames = [cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) for frame in batch_frames]
+        batch_gray_frames_next = [cv2.cvtColor(frame_next, cv2.COLOR_BGR2GRAY) for frame_next in batch_frames_next]
+        batch_gray_frames = [frame.cpu().numpy() if isinstance(frame, torch.Tensor) else frame for frame in batch_gray_frames]
+        batch_gray_frames_next = [frame.cpu().numpy() if isinstance(frame, torch.Tensor) else frame for frame in batch_gray_frames_next]
+        flows = [cv2.calcOpticalFlowFarneback(batch_gray_frames[j], batch_gray_frames_next[j], None, 0.5, 3, 15, 3, 5, 1.2,0) for j in range(len(batch_gray_frames))]
+
+        for j in range(batch_tensors.size(0)):
+            '''sampled video frames'''
+            frame_tensor = batch_tensors[j].unsqueeze(0)
+            frame_rgb_tensor = batch_rgb_tensors[j].unsqueeze(0)
+            # frame_next_tensor = batch_tensors_next[j].unsqueeze(0)
+            frame_number = i + j + 1
+
+            # ResNet50 layer-stack features
+            activations_dict_resnet, _, _ = get_deep_feature('resnet50', video_name, frame_rgb_tensor, frame_number, resnet50, device, 'layerstack')
+            all_frame_activations_resnet.append(activations_dict_resnet)
+            # ViT pooling features
+            activations_dict_vit, _, _ = get_deep_feature('vit', video_name, frame_rgb_tensor, frame_number, vit, device, 'pool')
+            all_frame_activations_vit.append(activations_dict_vit)
+
+
+            '''residual video frames'''
+            residual = residuals[j].unsqueeze(0)
+            flow = flows[j]
+            original_path = os.path.join(sampled_fragment_path, f'{video_name}_{frame_number}.png')
+
+            # Frame Differencing
+            residual_frag_path, diff_frag, positions = process_patches(original_path, 'frame_diff', residual, patch_size, target_size, top_n)
+            # Frame fragment
+            frame_patches = get_frame_patches(frame_tensor, positions, patch_size, target_size)
+            # Optical Flow
+            opticalflow_rgb = flow_to_rgb(flow)
+            opticalflow_rgb_tensor = transforms.ToTensor()(opticalflow_rgb).unsqueeze(0).to(device)
+            opticalflow_frag_path, flow_frag, _ = process_patches(original_path, 'optical_flow', opticalflow_rgb_tensor, patch_size, target_size, top_n)
+
+            merged_frag = merge_fragments(diff_frag, flow_frag)
+
+            # fragments ResNet50 features
+            sampled_frag_activations_resnet, _, _ = get_deep_feature('resnet50', video_name, frame_patches, frame_number, resnet50, device, 'layerstack')
+            merged_frag_activations_resnet, _, _ = get_deep_feature('resnet50', video_name, merged_frag, frame_number, resnet50, device, 'pool')
+            all_frame_activations_sampled_resnet.append(sampled_frag_activations_resnet)
+            all_frame_activations_merged_resnet.append(merged_frag_activations_resnet)
+            # fragments ViT features
+            sampled_frag_activations_vit,_, _ = get_deep_feature('vit', video_name, frame_patches, frame_number, vit, device, 'pool')
+            merged_frag_activations_vit, _, _ = get_deep_feature('vit', video_name, merged_frag, frame_number, vit, device, 'pool')
+            all_frame_activations_sampled_vit.append(sampled_frag_activations_vit)
+            all_frame_activations_merged_vit.append(merged_frag_activations_vit)
+
+    print(f'video frame number: {len(all_frame_activations_resnet)}')
+    averaged_frames_resnet = process_video_feature(all_frame_activations_resnet, 'resnet50', 'layerstack')
+    averaged_frames_vit = process_video_feature(all_frame_activations_vit, 'vit', 'pool')
+    # print("ResNet50 layer-stacking feature shape:", averaged_frames_resnet.shape)
+    # print("ViT pooling feature shape:", averaged_frames_vit.shape)
+    averaged_frames_sampled_resnet = process_video_feature(all_frame_activations_sampled_resnet, 'resnet50', 'layerstack')
+    averaged_frames_merged_resnet = process_video_feature(all_frame_activations_merged_resnet, 'resnet50', 'pool')
+    averaged_combined_feature_resnet = concatenate_features(averaged_frames_sampled_resnet, averaged_frames_merged_resnet)
+    # print("Sampled fragments ResNet50 features shape:", averaged_frames_sampled_resnet.shape)
+    # print("Merged fragments ResNet50 features shape:", averaged_frames_merged_resnet.shape)
+    averaged_frames_sampled_vit = process_video_feature(all_frame_activations_sampled_vit, 'vit', 'pool')
+    averaged_frames_merged_vit = process_video_feature(all_frame_activations_merged_vit, 'vit', 'pool')
+    averaged_combined_feature_vit = concatenate_features(averaged_frames_sampled_vit, averaged_frames_merged_vit)
+    # print("Sampled fragments ViT features shape:", averaged_frames_sampled_vit.shape)
+    # print("Merged fragments ResNet50 features shape:", averaged_frames_merged_vit.shape)
+
+    # remove tmp folders
+    shutil.rmtree(sampled_fragment_path)
+
+    # concatenate features
+    combined_features = torch.cat([torch.mean(averaged_frames_resnet, dim=0), torch.mean(averaged_frames_vit, dim=0),
+                                   torch.mean(averaged_combined_feature_resnet, dim=0), torch.mean(averaged_combined_feature_vit, dim=0)], dim=0).view(1, -1)
+    imputer = load(f'{save_path}/scaler/{video_type}_imputer.pkl')
+    scaler = load(f'{save_path}/scaler/{video_type}_scaler.pkl')
+    X_test_processed, _, _, _ = preprocess_data(combined_features, None, imp=imputer, scaler=scaler)
+    feature_tensor = X_test_processed
+
+    # evaluation for test video
+    model_mlp.eval()
+
+    with torch.no_grad():
+        with torch.cuda.amp.autocast():
+            prediction = model_mlp(feature_tensor)
+            predicted_score = prediction.item()
+            # print(f"Raw Predicted Quality Score: {predicted_score}")
+            run_time = time.time() - start_time
+
+            if not is_finetune:
+                if video_type in ['konvid_1k', 'youtube_ugc']:
+                    scaled_prediction = ((predicted_score - 1) / (99 / 4)) + 1.0
+                    # print(f"Scaled Predicted Quality Score (1-5): {scaled_prediction}")
+                    return scaled_prediction, run_time
+                else:
+                    scaled_prediction = predicted_score
+                    return scaled_prediction, run_time
+            else:
+                return predicted_score, run_time
+
+def parse_arguments():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-device', type=str, default='gpu', help='cpu or gpu')
+    parser.add_argument('-model_name', type=str, default='Mlp', help='Name of the regression model')
+    parser.add_argument('-select_criteria', type=str, default='byrmse', help='Selection criteria')
+    parser.add_argument('-train_data_name', type=str, default='lsvq_train', help='Name of the training data')
+    parser.add_argument('-is_finetune', type=bool, default=False, help='With or without finetune')
+    parser.add_argument('-save_path', type=str, default='model/', help='Path to save models')
+    parser.add_argument('-video_type', type=str, default='konvid_1k', help='Type of video')
+    parser.add_argument('-video_name', type=str, default='5636101558_540p', help='Name of the video')
+    parser.add_argument('-framerate', type=float, default=24, help='Frame rate of the video')
+
+    args = parser.parse_args()
+    return args
+
+if __name__ == '__main__':
+    args = parse_arguments()
+    config = vars(args)
+    if config['device'] == "gpu":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device("cpu")
+    print(f"Running on {'GPU' if device.type == 'cuda' else 'CPU'}")
+
+    # load models to device
+    resnet50 = models.resnet50(pretrained=True).to(device)
+    vit = VitGenerator('vit_base', 16, device, evaluate=True, random=False, verbose=True)
+    model_mlp = load_model(config, device)
+
+    total_time = 0
+    num_runs = 1
+    for i in range(num_runs):
+        quality_prediction, run_time = evaluate_video_quality(config, resnet50, vit, model_mlp, device)
+        print(f"Run {i + 1} - Time taken: {run_time:.4f} seconds")
+
+        total_time += run_time
+    average_time = total_time / num_runs
+
+    print(f"Average running time over {num_runs} runs: {average_time:.4f} seconds")
+    print("Predicted Quality Score:", quality_prediction)

model_regression.py

@@ -0,0 +1,693 @@
+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)

relax_vqa.py

@@ -0,0 +1,159 @@
+import torch
+import os
+import cv2
+import numpy as np
+from extractor import visualise_resnet, visualise_resnet_layer, visualise_vit_layer
+
+
+def get_deep_feature(network_name, video_name, frame, frame_number, model, device, layer_name):
+    if network_name == 'resnet50':
+        if layer_name == 'layerstack':
+            all_layers = ['resnet50.conv1',
+                          'resnet50.layer1[0]', 'resnet50.layer1[1]', 'resnet50.layer1[2]',
+                          'resnet50.layer2[0]', 'resnet50.layer2[1]', 'resnet50.layer2[2]', 'resnet50.layer2[3]',
+                          'resnet50.layer3[0]', 'resnet50.layer3[1]', 'resnet50.layer3[2]', 'resnet50.layer3[3]',
+                          'resnet50.layer4[0]', 'resnet50.layer4[1]', 'resnet50.layer4[2]']
+            resnet50 = model
+            activations_dict, _, total_flops, total_params = visualise_resnet.process_video_frame(video_name, frame, frame_number, all_layers, resnet50, device)
+
+        elif layer_name == 'pool':
+            visual_layer = 'resnet50.avgpool' # before avg_pool
+            resnet50 = model
+            activations_dict, _, total_flops, total_params = visualise_resnet_layer.process_video_frame(video_name, frame, frame_number, visual_layer, resnet50, device)
+
+    elif network_name == 'vit':
+        patch_size = 16
+        activations_dict, _, total_flops, total_params = visualise_vit_layer.process_video_frame(video_name, frame, frame_number, model, patch_size, device)
+
+    return activations_dict, total_flops, total_params
+
+
+def process_video_feature(video_feature, network_name, layer_name):
+    # initialize an empty list to store processed frames
+    averaged_frames = []
+    # iterate through each frame in the video_feature
+    for frame in video_feature:
+        frame_features = []
+
+        if network_name == 'vit':
+            # global mean and std
+            global_mean = torch.mean(frame, dim=0)
+            global_max = torch.max(frame, dim=0)[0]
+            global_std = torch.std(frame, dim=0)
+            # concatenate all pooling
+            combined_features = torch.hstack([global_mean, global_max, global_std])
+            frame_features.append(combined_features)
+
+        elif network_name == 'resnet50':
+            if layer_name == 'layerstack':
+                # iterate through each layer in the current framex
+                for layer_array in frame.values():
+                    # calculate the mean along the specified axes (1 and 2) for each layer
+                    layer_mean = torch.mean(layer_array, dim=(1, 2))
+                    # append the calculated mean to the list for the current frame
+                    frame_features.append(layer_mean)
+            elif layer_name == 'pool':
+                frame = torch.squeeze(torch.tensor(frame))
+                # global mean and std
+                global_mean = torch.mean(frame, dim=0)
+                global_max = torch.max(frame, dim=0)[0]
+                global_std = torch.std(frame, dim=0)
+                # concatenate all pooling
+                combined_features = torch.hstack([frame, global_mean, global_max, global_std])
+                frame_features.append(combined_features)
+
+        # concatenate the layer means horizontally to form the processed frame
+        processed_frame = torch.hstack(frame_features)
+        averaged_frames.append(processed_frame)
+
+    averaged_frames = torch.stack(averaged_frames)
+    return averaged_frames
+
+
+def flow_to_rgb(flow):
+    mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1])
+    mag = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)
+    # convert angle to hue
+    hue = ang * 180 / np.pi / 2
+
+    # create HSV
+    hsv = np.zeros((flow.shape[0], flow.shape[1], 3), dtype=np.uint8)
+    hsv[..., 0] = hue
+    hsv[..., 1] = 255
+    hsv[..., 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)
+    # convert HSV to RGB
+    rgb = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
+    return rgb
+
+def get_patch_diff(residual_frame, patch_size):
+    h, w = residual_frame.shape[2:]  # Assuming (1, C, H, W) shape
+    h_adj = (h // patch_size) * patch_size
+    w_adj = (w // patch_size) * patch_size
+    residual_frame_adj = residual_frame[:, :, :h_adj, :w_adj]
+    # calculate absolute patch difference
+    diff = torch.zeros((h_adj // patch_size, w_adj // patch_size), device=residual_frame.device)
+    for i in range(0, h_adj, patch_size):
+        for j in range(0, w_adj, patch_size):
+            patch = residual_frame_adj[:, :, i:i + patch_size, j:j + patch_size]
+            # absolute sum
+            diff[i // patch_size, j // patch_size] = torch.sum(torch.abs(patch))
+    return diff
+
+def extract_important_patches(residual_frame, diff, patch_size=16, target_size=224, top_n=196):
+    # find top n patches indices
+    patch_idx = torch.argsort(-diff.view(-1))
+    top_patches = [(idx // diff.shape[1], idx % diff.shape[1]) for idx in patch_idx[:top_n]]
+    sorted_idx = sorted(top_patches, key=lambda x: (x[0], x[1]))
+
+    imp_patches_img = torch.zeros((residual_frame.shape[1], target_size, target_size), dtype=residual_frame.dtype, device=residual_frame.device)
+    patches_per_row = target_size // patch_size  # 14
+    # order the patch in the original location relation
+    positions = []
+    for idx, (y, x) in enumerate(sorted_idx):
+        patch = residual_frame[:, :, y * patch_size:(y + 1) * patch_size, x * patch_size:(x + 1) * patch_size]
+        # new patch location
+        row_idx = idx // patches_per_row
+        col_idx = idx % patches_per_row
+        start_y = row_idx * patch_size
+        start_x = col_idx * patch_size
+        imp_patches_img[:, start_y:start_y + patch_size, start_x:start_x + patch_size] = patch
+        positions.append((y.item(), x.item()))
+    return imp_patches_img, positions
+
+def get_frame_patches(frame, positions, patch_size, target_size):
+    imp_patches_img = torch.zeros((frame.shape[1], target_size, target_size), dtype=frame.dtype, device=frame.device)
+    patches_per_row = target_size // patch_size
+
+    for idx, (y, x) in enumerate(positions):
+        start_y = y * patch_size
+        start_x = x * patch_size
+        end_y = start_y + patch_size
+        end_x = start_x + patch_size
+
+        patch = frame[:, :, start_y:end_y, start_x:end_x]
+        row_idx = idx // patches_per_row
+        col_idx = idx % patches_per_row
+        target_start_y = row_idx * patch_size
+        target_start_x = col_idx * patch_size
+
+        imp_patches_img[:, target_start_y:target_start_y + patch_size,
+        target_start_x:target_start_x + patch_size] = patch.squeeze(0)
+    return imp_patches_img
+
+def process_patches(original_path, frag_name, residual, patch_size, target_size, top_n):
+    diff = get_patch_diff(residual, patch_size)
+    imp_patches, positions = extract_important_patches(residual, diff, patch_size, target_size, top_n)
+    if frag_name == 'frame_diff':
+        frag_path = original_path.replace('.png', '_residual_imp.png')
+    elif frag_name == 'optical_flow':
+        frag_path = original_path.replace('.png', '_residual_of_imp.png')
+    # cv2.imwrite(frag_path, imp_patches)
+    return frag_path, imp_patches, positions
+
+def merge_fragments(diff_fragment, flow_fragment):
+    alpha = 0.5
+    merged_fragment = diff_fragment * alpha + flow_fragment * (1 - alpha)
+    return merged_fragment
+
+def concatenate_features(frame_feature, residual_feature):
+    return torch.cat((frame_feature, residual_feature), dim=-1)

requirements.txt

@@ -0,0 +1,15 @@
+gradio
+torch
+torchvision
+torchaudio
+opencv-python
+joblib
+scikit-learn
+scipy
+numpy
+pandas
+matplotlib
+ipywidgets
+thop
+PyYAML
+seaborn