Upload 6 files

app.py

@@ -0,0 +1,176 @@
+from spaces import GPU
+import gradio as gr
+import torch
+import os
+import pandas as pd
+from types import SimpleNamespace
+import clip
+from transformers import Blip2Processor, Blip2ForConditionalGeneration
+
+from extractor.extract_frag import VideoDataset_feature
+from extractor.extract_slowfast_clip import SlowFast
+from extractor.extract_swint_clip import SwinT
+from demo_test import get_transform, load_prompts, get_video_metadata, load_model, evaluate_video_quality
+
+
+model_cache = {}
+
+@GPU
+def run_camp_vqa(video_path, intra_cross_experiment, is_finetune, train_data_name, test_data_name, network_name):
+    if not os.path.exists(video_path):
+        return "❌ No video uploaded or the uploaded file has expired. Please upload again."
+
+    print("CUDA available:", torch.cuda.is_available())
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if device.type == "cuda":
+        print("Current device:", torch.cuda.current_device())
+    else:
+        print("Running on CPU")
+
+    config = SimpleNamespace(**{
+        'model_name': 'Mlp',
+        'select_criteria': 'byrmse',
+        'intra_cross_experiment': intra_cross_experiment,
+        'is_finetune': is_finetune,
+        'save_model_path': 'model/',
+        'prompt_path': './config/prompts.json',
+        'train_data_name': train_data_name,
+        'test_data_name': test_data_name,
+        'test_video_path': video_path,
+        'prediction_mode': 50,
+        'network_name': network_name,
+        'num_workers': 2,
+        'resize': 224,
+        'patch_size': 16,
+        'target_size': 224,
+    })
+    print(f"Test video path: {config.test_video_path}")
+
+    # test demo video
+    resize_transform = get_transform(config.resize)
+    top_n = int(config.target_size / config.patch_size) ** 2
+
+    width, height, bitrate, bitdepth, framerate = get_video_metadata(config.test_video_path)
+
+    data = {'vid': [os.path.splitext(os.path.basename(config.test_video_path))[0]],
+        'test_data_name': [config.test_data_name],
+        'test_video_path': [config.test_video_path],
+        'prediction_mode': [config.prediction_mode],
+        'width': [width], 'height': [height], 'bitrate': [bitrate], 'bitdepth': [bitdepth], 'framerate': [framerate]}
+    videos_dir = os.path.dirname(config.test_video_path)
+    test_df = pd.DataFrame(data)
+    print(test_df.T)
+    print(f"Experiment Setting: {config.intra_cross_experiment}, {config.train_data_name} -> {config.test_data_name}")
+    if config.intra_cross_experiment == 'cross':
+        if config.train_data_name == 'lsvq_train':
+            print(f"Fine-tune: {config.is_finetune}")
+
+    # load models to device
+    global model_cache
+    if not model_cache:
+        print("Loading models into cache (first time)...")
+        model_cache["slowfast"] = SlowFast().to(device)
+        model_cache["swint"] = SwinT(model_name='swin_large_patch4_window7_224', global_pool='avg', pretrained=True).to(device)
+
+        model_cache["clip"], model_cache["clip_preprocess"] = clip.load("ViT-B/32", device=device)
+        model_cache["blip_processor"] = Blip2Processor.from_pretrained("Salesforce/blip2-flan-t5-xl", use_fast=True)
+        model_cache["blip_model"] = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-flan-t5-xl").to(device)
+        print("Model cache initialized.")
+
+    # get model from cache
+    model_slowfast = model_cache["slowfast"]
+    model_swint = model_cache["swint"]
+    clip_model, clip_preprocess = model_cache["clip"], model_cache["clip_preprocess"]
+    blip_processor = model_cache["blip_processor"]
+    blip_model = model_cache["blip_model"]
+
+    input_features = 13056
+    if config.intra_cross_experiment == 'intra':
+        if config.train_data_name == 'lsvq_train':
+            from model_regression_lsvq import Mlp, preprocess_data
+        else:
+            from model_regression import Mlp, preprocess_data
+    elif config.intra_cross_experiment == 'cross':
+        from model_regression_lsvq import Mlp, preprocess_data
+    model_mlp = load_model(config, device, Mlp, input_features)
+    prompts = load_prompts(config.prompt_path)
+
+    dataset = VideoDataset_feature(test_df, videos_dir, config.test_data_name, resize_transform, config.resize, config.patch_size, config.target_size, top_n)
+    data_loader = torch.utils.data.DataLoader(
+        dataset, batch_size=1, shuffle=False, num_workers=min(config.num_workers, os.cpu_count() or 1), pin_memory=(device.type == "cuda")
+    )
+
+    try:
+        score = evaluate_video_quality(
+            preprocess_data,
+            data_loader,
+            model_slowfast,
+            model_swint,
+            clip_model,
+            clip_preprocess,
+            blip_processor,
+            blip_model,
+            prompts,
+            model_mlp,
+            device
+        )
+        return f"**Predicted Perceptual Quality Score:** {score:.4f} / 100"
+
+    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("# 📹 CAMP-VQA Online Demo")
+    gr.Markdown(
+        "Upload a short video and get its perceptual quality score predicted by CAMP-VQA."
+        "You can try our test video"
+        "<a href='https://huggingface.co/spaces/xinyiW915/CAMP-VQA/blob/main/ugc_original_videos/0_16_07_500001604801190-yase.mp4' target='_blank'>demo video</a>. "
+        "<br><br>"
+        # "⚙️ This demo is currently running on <strong>Hugging Face CPU Basic</strong>: 2 vCPU • 16 GB RAM."
+        "⚙️ 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)")
+            intra_cross_experiment = gr.Dropdown(
+                label="Intra or Cross experiment",
+                choices=["intra", "cross"],
+                value="cross"
+            )
+            is_finetune_checkbox = gr.Checkbox(label="Use Finetuning?", value=False)
+            train_dataset = gr.Dropdown(
+                label="Train Dataset",
+                choices=["lsvq_train", "cvd_2014", "konvid_1k", "live_vqc", "youtube_ugc", "finevd", "live_yt_gaming", "kvq"],
+                value="lsvq_train"
+            )
+            test_dataset = gr.Dropdown(
+                label="Test Dataset for Finetuning",
+                choices=["lsvq_test", "lsvq_test_1080p", "cvd_2014", "konvid_1k", "live_vqc", "youtube_ugc", "finevd", "live_yt_gaming", "kvq"],
+                value="finevd"
+            )
+            model_dropdown = gr.Dropdown(
+                label="Our Models",
+                choices=["camp-vqa"],
+                value="camp-vqa"
+            )
+            run_button = gr.Button("Run Prediction")
+
+        with gr.Column(scale=1):
+            output_box = gr.Textbox(label="Predicted Quality Score (0–100)", lines=5)
+
+    run_button.click(
+        fn=run_camp_vqa,
+        inputs=[video_input, intra_cross_experiment, is_finetune_checkbox, train_dataset, test_dataset, model_dropdown],
+        outputs=output_box,
+        api_name="run",
+        queue=True
+    )
+
+demo.launch()

demo_test.py

@@ -0,0 +1,220 @@
+import argparse
+import os
+import sys
+import subprocess
+import json
+import ffmpeg
+import pandas as pd
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+from torchvision import transforms
+import clip
+from transformers import Blip2Processor, Blip2ForConditionalGeneration
+
+from extractor.extract_frag import VideoDataset_feature
+from extractor.extract_clip_embeds import extract_features_clip_embed
+from extractor.extract_slowfast_clip import SlowFast, extract_features_slowfast_pool
+from extractor.extract_swint_clip import SwinT, extract_features_swint_pool
+from model_finetune import fix_state_dict
+
+
+def get_transform(resize):
+    return transforms.Compose([transforms.Resize([resize, resize]),
+                               transforms.ToTensor(),
+                               transforms.Normalize(mean=[0.45, 0.45, 0.45], std=[0.225, 0.225, 0.225])])
+
+def setup_device(config):
+    if config.device == "gpu":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        if device.type == "cuda":
+            torch.cuda.set_device(0)
+    else:
+        device = torch.device("cpu")
+    print(f"Running on {'GPU' if device.type == 'cuda' else 'CPU'}")
+    return device
+
+def load_prompts(json_path):
+    with open(json_path, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+def load_model(config, device, Mlp, input_features=13056):
+    model = Mlp(input_features=input_features, out_features=1, drop_rate=0.1, act_layer=nn.GELU).to(device)
+
+    if config.intra_cross_experiment == 'intra':
+        if config.train_data_name == 'lsvq_train':
+            if config.test_data_name  == 'lsvq_test':
+                model_path = os.path.join(config.save_model_path, f"wo_finetune/{config.train_data_name}_{config.network_name}_{config.model_name}_{config.select_criteria}_trained_model_kfold.pth")
+            elif config.test_data_name  == 'lsvq_test_1080p':
+                model_path = os.path.join(config.save_model_path, f"wo_finetune/{config.train_data_name}_{config.network_name}_{config.model_name}_{config.select_criteria}_trained_model_1080p.pth")
+            else:
+                print("Please use a cross-dataset experiment setting for the lsvq_train model to test it on another dataset, please try using the input 'cross' for 'intra_cross_experiment'.")
+                sys.exit(1)
+        else:
+            model_path = os.path.join(config.save_model_path, f"wo_finetune/{config.train_data_name}_{config.network_name}_{config.model_name}_{config.select_criteria}_trained_model.pth")
+
+    elif config.intra_cross_experiment == 'cross':
+        if config.train_data_name == 'lsvq_train':
+            if config.is_finetune:
+                model_path = os.path.join(config.save_model_path, f"finetune/{config.test_data_name}_{config.network_name}_fine_tuned_model.pth")
+            else:
+                model_path = os.path.join(config.save_model_path, f"wo_finetune/{config.train_data_name}_{config.network_name}_{config.model_name}_{config.select_criteria}_trained_model_kfold.pth")
+        else:
+            print("Invalid training data name for cross-experiment. We provided the lsvq_train model for the cross-experiment, please try using the input 'lsvq_train' for 'train_data_name'.")
+            sys.exit(1)
+
+    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(preprocess_data, data_loader, model_slowfast, model_swint, clip_model, clip_preprocess, blip_processor, blip_model, prompts, model_mlp, device):
+    # get video features
+    model_slowfast.eval()
+    model_swint.eval()
+    clip_model.eval()
+    blip_model.eval()
+    with torch.no_grad():
+        for i, (video_segments, video_res_frag_all, video_frag_all, video_name, frames_info, metadata) in enumerate(tqdm(data_loader, desc="Processing Videos")):
+            # slowfast features
+            _, _, slowfast_frame_feats = extract_features_slowfast_pool(video_segments, model_slowfast, device)
+            _, _, slowfast_res_frag_feats = extract_features_slowfast_pool(video_res_frag_all, model_slowfast, device)
+            _, _, slowfast_frame_frag_feats = extract_features_slowfast_pool(video_frag_all, model_slowfast, device)
+            slowfast_frame_feats_avg = slowfast_frame_feats.mean(dim=0)
+            slowfast_res_frag_feats_avg = slowfast_res_frag_feats.mean(dim=0)
+            slowfast_frame_frag_feats_avg = slowfast_frame_frag_feats.mean(dim=0)
+
+            # swinT feature
+            swint_frame_feats = extract_features_swint_pool(video_segments, model_swint, device)
+            swint_res_frag_feats = extract_features_swint_pool(video_res_frag_all, model_swint, device)
+            swint_frame_frag_feats = extract_features_swint_pool(video_frag_all, model_swint, device)
+            swint_frame_feats_avg = swint_frame_feats.mean(dim=0)
+            swint_res_frag_feats_avg = swint_res_frag_feats.mean(dim=0)
+            swint_frame_frag_feats_avg = swint_frame_frag_feats.mean(dim=0)
+
+            # semantic features
+            image_embedding, quality_embedding, artifact_embedding = extract_features_clip_embed(frames_info, metadata, clip_model, clip_preprocess, blip_processor, blip_model, prompts, device)
+            image_embedding_avg = image_embedding.mean(dim=0)
+            quality_embedding_avg = quality_embedding.mean(dim=0)
+            artifact_embedding_avg = artifact_embedding.mean(dim=0)
+
+            # frame + residual fragment + frame fragment features
+            slowfast_features = torch.cat((slowfast_frame_feats_avg, slowfast_res_frag_feats_avg, slowfast_frame_frag_feats_avg), dim=0)
+            swint_features = torch.cat((swint_frame_feats_avg, swint_res_frag_feats_avg, swint_frame_frag_feats_avg), dim=0)
+            clip_features = torch.cat((image_embedding_avg, quality_embedding_avg, artifact_embedding_avg), dim=0)
+            vqa_feats = torch.cat((slowfast_features, swint_features, clip_features), dim=0)
+
+    vqa_feats = vqa_feats
+    feature_tensor, _ = preprocess_data(vqa_feats, None)
+    feature_tensor = feature_tensor.unsqueeze(0) if feature_tensor.dim() == 1 else feature_tensor
+    print(f"Feature tensor shape before MLP: {feature_tensor.shape}")
+
+    model_mlp.eval()
+    with torch.no_grad():
+        with torch.amp.autocast(device_type=device.type if device.type == 'cuda' else 'cpu'):
+            prediction = model_mlp(feature_tensor)
+            predicted_score = prediction.item()
+            return predicted_score
+
+def parse_framerate(framerate_str):
+    num, den = framerate_str.split('/')
+    framerate = float(num)/float(den)
+    return framerate
+
+def get_video_metadata(video_path):
+    print(video_path)
+    ffprobe_path = 'ffprobe'
+    cmd = f'{ffprobe_path} -v error -select_streams v:0 -show_entries stream=width,height,nb_frames,r_frame_rate,bit_rate,bits_per_raw_sample,pix_fmt -of json {video_path}'
+    try:
+        result = subprocess.run(cmd, shell=True, capture_output=True, check=True)
+        info = json.loads(result.stdout)
+    except Exception as e:
+        print(f"Error processing file {video_path}: {e}")
+        return {}
+
+    width = info['streams'][0]['width']
+    height = info['streams'][0]['height']
+    bitrate = info['streams'][0].get('bit_rate', 0)
+    bitdepth = info['streams'][0].get('bits_per_raw_sample', 0)
+    framerate = info['streams'][0]['r_frame_rate']
+    framerate = parse_framerate(framerate)
+    return width, height, bitrate, bitdepth, framerate
+
+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')
+    parser.add_argument('--select_criteria', type=str, default='byrmse')
+    parser.add_argument('--intra_cross_experiment', type=str, default='cross', help='intra or cross')
+    parser.add_argument('--is_finetune', type=bool, default=True, help='True or False')
+    parser.add_argument('--save_model_path', type=str, default='../model/')
+    parser.add_argument('--prompt_path', type=str, default="./config/prompts.json")
+
+    parser.add_argument('--train_data_name', type=str, default='lsvq_train', help='Name of the training data')
+    parser.add_argument('--test_data_name', type=str, default='finevd', help='Name of the testing data')
+    parser.add_argument('--test_video_path', type=str, default='../test_videos/0_16_07_500001604801190-yase.mp4', help='demo test video')
+    parser.add_argument('--prediction_mode', type=float, default=50, help='default for inference')
+
+    parser.add_argument('--network_name', type=str, default='camp-vqa')
+    parser.add_argument('--num_workers', type=int, default=4)
+    parser.add_argument('--resize', type=int, default=224)
+    parser.add_argument('--patch_size', type=int, default=16)
+    parser.add_argument('--target_size', type=int, default=224)
+    args = parser.parse_args()
+    return args
+
+if __name__ == '__main__':
+    config = parse_arguments()
+    device = setup_device(config)
+    prompts = load_prompts(config.prompt_path)
+
+    # test demo video
+    resize_transform = get_transform(config.resize)
+    top_n = int(config.target_size /config. patch_size) * int(config.target_size / config.patch_size)
+
+    width, height, bitrate, bitdepth, framerate = get_video_metadata(config.test_video_path)
+
+    data = {'vid': [os.path.splitext(os.path.basename(config.test_video_path))[0]],
+        'test_data_name': [config.test_data_name],
+        'test_video_path': [config.test_video_path],
+        'prediction_mode': [config.prediction_mode],
+        'width': [width], 'height': [height], 'bitrate': [bitrate], 'bitdepth': [bitdepth], 'framerate': [framerate]}
+    videos_dir = os.path.dirname(config.test_video_path)
+    test_df = pd.DataFrame(data)
+    print(test_df.T)
+    print(f"Experiment Setting: {config.intra_cross_experiment}, {config.train_data_name} -> {config.test_data_name}")
+    if config.intra_cross_experiment == 'cross':
+        if config.train_data_name == 'lsvq_train':
+            print(f"Fine-tune: {config.is_finetune}")
+
+    dataset = VideoDataset_feature(test_df, videos_dir, config.test_data_name, resize_transform, config.resize, config.patch_size, config.target_size, top_n)
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset, batch_size=1, shuffle=False, num_workers = min(config.num_workers, os.cpu_count() or 1), pin_memory = device.type == "cuda"
+    )
+    print(f"Model: {config.network_name} | Dataset: {config.test_data_name} | Device: {device}")
+
+    # load models to device
+    model_slowfast = SlowFast().to(device)
+    model_swint = SwinT(model_name='swin_large_patch4_window7_224', global_pool='avg', pretrained=True).to(device)
+
+    clip_model, clip_preprocess = clip.load("ViT-B/32", device=device)
+    blip_processor = Blip2Processor.from_pretrained("Salesforce/blip2-flan-t5-xl", use_fast=True)
+    blip_model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-flan-t5-xl").to(device)
+
+    input_features = 13056
+    if config.intra_cross_experiment == 'intra':
+        if config.train_data_name == 'lsvq_train':
+            from model_regression_lsvq import Mlp, preprocess_data
+        else:
+            from model_regression import Mlp, preprocess_data
+    elif config.intra_cross_experiment == 'cross':
+        from model_regression_lsvq import Mlp, preprocess_data
+    model_mlp = load_model(config, device, Mlp, input_features)
+
+    quality_prediction = evaluate_video_quality(preprocess_data, data_loader, model_slowfast, model_swint, clip_model, clip_preprocess, blip_processor, blip_model, prompts, model_mlp, device)
+    print("Predicted Quality Score:", quality_prediction)

model_finetune.py

@@ -0,0 +1,326 @@
+import argparse
+import pandas as pd
+import numpy as np
+import math
+import os
+import scipy.io
+import scipy.stats
+from scipy.optimize import curve_fit
+from sklearn.model_selection import train_test_split
+import seaborn as sns
+import matplotlib.pyplot as plt
+import copy
+from joblib import dump, load
+import torch
+import torch.nn as nn
+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 model_regression_lsvq import Mlp, MAEAndRankLoss, preprocess_data, compute_correlation_metrics, logistic_func, plot_results
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+if device.type == "cuda":
+    torch.cuda.set_device(0)
+
+def create_results_dataframe(data_list, network_name, srcc_list, krcc_list, plcc_list, rmse_list, select_criteria_list):
+    df_results = pd.DataFrame(columns=['DATASET', 'MODEL', 'SRCC', 'KRCC', 'PLCC', 'RMSE', 'SELECT_CRITERIC'])
+    df_results['DATASET'] = data_list
+    df_results['MODEL'] = network_name
+    df_results['SRCC'] = srcc_list
+    df_results['KRCC'] = krcc_list
+    df_results['PLCC'] = plcc_list
+    df_results['RMSE'] = rmse_list
+    df_results['SELECT_CRITERIC'] = select_criteria_list
+    return df_results
+
+def process_test_set(test_data_name, metadata_path, feature_path, network_name):
+    test_df = pd.read_csv(f'{metadata_path}/{test_data_name.upper()}_metadata.csv')
+
+    test_vids = test_df['vid']
+    mos = torch.tensor(test_df['mos'].astype(float), dtype=torch.float32)
+    if test_data_name in ('konvid_1k', 'youtube_ugc_h264'):
+        test_scores = ((mos - 1) * (99 / 4) + 1.0)
+    else:
+        test_scores = mos
+
+    sorted_test_df = pd.DataFrame({
+        'vid': test_df['vid'],
+        'framerate': test_df['framerate'],
+        'MOS': test_scores,
+        'MOS_raw': mos
+    })
+    test_features = torch.load(f'{feature_path}/{network_name}_{test_data_name}_features.pt')
+    print(f'num of {test_data_name} features: {len(test_features)}')
+    return test_features, test_vids, test_scores, sorted_test_df
+
+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 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(X, 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 fine_tune_model(model, device, model_path, X_fine_tune, y_fine_tune, save_path, batch_size, epochs, loss_type, optimizer_type, initial_lr, weight_decay, use_swa, l1_w, rank_w):
+    state_dict = torch.load(model_path)
+    fixed_state_dict = fix_state_dict(state_dict)
+    try:
+        model.load_state_dict(fixed_state_dict)
+    except RuntimeError as e:
+        print(e)
+
+    for param in model.parameters():
+        param.requires_grad = True
+    model.train().to(device) # to gpu
+
+    fine_tune_dataset = TensorDataset(X_fine_tune, y_fine_tune)
+    fine_tune_loader = DataLoader(dataset=fine_tune_dataset, batch_size=batch_size, shuffle=False)
+
+    # initialisation of loss function, optimiser
+    if loss_type == 'MAERankLoss':
+        criterion = MAEAndRankLoss()
+        criterion.l1_w = l1_w
+        criterion.rank_w = rank_w
+    else:
+        criterion = 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_min=1e-5
+    else:
+        optimizer = optim.AdamW(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')
+    swa_start = int(epochs * 0.75) if use_swa else epochs  # SWA starts after 75% of total epochs, only set SWA start if SWA is used
+
+    best_loss = float('inf')
+    for epoch in range(epochs):
+        model.train()
+        epoch_loss = 0.0
+        for inputs, labels in fine_tune_loader:
+            inputs, labels = inputs.to(device), labels.to(device)
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            loss = criterion(outputs, labels.view(-1, 1))
+            loss.backward()
+            optimizer.step()
+            epoch_loss += loss.item() * inputs.size(0)
+
+        scheduler.step()
+        if use_swa and epoch >= swa_start:
+            swa_model.update_parameters(model)
+            swa_scheduler.step()
+            print(f"Current learning rate with SWA: {swa_scheduler.get_last_lr()}")
+        avg_loss = epoch_loss / len(fine_tune_loader.dataset)
+        if (epoch + 1) % 5 == 0:
+            print(f"Epoch {epoch+1}, Loss: {avg_loss:.4f}")
+
+        # decide which model to evaluate: SWA model or regular model
+        current_model = swa_model if use_swa and epoch >= swa_start else model
+        # Save best model state
+        if avg_loss < best_loss:
+            best_loss = avg_loss
+            best_model = copy.deepcopy(current_model)
+
+    # decide which model to evaluate: SWA model or regular model
+    if use_swa and epoch >= swa_start:
+        train_loader = DataLoader(dataset=fine_tune_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)
+    # model_path_new = os.path.join(save_path, f"{test_data_name}_diva-vqa_fine_tuned_model.pth")
+    # torch.save(best_model.state_dict(), model_path_new)  # save finetuned model
+    return best_model
+
+def fine_tuned_model_test(model, device, X_test, y_test, test_data_name):
+    model.eval()
+    y_test_pred = model_test(model, X_test, y_test, device)
+    y_test_pred = torch.tensor(list(y_test_pred), dtype=torch.float32)
+    if test_data_name in ('konvid_1k', 'youtube_ugc_h264'):
+        y_test_convert = ((y_test - 1) / (99 / 4) + 1.0)
+        y_test_pred_convert = ((y_test_pred - 1) / (99 / 4) + 1.0)
+    else:
+        y_test_convert = y_test
+        y_test_pred_convert = y_test_pred
+
+    y_test_pred_logistic, plcc_test, rmse_test, srcc_test, krcc_test = compute_correlation_metrics(y_test_convert.cpu().numpy(), y_test_pred_convert.cpu().numpy())
+    test_pred_score = {'MOS': y_test_convert, 'y_test_pred': y_test_pred_convert, 'y_test_pred_logistic': y_test_pred_logistic}
+    df_test_pred = pd.DataFrame(test_pred_score)
+    return df_test_pred, y_test_convert, y_test_pred_logistic, plcc_test, rmse_test, srcc_test, krcc_test
+
+def wo_fine_tune_model(model, device, model_path, X_test, y_test, loss_type, test_data_name):
+    state_dict = torch.load(model_path)
+    fixed_state_dict = fix_state_dict(state_dict)
+    try:
+        model.load_state_dict(fixed_state_dict)
+    except RuntimeError as e:
+        print(e)
+    model.eval().to(device) # to gpu
+
+    if loss_type == 'MAERankLoss':
+        criterion = MAEAndRankLoss()
+    else:
+        criterion = torch.nn.MSELoss()
+
+    # evaluate the model
+    test_dataset = TensorDataset(X_test, y_test)
+    test_loader = DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)
+    test_loss = 0.0
+    for inputs, labels in test_loader:
+        inputs, labels = inputs.to(device), labels.to(device)
+        outputs = model(inputs)
+        loss = criterion(outputs, labels.view(-1, 1))
+        test_loss += loss.item() * inputs.size(0)
+    average_loss = test_loss / len(test_loader.dataset)
+    print(f"Test Loss: {average_loss}")
+
+    y_test_pred = model_test(model, X_test, y_test, device)
+    y_test_pred = torch.tensor(list(y_test_pred), dtype=torch.float32)
+    if test_data_name in ('konvid_1k', 'youtube_ugc_h264'):
+        y_test_convert = ((y_test - 1) / (99 / 4) + 1.0)
+        y_test_pred_convert = ((y_test_pred - 1) / (99 / 4) + 1.0)
+    else:
+        y_test_convert = y_test
+        y_test_pred_convert = y_test_pred
+
+    y_test_pred_logistic, plcc_test, rmse_test, srcc_test, krcc_test = compute_correlation_metrics(y_test_convert.cpu().numpy(), y_test_pred_convert.cpu().numpy())
+    test_pred_score = {'MOS': y_test_convert, 'y_test_pred': y_test_pred_convert, 'y_test_pred_logistic': y_test_pred_logistic}
+    df_test_pred = pd.DataFrame(test_pred_score)
+    return df_test_pred, y_test_convert, y_test_pred_logistic, plcc_test, rmse_test, srcc_test, krcc_test
+
+def run(args):
+    data_list, srcc_list, krcc_list, plcc_list, rmse_list, select_criteria_list = [], [], [], [], [], []
+
+    os.makedirs(os.path.join(args.report_path, 'fine_tune'), exist_ok=True)
+    if args.is_finetune:
+        csv_name = f'{args.report_path}/fine_tune/{args.test_data_name}_{args.network_name}_{args.select_criteria}_finetune.csv'
+    else:
+        csv_name = f'{args.report_path}/fine_tune/{args.test_data_name}_{args.network_name}_{args.select_criteria}_wo_finetune.csv'
+    print(f'Test dataset: {args.test_data_name}')
+    test_features, test_vids, test_scores, sorted_test_df = process_test_set(args.test_data_name, args.metadata_path, args.feature_path, args.network_name)
+    X_test, y_test = preprocess_data(test_features, test_scores)
+
+    # get save model param
+    model = Mlp(input_features=X_test.shape[1], out_features=1, drop_rate=0.2, act_layer=nn.GELU)
+    model = model.to(device)
+    model_path = os.path.join(args.model_path, f"{args.train_data_name}_{args.network_name}_{args.model_name}_{args.select_criteria}_trained_model_kfold.pth")
+
+    model_results = []
+    for i in range(1, args.n_repeats + 1):
+        print(f"{i}th repeated 80-20 hold out test")
+        X_fine_tune, X_final_test, y_fine_tune, y_final_test = train_test_split(X_test, y_test, test_size=0.2, random_state=math.ceil(8.8 * i))
+        if args.is_finetune:
+            # test fine tuned model on the test dataset
+            ft_model = fine_tune_model(model, device, model_path, X_fine_tune, y_fine_tune, args.report_path, args.batch_size,
+                                       args.epochs, args.loss_type, args.optimizer_type, args.initial_lr, args.weight_decay, args.use_swa, args.l1_w, args.rank_w)
+            df_test_pred, y_test_convert, y_test_pred_logistic, plcc_test, rmse_test, srcc_test, krcc_test = fine_tuned_model_test(ft_model, device, X_final_test, y_final_test, args.test_data_name)
+            best_model = copy.deepcopy(ft_model)
+        else:
+            # without fine tune on the test dataset
+            df_test_pred, y_test_convert, y_test_pred_logistic, plcc_test, rmse_test, srcc_test, krcc_test = wo_fine_tune_model(model, device, model_path, X_test, y_test, args.loss_type, args.test_data_name)
+            print(y_test_pred_logistic)
+            best_model = copy.deepcopy(model)
+
+        model_results.append({
+            'model': best_model,
+            'srcc': srcc_test,
+            'krcc': krcc_test,
+            'plcc': plcc_test,
+            'rmse': rmse_test,
+            'df_pred': df_test_pred
+        })
+        print('\n')
+
+    if args.select_criteria == 'byrmse':
+        sorted_results = sorted(model_results, key=lambda x: x['rmse'])
+    elif args.select_criteria == 'bykrcc':
+        sorted_results = sorted(model_results, key=lambda x: x['krcc'], reverse=True)
+    else:
+        raise ValueError(f"Unknown select_criteria: {args.select_criteria}")
+    median_index = len(sorted_results) // 2
+    median_result = sorted_results[median_index]
+    median_model = median_result['model']
+    median_df_test_pred = median_result['df_pred']
+    median_srcc_test = median_result['srcc']
+    median_krcc_test = median_result['krcc']
+    median_plcc_test = median_result['plcc']
+    median_rmse_test = median_result['rmse']
+    data_list.append(args.test_data_name)
+    srcc_list.append(median_srcc_test)
+    krcc_list.append(median_krcc_test)
+    plcc_list.append(median_plcc_test)
+    rmse_list.append(median_rmse_test)
+    select_criteria_list.append(args.select_criteria)
+    median_df_test_pred.head()
+
+    # save finetuned model
+    if args.is_finetune:
+        model_path_new = os.path.join(args.report_path, f"{args.test_data_name}_{args.network_name}_fine_tuned_model.pth")
+        torch.save(median_model.state_dict(), model_path_new)
+        print(f"Median model select {args.select_criteria} saved to {model_path_new}")
+
+    df_results = create_results_dataframe(data_list, args.network_name, srcc_list, krcc_list, plcc_list, rmse_list, select_criteria_list)
+    print(df_results.T)
+    df_results.to_csv(csv_name, index=None, encoding="UTF-8")
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+
+    # input parameters
+    parser.add_argument('--train_data_name', type=str, default='lsvq_train')
+    parser.add_argument('--test_data_name', type=str, default='finevd')
+    parser.add_argument('--network_name', type=str, default='camp-vqa')
+    parser.add_argument('--model_name', type=str, default='Mlp')
+    parser.add_argument('--select_criteria', type=str, default='byrmse', choices=['byrmse', 'bykrcc'])
+
+    # paths
+    parser.add_argument('--metadata_path', type=str, default='../metadata/')
+    parser.add_argument('--feature_path', type=str, default=None)
+    parser.add_argument('--model_path', type=str, default='../model/')
+    parser.add_argument('--report_path', type=str, default='../log/')
+
+    # training parameters
+    parser.add_argument('--is_finetune', action='store_true', help="Enable fine-tuning")
+    parser.add_argument('--n_repeats', type=int, default=21)
+    parser.add_argument('--batch_size', type=int, default=256)
+    parser.add_argument('--epochs', type=int, default=200)
+
+    # misc
+    parser.add_argument('--loss_type', type=str, default='MAERankLoss')
+    parser.add_argument('--optimizer_type', type=str, default='sgd')
+    parser.add_argument('--initial_lr', type=float, default=1e-2)
+    parser.add_argument('--weight_decay', type=float, default=0.0005)
+    parser.add_argument('--use_swa', type=bool, default=True, help="Enable SWA (default: True)")
+    parser.add_argument('--l1_w', type=float, default=0.6)
+    parser.add_argument('--rank_w', type=float, default=1.0)
+
+    args = parser.parse_args()
+    if args.feature_path is None:
+        args.feature_path = f'../features/{args.network_name}/'
+    print(f"[Paths] metadata: {args.metadata_path}; features: {args.feature_path}; model: {args.model_path}; report: {args.report_path}")
+    run(args)

model_regression.py

@@ -0,0 +1,656 @@
+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 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 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, data, data_name, set_name):
+    try:
+        df = pd.read_csv(csv, skiprows=[], header=None)
+    except Exception as e:
+        logging.error(f'Read CSV file error: {e}')
+        raise
+
+    y_data = df.values[1:, 2].astype(float)
+    y = torch.tensor(y_data, dtype=torch.float32)
+
+    if set_name == 'test':
+        print(f"Modified y_true: {y}")
+    X = data
+    return X, y
+
+def preprocess_data(X, y):
+    X[torch.isnan(X)] = 0
+    X[torch.isinf(X)] = 0
+
+    # MinMaxScaler (use PyTorch implementation)
+    X_min = X.min(dim=0, keepdim=True).values
+    X_max = X.max(dim=0, keepdim=True).values
+    X = (X - X_min) / (X_max - X_min)
+    y = y.view(-1, 1).squeeze()
+    return X, y
+
+# 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
+    if isinstance(y_test, torch.Tensor):
+        mos1 = y_test.numpy()
+    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' or data_name == 'live_yt_gaming' or data_name == "finevd":
+        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}_{select_criteria}.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}_{select_criteria}.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 == 'lsvq_train':
+        train_csv = os.path.join(metadata_path, f'mos_files/{data_name}_MOS_train.csv')
+        test_csv = os.path.join(metadata_path, f'mos_files/{data_name}_MOS_test.csv')
+        X_train, y_train = load_data(train_csv, train_features, data_name, 'train')
+        X_test, y_test = load_data(test_csv, test_features, data_name, 'test')
+
+    else:
+        train_csv = os.path.join(metadata_path, f'mos_files/{data_name}_MOS_train.csv')
+        test_csv = os.path.join(metadata_path, f'mos_files/{data_name}_MOS_test.csv')
+        train_data = torch.load(f'{feature_path}split_train_test/{network_name}_{data_name}_train_features.pt')
+        test_data = torch.load(f'{feature_path}split_train_test/{network_name}_{data_name}_test_features.pt')
+        X_train, y_train = load_data(train_csv, train_data, data_name, 'train')
+        X_test, y_test = load_data(test_csv, test_data, data_name, 'test')
+
+    # standard min-max normalization of training 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 = []
+    y_val_true = []
+    with torch.no_grad():
+        for inputs, targets in val_loader:
+            inputs, targets = inputs.to(device), targets.to(device)
+
+            outputs = model(inputs)
+            y_val_pred.append(outputs)
+            y_val_true.append(targets)
+            loss = criterion(outputs, targets.view(-1, 1))
+            val_loss += loss.item() * inputs.size(0)
+
+    val_loss /= len(val_loader.dataset)
+    y_val_pred = torch.cat(y_val_pred, dim=0)
+    y_val_true = torch.cat(y_val_true, dim=0)
+    return val_loss, y_val_pred, y_val_true
+
+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']
+    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}, 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}')
+
+    # Split data into train and validation
+    X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=0.2, 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 = []
+
+    # initialisation of model, loss function, optimiser
+    model = Mlp(input_features=X_train.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:
+        criterion = 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_min=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(X_train, y_train)
+    val_dataset = TensorDataset(X_val, y_val)
+    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, y_val_true = evaluate(current_model, val_loader, criterion, device)
+        val_losses.append(val_loss)
+        print(f"Epoch {epoch + 1}, Training Loss: {train_loss}, Validation Loss: {val_loss}")
+
+        y_val_pred = torch.cat([pred for pred in y_val_pred])
+        _, _, rmse_val, _, krcc_val = compute_correlation_metrics(y_val.cpu().numpy(), y_val_pred.cpu().numpy())
+        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}:")
+            y_val_pred_logistic_tmp, plcc_valid_tmp, rmse_valid_tmp, srcc_valid_tmp, krcc_valid_tmp = compute_correlation_metrics(y_val.cpu().numpy(), y_val_pred.cpu().numpy())
+            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 = X_train
+            y_tra_pred_tmp = best_model(X_train_fold_tensor).detach().cpu().squeeze()
+            y_tra_pred_logistic_tmp, plcc_train_tmp, rmse_train_tmp, srcc_train_tmp, krcc_train_tmp = compute_correlation_metrics(y_train.cpu().numpy(), y_tra_pred_tmp.cpu().numpy())
+            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)
+            best_model = best_model.to(device)
+            best_model.eval()
+            torch.optim.swa_utils.update_bn(train_loader, best_model)
+
+        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(X, 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}_{model_name}_{select_criteria}.mat'
+    pred_score_filename = os.path.join(score_path, f"{data_name}_{network_name}_{model_name}_Predicted_Score_{select_criteria}.csv")
+    file_path = os.path.join(save_path, f"{data_name}_{network_name}_{model_name}_{select_criteria}_trained_model.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)
+        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 = torch.tensor(list(y_train_pred), dtype=torch.float32)
+        y_train_pred_logistic, plcc_train, rmse_train, srcc_train, krcc_train = compute_correlation_metrics(y_train.cpu().numpy(), y_train_pred.cpu().numpy())
+
+        # test best model on the test dataset
+        y_test_pred = model_test(best_model, X_test, y_test, device)
+        y_test_pred = torch.tensor(list(y_test_pred), dtype=torch.float32)
+        y_test_pred_logistic, plcc_test, rmse_test, srcc_test, krcc_test = compute_correlation_metrics(y_test.cpu().numpy(), y_test_pred.cpu().numpy())
+
+        # 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 = torch.tensor(SRCC_all_repeats_test, dtype=torch.float32)
+    KRCC_all_repeats_test = torch.tensor(KRCC_all_repeats_test, dtype=torch.float32)
+    PLCC_all_repeats_test = torch.tensor(PLCC_all_repeats_test, dtype=torch.float32)
+    PLCC_all_repeats_test = PLCC_all_repeats_test[~torch.isnan(PLCC_all_repeats_test)]
+    RMSE_all_repeats_test = torch.tensor(RMSE_all_repeats_test, dtype=torch.float32)
+    SRCC_all_repeats_train = torch.tensor(SRCC_all_repeats_train, dtype=torch.float32)
+    KRCC_all_repeats_train = torch.tensor(KRCC_all_repeats_train, dtype=torch.float32)
+    PLCC_all_repeats_train = torch.tensor(PLCC_all_repeats_train, dtype=torch.float32)
+    RMSE_all_repeats_train = torch.tensor(RMSE_all_repeats_train, dtype=torch.float32)
+
+    logging.info('======================================================')
+    logging.info('Average training results among all repeated 80-20 holdouts:')
+    logging.info('SRCC: %f (std: %f)', torch.median(SRCC_all_repeats_train).item(), torch.std(SRCC_all_repeats_train).item())
+    logging.info('KRCC: %f (std: %f)', torch.median(KRCC_all_repeats_train).item(), torch.std(KRCC_all_repeats_train).item())
+    logging.info('PLCC: %f (std: %f)', torch.median(PLCC_all_repeats_train).item(), torch.std(PLCC_all_repeats_train).item())
+    logging.info('RMSE: %f (std: %f)', torch.median(RMSE_all_repeats_train).item(), torch.std(RMSE_all_repeats_train).item())
+    logging.info('======================================================')
+    logging.info('Average testing results among all repeated 80-20 holdouts:')
+    logging.info('SRCC: %f (std: %f)', torch.median(SRCC_all_repeats_test).item(), torch.std(SRCC_all_repeats_test).item())
+    logging.info('KRCC: %f (std: %f)', torch.median(KRCC_all_repeats_test).item(), torch.std(KRCC_all_repeats_test).item())
+    logging.info('PLCC: %f (std: %f)', torch.median(PLCC_all_repeats_test).item(), torch.std(PLCC_all_repeats_test).item())
+    logging.info('RMSE: %f (std: %f)', torch.median(RMSE_all_repeats_test).item(), torch.std(RMSE_all_repeats_test).item())
+    logging.info('======================================================')
+    logging.info('\n')
+
+    # find the median model and the index of the median
+    print('======================================================')
+    if select_criteria == 'byrmse':
+        median_metrics = torch.median(RMSE_all_repeats_test).item()
+        indices = (RMSE_all_repeats_test == median_metrics).nonzero(as_tuple=True)[0].tolist()
+        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 = torch.median(KRCC_all_repeats_test).item()
+        indices = (KRCC_all_repeats_test == median_metrics).nonzero(as_tuple=True)[0].tolist()
+        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': SRCC_all_repeats_train.numpy(),
+                                         'KRCC_train': KRCC_all_repeats_train.numpy(),
+                                         'PLCC_train': PLCC_all_repeats_train.numpy(),
+                                         'RMSE_train': RMSE_all_repeats_train.numpy(),
+                                         'SRCC_test': SRCC_all_repeats_test.numpy(),
+                                         'KRCC_test': KRCC_all_repeats_test.numpy(),
+                                         'PLCC_test': PLCC_all_repeats_test.numpy(),
+                                         'RMSE_test': RMSE_all_repeats_test.numpy(),
+                                         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='konvid_1k')
+    parser.add_argument('--network_name', type=str, default='camp-vqa')
+
+    parser.add_argument('--metadata_path', type=str, default='../metadata/')
+    parser.add_argument('--feature_path', type=str, default=f'../features/camp-vqa/')
+    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')
+    parser.add_argument('--n_repeats', type=int, default=21)
+    parser.add_argument('--batch_size', type=int, default=256)
+    parser.add_argument('--epochs', type=int, default=200)
+    parser.add_argument('--hidden_features', type=int, default=256)
+    parser.add_argument('--drop_rate', type=float, default=0.1)
+    # misc
+    parser.add_argument('--loss_type', type=str, default='MAERankLoss')
+    parser.add_argument('--optimizer_type', type=str, default='sgd')
+    parser.add_argument('--initial_lr', type=float, default=1e-2)
+    parser.add_argument('--weight_decay', type=float, default=0.0005)
+    parser.add_argument('--patience', type=int, default=5)
+    parser.add_argument('--use_swa', type=bool, default=True)
+    parser.add_argument('--l1_w', type=float, default=0.6)
+    parser.add_argument('--rank_w', type=float, default=1.0)
+
+    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)

model_regression_lsvq.py

@@ -0,0 +1,666 @@
+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 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, data, data_name, set_name):
+    try:
+        df = pd.read_csv(csv, skiprows=[], header=None)
+    except Exception as e:
+        logging.error(f'Read CSV file error: {e}')
+        raise
+
+    y_data = df.values[1:, 2].astype(float)
+    y = torch.tensor(y_data, dtype=torch.float32)
+
+    if set_name == 'test':
+        print(f"Modified y_true: {y}")
+    X = data
+    return X, y
+
+def preprocess_data(X, y):
+    X[torch.isnan(X)] = 0
+    X[torch.isinf(X)] = 0
+
+    # MinMaxScaler (use PyTorch implementation)
+    X_min = X.min(dim=0, keepdim=True).values
+    X_max = X.max(dim=0, keepdim=True).values
+    X = (X - X_min) / (X_max - X_min)
+    if y is not None:
+        y = y.view(-1, 1).squeeze()
+    return X, y
+
+# 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
+    if isinstance(y_test, torch.Tensor):
+        mos1 = y_test.numpy()
+    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}_{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}_{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 == 'lsvq_train':
+        train_csv = os.path.join(metadata_path, f'mos_files/{data_name}_MOS_train.csv')
+        test_csv = os.path.join(metadata_path, f'mos_files/{data_name}_MOS_test.csv')
+        X_train, y_train = load_data(train_csv, train_features, data_name, 'train')
+        X_test, y_test = load_data(test_csv, test_features, data_name, 'test')
+
+    else:
+        train_csv = os.path.join(metadata_path, f'mos_files/{data_name}_MOS_train.csv')
+        test_csv = os.path.join(metadata_path, f'mos_files/{data_name}_MOS_test.csv')
+        train_data = torch.load(f'{feature_path}split_train_test/{network_name}_{data_name}_train_features.pt')
+        test_data = torch.load(f'{feature_path}split_train_test/{network_name}_{data_name}_test_features.pt')
+        X_train, y_train = load_data(train_csv, train_data, data_name, 'train')
+        X_test, y_test = load_data(test_csv, test_data, data_name, 'test')
+
+    # standard min-max normalization of training 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 = []
+    y_val_true = []
+    with torch.no_grad():
+        for inputs, targets in val_loader:
+            inputs, targets = inputs.to(device), targets.to(device)
+
+            outputs = model(inputs)
+            y_val_pred.append(outputs)
+            y_val_true.append(targets)
+            loss = criterion(outputs, targets.view(-1, 1))
+            val_loss += loss.item() * inputs.size(0)
+
+    val_loss /= len(val_loader.dataset)
+    y_val_pred = torch.cat(y_val_pred, dim=0)
+    y_val_true = torch.cat(y_val_true, dim=0)
+    return val_loss, y_val_pred, y_val_true
+
+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:
+            criterion = 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(X_train_fold, y_train_fold)
+        val_dataset = TensorDataset(X_val_fold, 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, y_val_true = 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 = torch.cat([pred for pred in y_val_pred])
+            _, _, rmse_val, _, krcc_val = compute_correlation_metrics(y_val_fold.cpu().numpy(), y_val_pred.cpu().numpy())
+            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.cpu().numpy(), y_val_pred.cpu().numpy())
+                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 = X_train_fold
+                y_tra_pred_tmp = best_model(X_train_fold_tensor).detach().cpu().squeeze()
+                y_tra_pred_logistic_tmp, plcc_train_tmp, rmse_train_tmp, srcc_train_tmp, krcc_train_tmp = compute_correlation_metrics(y_train_fold.cpu().numpy(), y_tra_pred_tmp.cpu().numpy())
+                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)
+            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(X, 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}_{model_name}_{select_criteria}_kfold.mat'
+    pred_score_filename = os.path.join(score_path, f"{data_name}_{network_name}_{model_name}_Predicted_Score_{select_criteria}_kfold.csv")
+    file_path = os.path.join(save_path, f"{data_name}_{network_name}_{model_name}_{select_criteria}_trained_model_kfold.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)
+        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 = torch.tensor(list(y_train_pred), dtype=torch.float32)
+        y_train_pred_logistic, plcc_train, rmse_train, srcc_train, krcc_train = compute_correlation_metrics(y_train.cpu().numpy(), y_train_pred.cpu().numpy())
+
+        # test best model on the test dataset
+        y_test_pred = model_test(best_model, X_test, y_test, device)
+        y_test_pred = torch.tensor(list(y_test_pred), dtype=torch.float32)
+        y_test_pred_logistic, plcc_test, rmse_test, srcc_test, krcc_test = compute_correlation_metrics(y_test.cpu().numpy(), y_test_pred.cpu().numpy())
+
+        # 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 = torch.tensor(SRCC_all_repeats_test, dtype=torch.float32)
+    KRCC_all_repeats_test = torch.tensor(KRCC_all_repeats_test, dtype=torch.float32)
+    PLCC_all_repeats_test = torch.tensor(PLCC_all_repeats_test, dtype=torch.float32)
+    RMSE_all_repeats_test = torch.tensor(RMSE_all_repeats_test, dtype=torch.float32)
+    SRCC_all_repeats_train = torch.tensor(SRCC_all_repeats_train, dtype=torch.float32)
+    KRCC_all_repeats_train = torch.tensor(KRCC_all_repeats_train, dtype=torch.float32)
+    PLCC_all_repeats_train = torch.tensor(PLCC_all_repeats_train, dtype=torch.float32)
+    RMSE_all_repeats_train = torch.tensor(RMSE_all_repeats_train, dtype=torch.float32)
+
+    logging.info('======================================================')
+    logging.info('Average training results among all repeated 80-20 holdouts:')
+    logging.info('SRCC: %f (std: %f)', torch.median(SRCC_all_repeats_train).item(), torch.std(SRCC_all_repeats_train).item())
+    logging.info('KRCC: %f (std: %f)', torch.median(KRCC_all_repeats_train).item(), torch.std(KRCC_all_repeats_train).item())
+    logging.info('PLCC: %f (std: %f)', torch.median(PLCC_all_repeats_train).item(), torch.std(PLCC_all_repeats_train).item())
+    logging.info('RMSE: %f (std: %f)', torch.median(RMSE_all_repeats_train).item(), torch.std(RMSE_all_repeats_train).item())
+    logging.info('======================================================')
+    logging.info('Average testing results among all repeated 80-20 holdouts:')
+    logging.info('SRCC: %f (std: %f)', torch.median(SRCC_all_repeats_test).item(), torch.std(SRCC_all_repeats_test).item())
+    logging.info('KRCC: %f (std: %f)', torch.median(KRCC_all_repeats_test).item(), torch.std(KRCC_all_repeats_test).item())
+    logging.info('PLCC: %f (std: %f)', torch.median(PLCC_all_repeats_test).item(), torch.std(PLCC_all_repeats_test).item())
+    logging.info('RMSE: %f (std: %f)', torch.median(RMSE_all_repeats_test).item(), torch.std(RMSE_all_repeats_test).item())
+    logging.info('======================================================')
+    logging.info('\n')
+
+    # find the median model and the index of the median
+    print('======================================================')
+    if select_criteria == 'byrmse':
+        median_metrics = torch.median(RMSE_all_repeats_test).item()
+        indices = (RMSE_all_repeats_test == median_metrics).nonzero(as_tuple=True)[0].tolist()
+        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 = torch.median(KRCC_all_repeats_test).item()
+        indices = (KRCC_all_repeats_test == median_metrics).nonzero(as_tuple=True)[0].tolist()
+        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': SRCC_all_repeats_train.numpy(),
+                                         'KRCC_train': KRCC_all_repeats_train.numpy(),
+                                         'PLCC_train': PLCC_all_repeats_train.numpy(),
+                                         'RMSE_train': RMSE_all_repeats_train.numpy(),
+                                         'SRCC_test': SRCC_all_repeats_test.numpy(),
+                                         'KRCC_test': KRCC_all_repeats_test.numpy(),
+                                         'PLCC_test': PLCC_all_repeats_test.numpy(),
+                                         'RMSE_test': RMSE_all_repeats_test.numpy(),
+                                         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')
+    parser.add_argument('--network_name', type=str, default='camp-vqa')
+
+    parser.add_argument('--metadata_path', type=str, default='../metadata/')
+    parser.add_argument('--feature_path', type=str, default=f'../features/camp-vqa/')
+    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')
+    parser.add_argument('--n_repeats', type=int, default=21)
+    parser.add_argument('--n_splits', type=int, default=10)
+    parser.add_argument('--batch_size', type=int, default=256)
+    parser.add_argument('--epochs', type=int, default=50)
+    parser.add_argument('--hidden_features', type=int, default=256)
+    parser.add_argument('--drop_rate', type=float, default=0.1)
+    # misc
+    parser.add_argument('--loss_type', type=str, default='MAERankLoss')
+    parser.add_argument('--optimizer_type', type=str, default='sgd')
+    parser.add_argument('--initial_lr', type=float, default=1e-1)
+    parser.add_argument('--weight_decay', type=float, default=0.005)
+    parser.add_argument('--patience', type=int, default=5)
+    parser.add_argument('--use_swa', type=bool, default=True)
+    parser.add_argument('--l1_w', type=float, default=0.6)
+    parser.add_argument('--rank_w', type=float, default=1.0)
+
+    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)

requirements.txt

@@ -0,0 +1,62 @@
+accelerate==1.6.0
+av==14.3.0
+certifi==2025.1.31
+charset-normalizer==3.4.1
+clip @ git+https://github.com/openai/CLIP.git@dcba3cb2e2827b402d2701e7e1c7d9fed8a20ef1
+ffmpeg==1.4
+filelock==3.18.0
+fsspec==2025.3.2
+ftfy==6.3.1
+fvcore==0.1.5.post20221221
+huggingface-hub==0.30.1
+idna==3.10
+iopath==0.1.10
+Jinja2==3.1.6
+joblib==1.4.2
+MarkupSafe==3.0.2
+matplotlib==3.10.0
+mpmath==1.3.0
+networkx==3.4.2
+nvidia-cublas-cu12==12.4.5.8
+nvidia-cuda-cupti-cu12==12.4.127
+nvidia-cuda-nvrtc-cu12==12.4.127
+nvidia-cuda-runtime-cu12==12.4.127
+nvidia-cudnn-cu12==9.1.0.70
+nvidia-cufft-cu12==11.2.1.3
+nvidia-curand-cu12==10.3.5.147
+nvidia-cusolver-cu12==11.6.1.9
+nvidia-cusparse-cu12==12.3.1.170
+nvidia-cusparselt-cu12==0.6.2
+nvidia-nccl-cu12==2.21.5
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-nvtx-cu12==12.4.127
+opencv-python==4.11.0.86
+parameterized==0.9.0
+peft==0.15.2
+portalocker==3.1.1
+psutil==7.0.0
+PyQt6==6.7.1
+pytorchvideo==0.1.5
+PyYAML==6.0.2
+regex==2024.11.6
+requests==2.32.3
+safetensors==0.5.3
+scikit-learn==1.6.1
+scipy==1.15.2
+seaborn==0.13.2
+sentencepiece==0.2.0
+sympy==1.13.1
+tabulate==0.9.0
+termcolor==3.1.0
+threadpoolctl==3.6.0
+timm==1.0.15
+tokenizers==0.21.1
+torch==2.6.0
+torchvision==0.21.0
+tqdm==4.67.1
+transformers==4.50.3
+triton==3.2.0
+typing_extensions==4.13.0
+urllib3==2.3.0
+wcwidth==0.2.13
+yacs==0.1.8