project files

.gitignore

@@ -0,0 +1,3 @@
+.gitignore
+.DS_Store
+test_videos/

README.md

@@ -0,0 +1,198 @@
+# FGSVQA
+![visitors](https://visitor-badge.laobi.icu/badge?page_id=xinyiW915/FGSVQA) 
+![GitHub Repo stars](https://img.shields.io/github/stars/xinyiW915/FGSVQA?logo=github)
+![Python](https://img.shields.io/badge/Python-3.8+-blue)
+[![arXiv](https://img.shields.io/badge/arXiv-2605.20016-b31b1b.svg)](http://arxiv.org/abs/2605.20016)
+
+Official Code for the following paper:
+
+**X. Wang, A. Katsenou, J.Shen and D. Bull**. [FGSVQA: Frequency-Guided Short-form Video Quality Assessment](http://arxiv.org/abs/2605.20016)
+
+[Our paper]() was accepted by the 18th International Conference on Quality of Multimedia Experience ([QoMEX 2026](https://qomex2026.itec.aau.at/)).
+
+---
+
+## Performance
+We validated our proposed method on two publicly available Short-form UGC datasets: KVQ and YouTube SFV+HDR dataset (YT-SFV).
+
+#### **Spearman’s Rank Correlation Coefficient (SRCC)**
+| **Model**                   | **KVQ**   | **YT-SFV (SDR)** | **YT-SFV (HDR2SDR)** |
+|----------------------------|-----------|------------------|----------------------|
+| FGSVQA  | 0.877    | 0.788           | 0.543           |
+
+#### **Pearson’s Linear Correlation Coefficient (PLCC)**  
+| **Model**                   | **KVQ**   | **YT-SFV (SDR)** | **YT-SFV (HDR2SDR)** |
+|----------------------------|-----------|------------------|----------------------|
+| FGSVQA  | 0.878    | 0.818            | 0.666           |
+
+#### **GPU runtime comparison (averaged over 10 runs) across different spatial resolutions on "SDR\_Animal\_5ngj.mp4".**
+| Method | Time(s)<br>540P | Time(s)<br>720P | Time(s)<br>1080P | Time(s)<br>2160P | Ground truth: 4.308<br>Predicted Score|
+|---|------------:|------------:|-------------:|---:|---:|
+| Fast-VQA |       0.599 |       0.673 |        0.909 | 2.217 | 3.319 |
+| FasterVQA |       0.489 |       0.547 |    **0.696** | **1.343** | 3.556 |
+| DOVER |       0.920 |       1.022 |        1.293 | 2.783 | 3.814 |
+| FGSVQA |   **0.313** |   **0.405** |        0.697 | 2.137 | **3.878** |
+
+More results can be found in **[correlation_result.ipynb](https://github.com/xinyiW915/FGSVQA/blob/main/src/correlation_result.ipynb)**.
+
+## Proposed Model
+Overview of the proposed model with the two branches: the frequency-guided weight map and the CLIP vision encoder.
+
+<img src="./SVQA.png" alt="proposed_FGSVQA_framework" width="800"/>
+
+## Usage
+### 📌 Install Requirement
+The repository is built with **Python 3.10** and can be installed via the following commands:
+
+```shell
+git clone https://github.com/xinyiW915/FGSVQA.git
+cd FGSVQA
+conda create -n fgsvqa python=3.10 -y
+conda activate fgsvqa
+pip install -r requirements.txt  
+```
+
+### 📥 Download UGC Datasets
+The corresponding UGC video datasets can be downloaded from the following sources:  
+[KVQ](https://lixinustc.github.io/projects/KVQ/), [YouTube SFV+HDR](https://media.withyoutube.com/sfv-hdr).  
+
+The metadata for the experimented UGC dataset is available under [`./metadata`](./metadata).  
+
+### 🎬 Test Demo
+Run the pre-trained model to evaluate the perceptual quality of a single video. The demo script reports the predicted quality score, runtime, and model complexity.
+
+The model checkpoint should be provided through `--ckpt_path`. Please use a full checkpoint file, such as `qd_model.best.pt`, which contains the saved model weights together with the training MOS mean and standard deviation.
+
+To evaluate a single video, run:
+```shell
+python demo_test.py \
+    --ckpt_path <MODEL_PATH> \
+    --db_path <VIDEO_FOLDER> \
+    --video_id <VIDEO_ID> \
+    --device <DEVICE>
+````
+For example:
+```shell
+python demo_test.py \
+    --ckpt_path ./checkpoints/lsvq/qd_model.best.pt \
+    --db_path ./test_videos/ \
+    --video_id SDR_Animal_5ngj \
+    --device cuda
+```
+
+### 🔁 Cross-Dataset Evaluation
+To evaluate a trained model on another dataset, use `transfer_test_only.py`. This script loads a trained checkpoint, reports the evaluation metrics, and saves the prediction results to a CSV file.
+
+Run:
+```shell
+python transfer_test_only.py \
+    --ckpt_path <MODEL_PATH> \
+    --csv_path <TEST_METADATA_CSV> \
+    --db_path <TEST_VIDEO_FOLDER> \
+    --device <DEVICE> \
+    --save_pred_csv <SAVE_PREDICTION_CSV>
+```
+For example:
+```shell
+python transfer_test_only.py \
+    --ckpt_path ./checkpoints/lsvq/qd_model.best.pt \
+    --csv_path ./metadata/KVQ_metadata.csv \
+    --db_path /path/to/KVQ/videos \
+    --device cuda \
+    --save_pred_csv /path/to/transfer_test_only_konvid_1k.csv
+```
+
+## Training
+Steps to train and fine-tune the model on different datasets.
+
+### Train Model
+Train the model using the metadata CSV file and the corresponding video folder. The metadata CSV file should contain `vid` and `mos` columns.
+
+```shell
+python train.py \
+    --csv_path <TRAIN_METADATA_CSV> \
+    --db_path <VIDEO_FOLDER> \
+    --save_dir <SAVE_DIR> \
+    --save_name qd_model.pt \
+    --device <DEVICE> \
+    --finetune_last_stage
+```
+For example:
+```shell
+python train.py \
+    --csv_path ./metadata/KVQ_TRAIN_metadata.csv \
+    --db_path /path/to/KVQ/videos \
+    --save_dir ./checkpoints/kvq \
+    --save_name qd_model.pt \
+    --device cuda \
+    --finetune_last_stage
+```
+The script saves the latest checkpoint and the best-performing checkpoint according to the validation SRCC.
+
+### Transfer Model
+To fine-tune a pre-trained model on a new dataset, run:
+
+```shell
+python transfer.py \
+    --mode finetune \
+    --pretrained <PRETRAINED_MODEL_PATH> \
+    --csv_path <TARGET_METADATA_CSV> \
+    --db_path <TARGET_VIDEO_FOLDER> \
+    --save_dir <SAVE_DIR> \
+    --save_name transfer.pt \
+    --device <DEVICE> \
+    --finetune_last_stage
+```
+For example:
+```shell
+python transfer.py \
+    --mode finetune \
+    --pretrained ./checkpoints/shorts-hdr-dataset_sdr/qd_model.best.pt \
+    --csv_path ./metadata/KVQ_TRAIN_metadata.csv \
+    --db_path /path/to/KVQ/videos \
+    --save_dir ./checkpoints_transfer/kvq \
+    --save_name transfer.pt \
+    --device cuda \
+    --finetune_last_stage
+```
+
+### Test Only
+To directly test a pre-trained model on another dataset, run:
+
+```shell
+python transfer.py \
+    --mode test_only \
+    --pretrained <PRETRAINED_MODEL_PATH> \
+    --csv_path <TEST_METADATA_CSV> \
+    --db_path <TEST_VIDEO_FOLDER> \
+    --device <DEVICE>
+```
+For example:
+```shell
+python transfer.py \
+    --mode test_only \
+    --pretrained ./checkpoints/shorts-hdr-dataset_sdr/qd_model.best.pt \
+    --csv_path ./metadata/KVQ_metadata.csv \
+    --db_path /path/to/KVQ/videos \
+    --device cuda
+```
+
+
+## Acknowledgment
+This work was funded by the UKRI MyWorld Strength in Places Programme (SIPF00006/1) as part of my PhD study.
+
+## Citation
+If you find this paper and the repo useful, please cite our paper 😊:
+
+```bibtex
+@article{wang2026fgsvqa,
+      title={FGSVQA: Frequency-Guided Short-form Video Quality Assessment},
+      author={Wang, Xinyi and Katsenou, Angeliki, Shen, Junxiao and Bull, David},
+      booktitle={2026 18th International Conference on Quality of Multimedia Experience (QoMEX)}, 
+      year={2026},
+      organization={IEEE}
+}
+```
+
+## Contact:
+Xinyi WANG, ```xinyi.wang@bristol.ac.uk```

requirements.txt

@@ -0,0 +1,75 @@
+accelerate==1.12.0
+annotated-doc==0.0.4
+anyio==4.12.1
+av==16.1.0
+bitsandbytes==0.49.1
+certifi==2026.1.4
+click==8.3.1
+contourpy==1.3.2
+cuda-bindings==12.9.4
+cuda-pathfinder==1.3.3
+cycler==0.12.1
+decord==0.6.0
+einops==0.8.2
+exceptiongroup==1.3.1
+filelock==3.20.0
+fonttools==4.61.1
+fsspec==2025.12.0
+h11==0.16.0
+hf-xet==1.2.0
+httpcore==1.0.9
+httpx==0.28.1
+huggingface_hub==1.4.1
+idna==3.11
+Jinja2==3.1.6
+kiwisolver==1.4.9
+mamba-ssm==2.3.0
+markdown-it-py==4.0.0
+MarkupSafe==2.1.5
+matplotlib==3.10.8
+mdurl==0.1.2
+mpmath==1.3.0
+networkx==3.4.2
+ninja==1.13.0
+numpy==2.2.6
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==9.1.0.70
+nvidia-cufft-cu12==11.0.2.54
+nvidia-cufile-cu12==1.13.1.3
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-cusparselt-cu12==0.7.1
+nvidia-nccl-cu12==2.21.5
+nvidia-nvjitlink-cu12==12.8.93
+nvidia-nvshmem-cu12==3.4.5
+nvidia-nvtx-cu12==12.1.105
+opencv-python==4.13.0.92
+packaging @ file:///home/task_176104877067765/conda-bld/packaging_1761049113113/work
+pillow==12.1.1
+psutil==7.2.2
+Pygments==2.19.2
+pyparsing==3.3.2
+python-dateutil==2.9.0.post0
+PyYAML==6.0.3
+regex==2026.1.15
+rich==14.3.2
+safetensors==0.7.0
+scipy==1.15.3
+shellingham==1.5.4
+six==1.17.0
+sympy==1.13.1
+tokenizers==0.22.2
+torch==2.5.1+cu121
+torchaudio==2.5.1+cu121
+torchvision==0.20.1+cu121
+tqdm==4.67.3
+transformers==5.1.0
+triton==3.1.0
+typer==0.23.0
+typer-slim==0.23.0
+typing_extensions==4.15.0
+xformers==0.0.34

src/correlation_result.ipynb

@@ -0,0 +1,914 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "id": "initial_id",
+   "metadata": {
+    "collapsed": true,
+    "ExecuteTime": {
+     "end_time": "2026-05-19T12:42:02.280702Z",
+     "start_time": "2026-05-19T12:42:02.276769Z"
+    }
+   },
+   "source": [
+    "from pathlib import Path\n",
+    "import re\n",
+    "import pandas as pd"
+   ],
+   "outputs": [],
+   "execution_count": 19
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": "Intra-dataset evaluation: Train on and test on the same <dataset>",
+   "id": "145ddc0eb7f59e45"
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-05-19T12:42:02.291070Z",
+     "start_time": "2026-05-19T12:42:02.288978Z"
+    }
+   },
+   "cell_type": "code",
+   "source": "LOG_ROOT = Path(\"./checkpoints\")",
+   "id": "eba6249c8ad34da8",
+   "outputs": [],
+   "execution_count": 20
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-05-19T12:42:02.315625Z",
+     "start_time": "2026-05-19T12:42:02.302579Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "filename_pattern = re.compile(r\"^train_(?P<dataset>.+)\\.log$\")\n",
+    "result_pattern = re.compile(\n",
+    "    r\"TEST\\s*\\|\\s*\"\n",
+    "    r\"loss=(?P<loss>[-+]?\\d*\\.?\\d+)\\s+\"\n",
+    "    r\"plcc=(?P<plcc>[-+]?\\d*\\.?\\d+)\\s+\"\n",
+    "    r\"srcc=(?P<srcc>[-+]?\\d*\\.?\\d+)\\s+\"\n",
+    "    r\"rmse=(?P<rmse>[-+]?\\d*\\.?\\d+)\"\n",
+    ")\n",
+    "records = []\n",
+    "for log_path in LOG_ROOT.rglob(\"*.log\"):\n",
+    "    filename_match = filename_pattern.match(log_path.name)\n",
+    "    dataset = filename_match.group(\"dataset\")\n",
+    "\n",
+    "    text = log_path.read_text(encoding=\"utf-8\", errors=\"ignore\")\n",
+    "    matches = list(result_pattern.finditer(text))\n",
+    "    if not matches:\n",
+    "        continue\n",
+    "    records.append({\n",
+    "        \"dataset\": dataset,\n",
+    "        \"srcc\": float(matches[-1].group(\"srcc\")),\n",
+    "        \"plcc\": float(matches[-1].group(\"plcc\")),\n",
+    "    })\n",
+    "train_df = (pd.DataFrame(records).sort_values(\"dataset\", ignore_index=True)[[\"dataset\", \"srcc\", \"plcc\"]])\n",
+    "train_df"
+   ],
+   "id": "4f2532e62a94f758",
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "                      dataset    srcc    plcc\n",
+       "0                      finevd  0.8596  0.8579\n",
+       "1                   konvid_1k  0.7970  0.8171\n",
+       "2                         kvq  0.8768  0.8781\n",
+       "3                    live_vqc  0.6386  0.7192\n",
+       "4              live_yt_gaming  0.8554  0.8755\n",
+       "5                        lsvq  0.8753  0.8760\n",
+       "6  shorts-hdr-dataset_hdr2sdr  0.5431  0.6657\n",
+       "7      shorts-hdr-dataset_sdr  0.7884  0.8182\n",
+       "8                 youtube_ugc  0.8325  0.8548"
+      ],
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>dataset</th>\n",
+       "      <th>srcc</th>\n",
+       "      <th>plcc</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>finevd</td>\n",
+       "      <td>0.8596</td>\n",
+       "      <td>0.8579</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>konvid_1k</td>\n",
+       "      <td>0.7970</td>\n",
+       "      <td>0.8171</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>kvq</td>\n",
+       "      <td>0.8768</td>\n",
+       "      <td>0.8781</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>live_vqc</td>\n",
+       "      <td>0.6386</td>\n",
+       "      <td>0.7192</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>live_yt_gaming</td>\n",
+       "      <td>0.8554</td>\n",
+       "      <td>0.8755</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5</th>\n",
+       "      <td>lsvq</td>\n",
+       "      <td>0.8753</td>\n",
+       "      <td>0.8760</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6</th>\n",
+       "      <td>shorts-hdr-dataset_hdr2sdr</td>\n",
+       "      <td>0.5431</td>\n",
+       "      <td>0.6657</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>7</th>\n",
+       "      <td>shorts-hdr-dataset_sdr</td>\n",
+       "      <td>0.7884</td>\n",
+       "      <td>0.8182</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>8</th>\n",
+       "      <td>youtube_ugc</td>\n",
+       "      <td>0.8325</td>\n",
+       "      <td>0.8548</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ]
+     },
+     "execution_count": 21,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "execution_count": 21
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": "Cross-dataset evaluation: Train on <trainOn> dataset and test on \"testOn\" dataset",
+   "id": "68dbaf20d93a49ed"
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-05-19T12:42:02.318434Z",
+     "start_time": "2026-05-19T12:42:02.316713Z"
+    }
+   },
+   "cell_type": "code",
+   "source": "LOG_ROOT = Path(\"./checkpoints_transfer\")",
+   "id": "e8583a08cc496a7c",
+   "outputs": [],
+   "execution_count": 22
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-05-19T12:42:02.328167Z",
+     "start_time": "2026-05-19T12:42:02.319595Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "filename_pattern = re.compile(r\"^transfer_(?P<trainOn>.+?)_test_on_(?P<testOn>.+?)\\.log$\")\n",
+    "result_pattern = re.compile(\n",
+    "    r\"TEST_ONLY\\s*\\|\\s*\"\n",
+    "    r\"loss=(?P<loss>[-+]?\\d*\\.?\\d+)\\s+\"\n",
+    "    r\"plcc=(?P<plcc>[-+]?\\d*\\.?\\d+)\\s+\"\n",
+    "    r\"srcc=(?P<srcc>[-+]?\\d*\\.?\\d+)\\s+\"\n",
+    "    r\"rmse=(?P<rmse>[-+]?\\d*\\.?\\d+)\"\n",
+    ")\n",
+    "records = []\n",
+    "for log_path in LOG_ROOT.rglob(\"transfer_*_test_on_*.log\"):\n",
+    "    filename_match = filename_pattern.match(log_path.name)\n",
+    "    trainOn = filename_match.group(\"trainOn\")\n",
+    "    testOn = filename_match.group(\"testOn\")\n",
+    "\n",
+    "    text = log_path.read_text(encoding=\"utf-8\", errors=\"ignore\")\n",
+    "    matches = list(result_pattern.finditer(text))\n",
+    "    if not matches:\n",
+    "        continue\n",
+    "    records.append({\n",
+    "            \"trainOn\": trainOn,\n",
+    "            \"testOn\": testOn,\n",
+    "            \"srcc\": float(matches[-1].group(\"srcc\")),\n",
+    "            \"plcc\": float(matches[-1].group(\"plcc\")),\n",
+    "    })\n",
+    "transfer_df = (pd.DataFrame(records).sort_values([\"trainOn\", \"testOn\"], ignore_index=True))\n",
+    "transfer_df"
+   ],
+   "id": "b196c863f89e0e0d",
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "                      trainOn                      testOn    srcc    plcc\n",
+       "0                         kvq  shorts-hdr-dataset_hdr2sdr  0.4761  0.5885\n",
+       "1                         kvq      shorts-hdr-dataset_sdr  0.7548  0.8065\n",
+       "2  shorts-hdr-dataset_hdr2sdr                         kvq  0.5976  0.5693\n",
+       "3  shorts-hdr-dataset_hdr2sdr      shorts-hdr-dataset_sdr  0.6166  0.6802\n",
+       "4      shorts-hdr-dataset_sdr                         kvq  0.7342  0.7451\n",
+       "5      shorts-hdr-dataset_sdr  shorts-hdr-dataset_hdr2sdr  0.5447  0.6958"
+      ],
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>trainOn</th>\n",
+       "      <th>testOn</th>\n",
+       "      <th>srcc</th>\n",
+       "      <th>plcc</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>kvq</td>\n",
+       "      <td>shorts-hdr-dataset_hdr2sdr</td>\n",
+       "      <td>0.4761</td>\n",
+       "      <td>0.5885</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>kvq</td>\n",
+       "      <td>shorts-hdr-dataset_sdr</td>\n",
+       "      <td>0.7548</td>\n",
+       "      <td>0.8065</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>shorts-hdr-dataset_hdr2sdr</td>\n",
+       "      <td>kvq</td>\n",
+       "      <td>0.5976</td>\n",
+       "      <td>0.5693</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>shorts-hdr-dataset_hdr2sdr</td>\n",
+       "      <td>shorts-hdr-dataset_sdr</td>\n",
+       "      <td>0.6166</td>\n",
+       "      <td>0.6802</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>shorts-hdr-dataset_sdr</td>\n",
+       "      <td>kvq</td>\n",
+       "      <td>0.7342</td>\n",
+       "      <td>0.7451</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5</th>\n",
+       "      <td>shorts-hdr-dataset_sdr</td>\n",
+       "      <td>shorts-hdr-dataset_hdr2sdr</td>\n",
+       "      <td>0.5447</td>\n",
+       "      <td>0.6958</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ]
+     },
+     "execution_count": 23,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "execution_count": 23
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": "Cross-dataset evaluation: Train on \"trainOn\" dataset and finetune on \"fintuneOn\" dataset",
+   "id": "e08409c56d2eaea3"
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-05-19T12:42:02.337383Z",
+     "start_time": "2026-05-19T12:42:02.329308Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "filename_pattern = re.compile(r\"^transfer_(?P<trainOn>.+?)_finetune_on_(?P<finetuneOn>.+?)\\.log$\")\n",
+    "finetune_result_pattern = re.compile(\n",
+    "    r\"FINETUNE TEST\\s*\\|\\s*\"\n",
+    "    r\"loss=(?P<loss>[-+]?\\d*\\.?\\d+)\\s+\"\n",
+    "    r\"plcc=(?P<plcc>[-+]?\\d*\\.?\\d+)\\s+\"\n",
+    "    r\"srcc=(?P<srcc>[-+]?\\d*\\.?\\d+)\\s+\"\n",
+    "    r\"rmse=(?P<rmse>[-+]?\\d*\\.?\\d+)\"\n",
+    ")\n",
+    "records = []\n",
+    "for log_path in LOG_ROOT.rglob(\"transfer_*_finetune_on_*.log\"):\n",
+    "    filename_match = filename_pattern.match(log_path.name)\n",
+    "    trainOn = filename_match.group(\"trainOn\")\n",
+    "    finetuneOn = filename_match.group(\"finetuneOn\")\n",
+    "\n",
+    "    text = log_path.read_text(encoding=\"utf-8\", errors=\"ignore\")\n",
+    "    matches = list(finetune_result_pattern.finditer(text))\n",
+    "    if not matches:\n",
+    "        continue\n",
+    "    records.append({\n",
+    "            \"trainOn\": trainOn,\n",
+    "            \"finetuneOn\": finetuneOn,\n",
+    "            \"srcc\": float(matches[-1].group(\"srcc\")),\n",
+    "            \"plcc\": float(matches[-1].group(\"plcc\")),\n",
+    "    })\n",
+    "finetune_df = (pd.DataFrame(records).sort_values([\"trainOn\", \"finetuneOn\"], ignore_index=True))\n",
+    "finetune_df"
+   ],
+   "id": "bfe4a5008377348d",
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "                      trainOn                  finetuneOn    srcc    plcc\n",
+       "0                         kvq  shorts-hdr-dataset_hdr2sdr  0.6412  0.7225\n",
+       "1                         kvq      shorts-hdr-dataset_sdr  0.8291  0.8683\n",
+       "2  shorts-hdr-dataset_hdr2sdr                         kvq  0.8743  0.8792\n",
+       "3  shorts-hdr-dataset_hdr2sdr      shorts-hdr-dataset_sdr  0.8183  0.8561\n",
+       "4      shorts-hdr-dataset_sdr                         kvq  0.8862  0.8883\n",
+       "5      shorts-hdr-dataset_sdr  shorts-hdr-dataset_hdr2sdr  0.6589  0.8013"
+      ],
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>trainOn</th>\n",
+       "      <th>finetuneOn</th>\n",
+       "      <th>srcc</th>\n",
+       "      <th>plcc</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>kvq</td>\n",
+       "      <td>shorts-hdr-dataset_hdr2sdr</td>\n",
+       "      <td>0.6412</td>\n",
+       "      <td>0.7225</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>kvq</td>\n",
+       "      <td>shorts-hdr-dataset_sdr</td>\n",
+       "      <td>0.8291</td>\n",
+       "      <td>0.8683</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>shorts-hdr-dataset_hdr2sdr</td>\n",
+       "      <td>kvq</td>\n",
+       "      <td>0.8743</td>\n",
+       "      <td>0.8792</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>shorts-hdr-dataset_hdr2sdr</td>\n",
+       "      <td>shorts-hdr-dataset_sdr</td>\n",
+       "      <td>0.8183</td>\n",
+       "      <td>0.8561</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>shorts-hdr-dataset_sdr</td>\n",
+       "      <td>kvq</td>\n",
+       "      <td>0.8862</td>\n",
+       "      <td>0.8883</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5</th>\n",
+       "      <td>shorts-hdr-dataset_sdr</td>\n",
+       "      <td>shorts-hdr-dataset_hdr2sdr</td>\n",
+       "      <td>0.6589</td>\n",
+       "      <td>0.8013</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ]
+     },
+     "execution_count": 24,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "execution_count": 24
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": "Complexity_test",
+   "id": "544f3bed13c8f3d9"
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-05-19T12:42:02.339580Z",
+     "start_time": "2026-05-19T12:42:02.337969Z"
+    }
+   },
+   "cell_type": "code",
+   "source": "complexity_csv_path = \"../test_videos/complexity_test/\"",
+   "id": "9d718a191ee6bdf2",
+   "outputs": [],
+   "execution_count": 25
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-05-19T12:42:02.347488Z",
+     "start_time": "2026-05-19T12:42:02.340191Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "resolution_df1 = pd.read_csv(f\"{complexity_csv_path}complexity_resolution_KVQ_0546.csv\").assign(video=\"KVQ_0546\")\n",
+    "resolution_df2 = pd.read_csv(f\"{complexity_csv_path}complexity_resolution_SDR_Animal_5ngj.csv\").assign(video=\"SDR_Animal_5ngj\")\n",
+    "resolution_df = pd.concat([resolution_df1, resolution_df2], ignore_index=True)\n",
+    "resolution_df"
+   ],
+   "id": "5c99ce5680f78c34",
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "       model   540P   720P  1080P  2160P    MOS            video\n",
+       "0   FAST-VQA  0.638  0.745  1.020  2.680  4.017         KVQ_0546\n",
+       "1  FasterVQA  0.428  0.480  0.573  1.077  4.356         KVQ_0546\n",
+       "2      DOVER  1.018  1.163  1.511  3.662  4.223         KVQ_0546\n",
+       "3        Our  0.364  0.485  0.753  2.437  3.650         KVQ_0546\n",
+       "4   FAST-VQA  0.599  0.673  0.909  2.217  3.319  SDR_Animal_5ngj\n",
+       "5  FasterVQA  0.489  0.547  0.696  1.343  3.556  SDR_Animal_5ngj\n",
+       "6      DOVER  0.920  1.022  1.293  2.783  3.814  SDR_Animal_5ngj\n",
+       "7        Our  0.313  0.405  0.697  2.137  3.878  SDR_Animal_5ngj"
+      ],
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>model</th>\n",
+       "      <th>540P</th>\n",
+       "      <th>720P</th>\n",
+       "      <th>1080P</th>\n",
+       "      <th>2160P</th>\n",
+       "      <th>MOS</th>\n",
+       "      <th>video</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>FAST-VQA</td>\n",
+       "      <td>0.638</td>\n",
+       "      <td>0.745</td>\n",
+       "      <td>1.020</td>\n",
+       "      <td>2.680</td>\n",
+       "      <td>4.017</td>\n",
+       "      <td>KVQ_0546</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>FasterVQA</td>\n",
+       "      <td>0.428</td>\n",
+       "      <td>0.480</td>\n",
+       "      <td>0.573</td>\n",
+       "      <td>1.077</td>\n",
+       "      <td>4.356</td>\n",
+       "      <td>KVQ_0546</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>DOVER</td>\n",
+       "      <td>1.018</td>\n",
+       "      <td>1.163</td>\n",
+       "      <td>1.511</td>\n",
+       "      <td>3.662</td>\n",
+       "      <td>4.223</td>\n",
+       "      <td>KVQ_0546</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>Our</td>\n",
+       "      <td>0.364</td>\n",
+       "      <td>0.485</td>\n",
+       "      <td>0.753</td>\n",
+       "      <td>2.437</td>\n",
+       "      <td>3.650</td>\n",
+       "      <td>KVQ_0546</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>FAST-VQA</td>\n",
+       "      <td>0.599</td>\n",
+       "      <td>0.673</td>\n",
+       "      <td>0.909</td>\n",
+       "      <td>2.217</td>\n",
+       "      <td>3.319</td>\n",
+       "      <td>SDR_Animal_5ngj</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5</th>\n",
+       "      <td>FasterVQA</td>\n",
+       "      <td>0.489</td>\n",
+       "      <td>0.547</td>\n",
+       "      <td>0.696</td>\n",
+       "      <td>1.343</td>\n",
+       "      <td>3.556</td>\n",
+       "      <td>SDR_Animal_5ngj</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6</th>\n",
+       "      <td>DOVER</td>\n",
+       "      <td>0.920</td>\n",
+       "      <td>1.022</td>\n",
+       "      <td>1.293</td>\n",
+       "      <td>2.783</td>\n",
+       "      <td>3.814</td>\n",
+       "      <td>SDR_Animal_5ngj</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>7</th>\n",
+       "      <td>Our</td>\n",
+       "      <td>0.313</td>\n",
+       "      <td>0.405</td>\n",
+       "      <td>0.697</td>\n",
+       "      <td>2.137</td>\n",
+       "      <td>3.878</td>\n",
+       "      <td>SDR_Animal_5ngj</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ]
+     },
+     "execution_count": 26,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "execution_count": 26
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2026-05-19T12:42:02.353757Z",
+     "start_time": "2026-05-19T12:42:02.348243Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "complexity_df = pd.read_csv(f\"{complexity_csv_path}complexity_test.csv\")\n",
+    "complexity_df"
+   ],
+   "id": "f639ed0f66a50c57",
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "        model              video   Time  Predicted_MOS  Scaled_MOS    MOS\n",
+       "0    FAST-VQA  SDR_Gameplay_wcq2  0.822       0.248790       1.995  4.156\n",
+       "1    FAST-VQA  SDR_Gameplay_s0pc  0.820       0.550190       3.201  4.264\n",
+       "2    FAST-VQA  SDR_Gameplay_z5fz  0.823       0.538800       3.155  4.221\n",
+       "3    FAST-VQA    SDR_Animal_5ngj  0.846       0.579850       3.319  4.308\n",
+       "4    FAST-VQA               0546  0.895       0.754290       4.017  2.389\n",
+       "5   FasterVQA  SDR_Gameplay_wcq2  0.502       0.299890       2.200  4.156\n",
+       "6   FasterVQA  SDR_Gameplay_s0pc  0.524       0.706330       3.825  4.264\n",
+       "7   FasterVQA  SDR_Gameplay_z5fz  0.499       0.400900       2.604  4.221\n",
+       "8   FasterVQA    SDR_Animal_5ngj  0.500       0.638990       3.556  4.308\n",
+       "9   FasterVQA               0546  0.484       0.839070       4.356  2.389\n",
+       "10      DOVER  SDR_Gameplay_wcq2  1.407       0.347305       2.389  4.156\n",
+       "11      DOVER  SDR_Gameplay_s0pc  1.345       0.582339       3.329  4.264\n",
+       "12      DOVER  SDR_Gameplay_z5fz  1.341       0.524504       3.098  4.221\n",
+       "13      DOVER    SDR_Animal_5ngj  1.326       0.703516       3.814  4.308\n",
+       "14      DOVER               0546  1.489       0.805786       4.223  2.389\n",
+       "15        Our  SDR_Gameplay_wcq2  0.594      59.864200       3.395  4.156\n",
+       "16        Our  SDR_Gameplay_s0pc  0.605      61.772900       3.471  4.264\n",
+       "17        Our  SDR_Gameplay_z5fz  0.602      61.077000       3.443  4.221\n",
+       "18        Our    SDR_Animal_5ngj  0.598      71.957800       3.878  4.308\n",
+       "19        Our               0546  0.734      66.257200       3.650  2.389"
+      ],
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>model</th>\n",
+       "      <th>video</th>\n",
+       "      <th>Time</th>\n",
+       "      <th>Predicted_MOS</th>\n",
+       "      <th>Scaled_MOS</th>\n",
+       "      <th>MOS</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>FAST-VQA</td>\n",
+       "      <td>SDR_Gameplay_wcq2</td>\n",
+       "      <td>0.822</td>\n",
+       "      <td>0.248790</td>\n",
+       "      <td>1.995</td>\n",
+       "      <td>4.156</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>FAST-VQA</td>\n",
+       "      <td>SDR_Gameplay_s0pc</td>\n",
+       "      <td>0.820</td>\n",
+       "      <td>0.550190</td>\n",
+       "      <td>3.201</td>\n",
+       "      <td>4.264</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>FAST-VQA</td>\n",
+       "      <td>SDR_Gameplay_z5fz</td>\n",
+       "      <td>0.823</td>\n",
+       "      <td>0.538800</td>\n",
+       "      <td>3.155</td>\n",
+       "      <td>4.221</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>FAST-VQA</td>\n",
+       "      <td>SDR_Animal_5ngj</td>\n",
+       "      <td>0.846</td>\n",
+       "      <td>0.579850</td>\n",
+       "      <td>3.319</td>\n",
+       "      <td>4.308</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>FAST-VQA</td>\n",
+       "      <td>0546</td>\n",
+       "      <td>0.895</td>\n",
+       "      <td>0.754290</td>\n",
+       "      <td>4.017</td>\n",
+       "      <td>2.389</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5</th>\n",
+       "      <td>FasterVQA</td>\n",
+       "      <td>SDR_Gameplay_wcq2</td>\n",
+       "      <td>0.502</td>\n",
+       "      <td>0.299890</td>\n",
+       "      <td>2.200</td>\n",
+       "      <td>4.156</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6</th>\n",
+       "      <td>FasterVQA</td>\n",
+       "      <td>SDR_Gameplay_s0pc</td>\n",
+       "      <td>0.524</td>\n",
+       "      <td>0.706330</td>\n",
+       "      <td>3.825</td>\n",
+       "      <td>4.264</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>7</th>\n",
+       "      <td>FasterVQA</td>\n",
+       "      <td>SDR_Gameplay_z5fz</td>\n",
+       "      <td>0.499</td>\n",
+       "      <td>0.400900</td>\n",
+       "      <td>2.604</td>\n",
+       "      <td>4.221</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>8</th>\n",
+       "      <td>FasterVQA</td>\n",
+       "      <td>SDR_Animal_5ngj</td>\n",
+       "      <td>0.500</td>\n",
+       "      <td>0.638990</td>\n",
+       "      <td>3.556</td>\n",
+       "      <td>4.308</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>9</th>\n",
+       "      <td>FasterVQA</td>\n",
+       "      <td>0546</td>\n",
+       "      <td>0.484</td>\n",
+       "      <td>0.839070</td>\n",
+       "      <td>4.356</td>\n",
+       "      <td>2.389</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>10</th>\n",
+       "      <td>DOVER</td>\n",
+       "      <td>SDR_Gameplay_wcq2</td>\n",
+       "      <td>1.407</td>\n",
+       "      <td>0.347305</td>\n",
+       "      <td>2.389</td>\n",
+       "      <td>4.156</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>11</th>\n",
+       "      <td>DOVER</td>\n",
+       "      <td>SDR_Gameplay_s0pc</td>\n",
+       "      <td>1.345</td>\n",
+       "      <td>0.582339</td>\n",
+       "      <td>3.329</td>\n",
+       "      <td>4.264</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>12</th>\n",
+       "      <td>DOVER</td>\n",
+       "      <td>SDR_Gameplay_z5fz</td>\n",
+       "      <td>1.341</td>\n",
+       "      <td>0.524504</td>\n",
+       "      <td>3.098</td>\n",
+       "      <td>4.221</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>13</th>\n",
+       "      <td>DOVER</td>\n",
+       "      <td>SDR_Animal_5ngj</td>\n",
+       "      <td>1.326</td>\n",
+       "      <td>0.703516</td>\n",
+       "      <td>3.814</td>\n",
+       "      <td>4.308</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>14</th>\n",
+       "      <td>DOVER</td>\n",
+       "      <td>0546</td>\n",
+       "      <td>1.489</td>\n",
+       "      <td>0.805786</td>\n",
+       "      <td>4.223</td>\n",
+       "      <td>2.389</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>15</th>\n",
+       "      <td>Our</td>\n",
+       "      <td>SDR_Gameplay_wcq2</td>\n",
+       "      <td>0.594</td>\n",
+       "      <td>59.864200</td>\n",
+       "      <td>3.395</td>\n",
+       "      <td>4.156</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>16</th>\n",
+       "      <td>Our</td>\n",
+       "      <td>SDR_Gameplay_s0pc</td>\n",
+       "      <td>0.605</td>\n",
+       "      <td>61.772900</td>\n",
+       "      <td>3.471</td>\n",
+       "      <td>4.264</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>17</th>\n",
+       "      <td>Our</td>\n",
+       "      <td>SDR_Gameplay_z5fz</td>\n",
+       "      <td>0.602</td>\n",
+       "      <td>61.077000</td>\n",
+       "      <td>3.443</td>\n",
+       "      <td>4.221</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>18</th>\n",
+       "      <td>Our</td>\n",
+       "      <td>SDR_Animal_5ngj</td>\n",
+       "      <td>0.598</td>\n",
+       "      <td>71.957800</td>\n",
+       "      <td>3.878</td>\n",
+       "      <td>4.308</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>19</th>\n",
+       "      <td>Our</td>\n",
+       "      <td>0546</td>\n",
+       "      <td>0.734</td>\n",
+       "      <td>66.257200</td>\n",
+       "      <td>3.650</td>\n",
+       "      <td>2.389</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ]
+     },
+     "execution_count": 27,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "execution_count": 27
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/demo_test.py

@@ -0,0 +1,264 @@
+import os
+os.environ.setdefault("DECORD_DUPLICATE_WARNING_THRESHOLD", "1.0")
+
+import argparse
+import time
+from pathlib import Path
+
+import pandas as pd
+import torch
+import torch.nn as nn
+from torch.amp import autocast
+from tqdm import tqdm
+from thop import profile
+from thop import clever_format
+
+from train import VQADataset
+from model.qd_model import QD_MODEL
+
+
+def load_checkpoint(ckpt_path, device):
+    ckpt = torch.load(str(ckpt_path), map_location=device, weights_only=True)
+    if isinstance(ckpt, dict) and "model" in ckpt:
+        return {
+            "state_dict": ckpt["model"],
+            "train_mos_mean": ckpt.get("mos_mean"),
+            "train_mos_std": ckpt.get("mos_std"),
+            "train_args": ckpt.get("args", {}),
+            "is_full_checkpoint": True,
+        }
+    if isinstance(ckpt, dict):
+        return {
+            "state_dict": ckpt,
+            "train_mos_mean": None,
+            "train_mos_std": None,
+            "train_args": {},
+            "is_full_checkpoint": False,
+        }
+    raise TypeError(f"Unsupported checkpoint type: {type(ckpt)!r}")
+
+class ForwardWrapper(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def forward(self, rgb, w_art, w_str):
+        yhat, _aux = self.model(rgb, w_art, w_str)
+        return yhat
+
+def count_parameters(model):
+    total = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    return total, trainable
+
+def build_dataset(args, mos_mean, mos_std):
+    rows = [(args.video_id, float(args.dummy_mos))]
+    dataset = VQADataset(
+        rows,
+        args.db_path,
+        clip_len=args.clip_len,
+        size=args.resize,
+        win=args.win,
+        win_step=args.win_step,
+        mos_mean=float(mos_mean),
+        mos_std=float(mos_std),
+    )
+    return dataset
+
+def prepare_single_sample(dataset, device):
+    rgb, w_art, w_str, y, vid = dataset[0]
+
+    rgb = rgb.unsqueeze(0).to(device, non_blocking=True)
+    w_art = w_art.unsqueeze(0).to(device, non_blocking=True)
+    w_str = w_str.unsqueeze(0).to(device, non_blocking=True)
+    y = y.unsqueeze(0).to(device, non_blocking=True).float()
+
+    if isinstance(vid, (list, tuple)):
+        video_id = vid[0]
+    else:
+        video_id = vid
+    return rgb, w_art, w_str, y, video_id
+
+
+@torch.no_grad()
+def predict_once(
+    model,
+    rgb,
+    w_art,
+    w_str,
+    *,
+    device,
+    amp,
+    train_mos_mean,
+    train_mos_std,
+):
+    model.eval()
+    device_type = "cuda" if str(device).startswith("cuda") else "cpu"
+    with autocast(device_type=device_type, enabled=(amp and device_type == "cuda")):
+        yhat, _aux = model(rgb, w_art, w_str)
+
+    pred_score = yhat.detach().float().cpu() * float(train_mos_std) + float(train_mos_mean)
+    return float(pred_score.squeeze().item())
+
+@torch.no_grad()
+def profile_with_thop(model, rgb, w_art, w_str):
+    macs, params = profile(model, inputs=(rgb, w_art, w_str), verbose=False)
+    flops = 2 * macs
+    macs, flops, params = clever_format([macs, flops, params], "%.3f")
+    return macs, flops, params
+
+@torch.no_grad()
+def benchmark_forward(model, rgb, w_art, w_str, *, device, amp, num_runs=10, warmup=3):
+    model.eval()
+    device_type = "cuda" if str(device).startswith("cuda") else "cpu"
+
+    for _ in range(max(0, warmup)):
+        with autocast(device_type=device_type, enabled=(amp and device_type == "cuda")):
+            _ = model(rgb, w_art, w_str)
+    if device_type == "cuda":
+        torch.cuda.synchronize()
+
+    start = time.perf_counter()
+    for _ in range(int(num_runs)):
+        with autocast(device_type=device_type, enabled=(amp and device_type == "cuda")):
+            _ = model(rgb, w_art, w_str)
+    if device_type == "cuda":
+        torch.cuda.synchronize()
+    return (time.perf_counter() - start) / max(1, int(num_runs))
+
+
+def run_end_to_end_once(args, model, train_mos_mean, train_mos_std, device, amp):
+    start = time.perf_counter()
+    dataset = build_dataset(args, train_mos_mean, train_mos_std)
+    rgb, w_art, w_str, _y, _video_id = prepare_single_sample(dataset, device)
+
+    pred_score = predict_once(
+        model,
+        rgb,
+        w_art,
+        w_str,
+        device=device,
+        amp=amp,
+        train_mos_mean=float(train_mos_mean),
+        train_mos_std=float(train_mos_std),
+    )
+    if str(device).startswith("cuda"):
+        torch.cuda.synchronize()
+    elapsed = time.perf_counter() - start
+    return elapsed, pred_score, rgb, w_art, w_str
+
+
+def parse_args():
+    ap = argparse.ArgumentParser(description="Demo-style single-video test for QD_MODEL")
+    # for complexity time test:
+    ap.add_argument("--ckpt_path", type=str, default="/home/xinyi/Project/FD-VQA/src/checkpoints/lsvq/qd_model.best.pt")
+    ap.add_argument("--db_path", type=str, default="/home/xinyi/Project/FD-VQA/test_videos/")
+    ap.add_argument("--video_id", type=str, default="SDR_Animal_5ngj")
+    # for resolution compelxity test:
+    # ap.add_argument("--ckpt_path", type=str, default="/home/xinyi/Project/FD-VQA/src/checkpoints/kvq/qd_model.best.pt")
+    # ap.add_argument("--db_path", type=str, default="/home/xinyi/Project/FD-VQA/test_videos/complexity_test/complexity_resolution/")
+    # ap.add_argument("--video_id", type=str, default="SDR_Animal_5ngj_540p")
+
+    ap.add_argument("--clip_len", type=int, default=16)
+    ap.add_argument("--resize", type=int, default=224)
+    ap.add_argument("--win", type=int, default=6)
+    ap.add_argument("--win_step", type=int, default=1)
+
+    ap.add_argument("--device", type=str, default="cuda")
+    ap.add_argument("--no_amp", action="store_true")
+
+    ap.add_argument("--dummy_mos", type=float, default=3.0, help="Only used to compute VQADataset, does not affect prediction")
+    ap.add_argument("--num_runs", type=int, default=10, help="Average N runs")
+    ap.add_argument("--warmup_runs", type=int, default=3)
+    ap.add_argument("--skip_profile", action="store_true")
+    return ap.parse_args()
+
+
+def main():
+    args = parse_args()
+    device = torch.device(args.device)
+    amp = not bool(args.no_amp)
+
+    print(f"Running on {'GPU' if device.type == 'cuda' else 'CPU'}")
+
+    display_path = str(Path(args.db_path) / args.video_id)
+    info = pd.DataFrame([
+        {
+            "vid": args.video_id,
+            "test_video_path": display_path,
+        }
+    ])
+    print(info)
+
+    dataset_preview = build_dataset(args, mos_mean=args.dummy_mos, mos_std=1.0)
+    print(f"Dataset loaded. Total videos: {len(dataset_preview)}, Total batches: 1")
+    print(f"Loading model from: {args.ckpt_path}")
+
+    ckpt_info = load_checkpoint(Path(args.ckpt_path), device)
+    train_mos_mean = ckpt_info["train_mos_mean"]
+    train_mos_std = ckpt_info["train_mos_std"]
+    if train_mos_mean is None or train_mos_std is None:
+        raise ValueError("Checkpoint does not contain mos_mean / mos_std. Please use a full checkpoint.")
+    if float(train_mos_std) <= 1e-8:
+        raise ValueError("train_mos_std must be > 0")
+
+    model = QD_MODEL(clip_model="openai/clip-vit-base-patch16").to(device)
+    model.load_state_dict(ckpt_info["state_dict"], strict=True)
+    model.eval()
+
+    run_times = []
+    pred_score = None
+    rgb = w_art = w_str = None
+
+    for i in range(args.num_runs):
+        for _ in tqdm(range(1), desc="Processing Videos"):
+            elapsed, pred_score, rgb, w_art, w_str = run_end_to_end_once(
+                args, model, train_mos_mean, train_mos_std, device, amp
+            )
+        run_times.append(elapsed)
+        print(f"Run {i + 1} - Time taken: {elapsed:.4f} seconds")
+
+    avg_total_time = sum(run_times) / max(1, len(run_times))
+    avg_forward_time = benchmark_forward(
+        model,
+        rgb,
+        w_art,
+        w_str,
+        device=device,
+        amp=amp,
+        num_runs=args.num_runs,
+        warmup=args.warmup_runs,
+    )
+
+    total_params, trainable_params = count_parameters(model)
+    macs = flops = params = None
+    if not args.skip_profile:
+        try:
+            macs, flops, params = profile_with_thop(model, rgb, w_art, w_str)
+        except Exception as e:
+            print(f"[WARN] THOP profiling failed: {e}")
+
+    print(f"Average running time over {args.num_runs} runs: {avg_total_time:.4f} seconds")
+    print(f"Predicted Quality Score: {pred_score:.4f}")
+    print("\n========== PROFILE SUMMARY ==========")
+    print(f"video_id              : {args.video_id}")
+    print(f"rgb shape             : {tuple(rgb.shape)}")
+    print(f"w_art shape           : {tuple(w_art.shape)}")
+    print(f"w_str shape           : {tuple(w_str.shape)}")
+    print(f"train mos mean/std    : {float(train_mos_mean):.6f} / {float(train_mos_std):.6f}")
+    print(f"predicted score       : {pred_score:.6f}")
+    print(f"params total          : {total_params:,} ({total_params / 1e6:.3f} M)")
+    print(f"params trainable      : {trainable_params:,} ({trainable_params / 1e6:.3f} M)")
+    if params is not None:
+        print(f"Params (THOP)         : {params} M")
+    if macs is not None:
+        print(f"MACs (THOP)           : {macs} G")
+    if flops is not None:
+        print(f"FLOPs (~2*MACs)       : {flops} G")
+    print(f"avg forward time      : {avg_forward_time:.6f} s  (runs={args.num_runs})")
+    print(f"avg end-to-end time   : {avg_total_time:.6f} s  (sample prep + forward)")
+    print("=====================================")
+
+
+if __name__ == "__main__":
+    main()

src/model/__init__.py

@@ -0,0 +1 @@
+# src/model/__init__.py

src/model/clip_dense_encoder.py

@@ -0,0 +1,61 @@
+import torch
+import torch.nn as nn
+from transformers import CLIPModel
+from transformers.utils import logging
+logging.set_verbosity_error()
+
+
+class CLIPDenseEncoder(nn.Module):
+    def __init__(self, model_name="openai/clip-vit-base-patch16"):
+        super().__init__()
+        self.clip = CLIPModel.from_pretrained(str(model_name), use_safetensors=True)
+
+        self.vision = self.clip.vision_model
+        vcfg = self.clip.config.vision_config
+        self.hidden_size = int(vcfg.hidden_size)
+        self.patch_size = int(vcfg.patch_size)
+        self.image_size = int(vcfg.image_size)
+
+        # OpenAI CLIP normalization constants
+        self.register_buffer("_mean", torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1))
+        self.register_buffer("_std", torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1))
+
+    def forward(self, x01):
+        x = (x01 - self._mean.to(dtype=x01.dtype)) / self._std.to(dtype=x01.dtype)
+        out = self.vision(pixel_values=x)
+        tokens = out.last_hidden_state[:, 1:, :]  # [B, 1+N, C] drop CLS -> [B, N, C]
+        b, n, c = tokens.shape
+        side = int(n**0.5)
+        if side * side != n:
+            raise RuntimeError(f"CLIP patch tokens N={n} not a square; cannot reshape to 2D grid.")
+
+        fmap = tokens.transpose(1, 2).contiguous().view(b, c, side, side)
+        return fmap
+
+    def freeze_all(self):
+        for p in self.clip.parameters():
+            p.requires_grad = False
+
+    def unfreeze_last_blocks(self, n_blocks=2, also_unfreeze_ln=True):
+        self.freeze_all()
+        layers = self.vision.encoder.layers
+        n_total = len(layers)
+
+        # all unfreeze
+        if n_blocks < 0 or n_blocks >= n_total:
+            for p in self.vision.parameters():
+                p.requires_grad = True
+        # unfreeze n blocks
+        else:
+            k = max(0, min(int(n_blocks), n_total))
+            for i in range(n_total - k, n_total):
+                for p in layers[i].parameters():
+                    p.requires_grad = True
+
+        if also_unfreeze_ln:
+            if hasattr(self.vision, "pre_layrnorm"):
+                for p in self.vision.pre_layrnorm.parameters():
+                    p.requires_grad = True
+            if hasattr(self.vision, "post_layernorm"):
+                for p in self.vision.post_layernorm.parameters():
+                    p.requires_grad = True

src/model/qd_model.py

@@ -0,0 +1,178 @@
+# ----------------------------
+# Model: CLIP encoder + w_art/w_str pooling -> weighted sequences -> seq model -> head fusion
+# ----------------------------
+import torch
+import torch.nn as nn
+
+from .clip_dense_encoder import CLIPDenseEncoder
+
+
+def weighted_pool_2d(fmap, wmap, eps=1e-6):
+    """
+    fmap: [B, C, H, W]
+    wmap: [B, 1, H, W]
+    -> [B, C]
+    """
+    w = wmap.clamp(0.0, 1.0)
+    w = w / (w.sum(dim=(2, 3), keepdim=True) + eps)
+    return (fmap * w).sum(dim=(2, 3))
+
+def default_gate_stats(w_art, w_str, fmap_bt, eps=1e-6):
+    """
+    w_art:   [B, 1, T, H, W]
+    w_str:   [B, 1, T, H, W]
+    fmap_bt: [B, T, C, H, W]
+    -> [B, 3]
+    """
+    mu_art = w_art.mean(dim=(1, 2, 3, 4))
+    mu_str = w_str.mean(dim=(1, 2, 3, 4))
+    mu_raw = fmap_bt.abs().mean(dim=(1, 2, 3, 4))
+
+    stats = torch.stack([mu_art, mu_str, mu_raw], dim=1)
+    return stats
+
+def two_gate_stats(w_art, w_str):
+    # [B, 2]
+    mu_art = w_art.mean(dim=(1, 2, 3, 4))
+    mu_str = w_str.mean(dim=(1, 2, 3, 4))
+    return torch.stack([mu_art, mu_str], dim=1)
+
+class QD_MODEL(nn.Module):
+    def __init__(
+        self,
+        *,
+        clip_model="openai/clip-vit-base-patch16",
+        head_hidden=384,
+        gate_hidden=32,
+        head_dropout=0.2,
+        gate_dropout=0.1,
+        ablation_mode="full",  # "full" | "art" | "str" | "raw"
+    ):
+        super().__init__()
+        self.ablation_mode = ablation_mode
+
+        self.encoder = CLIPDenseEncoder(model_name=str(clip_model))
+        c = int(self.encoder.hidden_size)
+
+        # 3 heads
+        self.head_art = nn.Sequential(
+            nn.Linear(c, head_hidden),
+            nn.ReLU(),
+            nn.Dropout(head_dropout),
+            nn.Linear(head_hidden, head_hidden // 2),
+            nn.ReLU(),
+            nn.Dropout(head_dropout),
+            nn.Linear(head_hidden // 2, 1),
+        )
+        self.head_str = nn.Sequential(
+            nn.Linear(c, head_hidden),
+            nn.ReLU(),
+            nn.Dropout(head_dropout),
+            nn.Linear(head_hidden, head_hidden // 2),
+            nn.ReLU(),
+            nn.Dropout(head_dropout),
+            nn.Linear(head_hidden // 2, 1),
+        )
+        self.head_fmap = nn.Sequential(
+            nn.Linear(c, head_hidden),
+            nn.ReLU(),
+            nn.Dropout(head_dropout),
+            nn.Linear(head_hidden, head_hidden // 2),
+            nn.ReLU(),
+            nn.Dropout(head_dropout),
+            nn.Linear(head_hidden // 2, 1),
+        )
+        # full model: 3-way gate
+        self.gate = nn.Sequential(
+            nn.Linear(3, gate_hidden),
+            nn.ReLU(),
+            nn.Dropout(gate_dropout),
+            nn.Linear(gate_hidden, 3),
+        )
+        # ablation: two gates only
+        self.gate_two = nn.Sequential(
+            nn.Linear(2, gate_hidden),
+            nn.ReLU(),
+            nn.Dropout(gate_dropout),
+            nn.Linear(gate_hidden, 2),
+        )
+
+    def freeze_clip_all(self):
+        self.encoder.freeze_all()
+
+    def unfreeze_clip_last_blocks(self, n_blocks=2, also_unfreeze_ln=True):
+        self.encoder.unfreeze_last_blocks(
+            n_blocks=n_blocks,
+            also_unfreeze_ln=also_unfreeze_ln,
+        )
+
+    def temporal_pool(self, x):
+        """
+        x: [B, T, C]
+        return: [B, C]
+        """
+        return x.mean(dim=1)
+
+    def forward(self, rgb, w_art, w_str, gate_stats=None):
+        """
+        rgb: [B,3,T,H,W]   (rgb in [0,1])
+        w_art/w_str: [B,1,T,H,W]   (0..1)
+        gate_stats: [B,3]
+        """
+        B, _C, T, H, W = rgb.shape
+
+        x2d = rgb.permute(0, 2, 1, 3, 4).contiguous().view(B * T, 3, H, W)
+        fmap2d = self.encoder(x2d)
+        _, C, Hp, Wp = fmap2d.shape  # (B*T, 768, 14, 14)
+
+        w_art_bt = w_art.transpose(1, 2)  # (B, T, 1, 14, 14)
+        w_str_bt = w_str.transpose(1, 2)
+        fmap_bt = fmap2d.view(B, T, C, Hp, Wp)  # (B, T, 768, 14, 14)
+
+        z_art = torch.stack([weighted_pool_2d(fmap_bt[:, i], w_art_bt[:, i]) for i in range(T)], dim=1)  # (B, T, 768)
+        z_str = torch.stack([weighted_pool_2d(fmap_bt[:, i], w_str_bt[:, i]) for i in range(T)], dim=1)
+        z_raw = fmap_bt.mean(dim=(-2, -1))
+
+        h_art = self.temporal_pool(z_art)  # [B, C]
+        h_str = self.temporal_pool(z_str)
+        h_fmap = self.temporal_pool(z_raw)
+
+        q_art = self.head_art(h_art)
+        q_str = self.head_str(h_str)
+        q_fmap = self.head_fmap(h_fmap)
+
+        mode = self.ablation_mode.lower()
+        if mode == "art":
+            y_hat = q_art
+            weights = torch.tensor([1.0, 0.0, 0.0], device=q_art.device).unsqueeze(0).repeat(B, 1)
+        elif mode == "str":
+            y_hat = q_str
+            weights = torch.tensor([0.0, 1.0, 0.0], device=q_art.device).unsqueeze(0).repeat(B, 1)
+        elif mode == "raw":
+            y_hat = q_fmap
+            weights = torch.tensor([0.0, 0.0, 1.0], device=q_art.device).unsqueeze(0).repeat(B, 1)
+        elif mode == "art+str":
+            if gate_stats is None:
+                gate_stats = two_gate_stats(w_art, w_str)  # [B, 2]
+            g_ar = self.gate_two(gate_stats)  # [B, 2]
+            a, b_ = torch.softmax(g_ar, dim=1).split(1, dim=1)
+            y_hat = a * q_art + b_ * q_str
+            zero = torch.zeros_like(a)
+            weights = torch.cat([a, b_, zero], dim=1)
+        elif mode == "full":
+            if gate_stats is None:
+                gate_stats = default_gate_stats(w_art, w_str, fmap_bt)
+            g = self.gate(gate_stats)
+            a, b_, c_ = torch.softmax(g, dim=1).split(1, dim=1)
+            y_hat = a * q_art + b_ * q_str + c_ * q_fmap
+            weights = torch.cat([a, b_, c_], dim=1)
+        else:
+            raise ValueError(f"Unknown ablation_mode: {self.ablation_mode}")
+
+        aux = (
+            q_art.squeeze(1),
+            q_str.squeeze(1),
+            q_fmap.squeeze(1),
+            weights,
+        )
+        return y_hat.squeeze(1), aux

src/module/__init__.py

@@ -0,0 +1 @@
+# src/module/__init__.py

src/module/compute_weight_map.py

@@ -0,0 +1,144 @@
+import numpy as np
+import csv
+from pathlib import Path
+import os
+os.environ["DECORD_DUPLICATE_WARNING_THRESHOLD"] = "1.0"
+from decord import VideoReader, cpu
+
+from module.frequency_dct import compute_twostream_dct
+from module.read_frame_decord import (
+    sample_frames_uniform,
+    collect_needed,
+    cache_needed_frames,
+    read_window_from_cache,
+)
+
+def process_video(
+    vr,
+    # resize
+    size=224,
+    # anchors + window
+    num_anchors=16,
+    win=6,
+    win_step=1,
+    # DCT / weights
+    block=16,
+):
+    """
+    Returns (frame_all, w_art_all, w_str_all, anchors_kept).
+    frame_all: anchor frames (RGB uint8, HxWx3)
+    w_art_all/w_str_all: maps (float32, HxW)
+    anchors_kept: frame indices used per anchor - list[list[int]]
+    """
+    total_frames = len(vr)
+    if total_frames <= 1:
+        raise RuntimeError(f"Video too short / invalid frame count: {total_frames}")
+
+    anchor_idxs = sample_frames_uniform(total_frames, num_anchors, win=win, win_step=win_step)
+    needed = collect_needed(anchor_idxs, total_frames, win, win_step)
+    cache = cache_needed_frames(vr, needed, size)
+
+    frame_all, w_art_all, w_str_all = [], [], []
+    anchors_kept = []
+
+    for anchor in anchor_idxs:
+        out = read_window_from_cache(cache, anchor, total_frames, win, win_step)
+        if out is None:
+            continue
+
+        anchor_frame, gray_seq, idxs = out
+
+        w_art, w_str, _dbg = compute_twostream_dct(
+            gray_seq,
+            block=block,
+        )
+
+        frame_all.append(anchor_frame)
+        w_art_all.append(w_art.astype(np.float32, copy=False))
+        w_str_all.append(w_str.astype(np.float32, copy=False))
+        anchors_kept.append([int(x) for x in idxs])
+
+    return frame_all, w_art_all, w_str_all, anchors_kept
+
+
+def compute_weight_map(frame_all, w_art_all, w_str_all):
+    if len(frame_all) == 0:
+        raise ValueError("No frames produced.")
+    if not (len(frame_all) == len(w_art_all) == len(w_str_all)):
+        raise ValueError(
+            f"Length mismatch: frames={len(frame_all)}, w_art_all={len(w_art_all)}, w_str_all={len(w_str_all)}"
+        )
+    frames_np = np.stack(frame_all, axis=0)  # (N,H,W,3) uint8
+    w_art_np = np.stack(w_art_all, axis=0)         # (N,H,W) float32
+    w_str_np = np.stack(w_str_all, axis=0)         # (N,H,W) float32
+    return frames_np, w_art_np, w_str_np
+
+def read_vid_mos_csv(csv_path):
+    rows = []
+    with open(csv_path, "r", encoding="utf-8") as f:
+        reader = csv.DictReader(f)
+        if not reader.fieldnames:
+            raise RuntimeError("CSV has no header")
+        for r in reader:
+            vid = str(r["vid"]).strip()
+            mos = float(r["mos"])
+            rows.append((vid, mos))
+    return rows
+
+def rng_rows(rows, seed=0):
+    rng = np.random.default_rng(int(seed))
+    idx = np.arange(len(rows))
+    rng.shuffle(idx)
+    train = [rows[i] for i in idx[:]]
+    return train
+
+if __name__ == "__main__":
+    csv_path = "/home/xinyi/Project/FD-VQA/metadata/TEST_metadata.csv"
+    db_path = "/home/xinyi/Project/FD-VQA/test_videos/"
+    # video_path = "/home/xinyi/Project/FD-VQA/test_videos/NesAirFortressIn4108.37ByTool23.mp4"
+
+    rows = read_vid_mos_csv(str(csv_path))
+    train = rng_rows(rows)
+    for i in range(len(train)):
+        vid, mos = train[i]
+        print(vid, mos)
+        # get video path
+        base_path = Path(db_path) / vid
+        video_path = None
+        for ext in ("mp4", "avi", "mkv"):
+            p = Path(str(base_path) + f".{ext}")
+            if p.exists():
+                video_path = str(p)
+                break
+        if video_path is None:
+            raise FileNotFoundError(f"Cannot find {vid} video")
+
+        try:
+            # read video
+            vr = VideoReader(video_path, ctx=cpu(0))
+            frame_all, w_art_all, w_str_all, anchors_kept = process_video(
+                vr,
+                size=224,
+                num_anchors=16,
+                win=6,
+                win_step=1,
+                block=16,
+            )
+            frames_np, w_art_np, w_str_np = compute_weight_map(frame_all, w_art_all, w_str_all)
+            print("frames_np:", frames_np.shape, frames_np.dtype)
+            print("w_art_np:", w_art_np.shape, w_art_np.dtype)
+            print("w_str_np:", w_str_np.shape, w_str_np.dtype)
+            print("anchors_kept:", len(anchors_kept), "example:", anchors_kept)
+        except Exception as e:
+            print("\n[DATA ERROR]")
+            print("idx:", i)
+            print("vid:", vid)
+            print("path:", video_path)
+            raise
+        finally:
+            # release decord video reader
+            try:
+                if vr is not None:
+                    del vr
+            except Exception:
+                pass

src/module/frequency_dct.py

@@ -0,0 +1,351 @@
+import os
+os.environ["DECORD_DUPLICATE_WARNING_THRESHOLD"] = "1.0"
+from decord import VideoReader, cpu
+import cv2
+import argparse
+import numpy as np
+from tqdm import tqdm
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+from module.read_frame_decord import sample_frames_uniform, collect_needed, cache_needed_frames, read_window_from_cache
+
+# ----------------------------
+# utils
+# ----------------------------
+def norm01(x, eps=1e-6):
+    x = x.astype(np.float32)
+    mn, mx = float(x.min()), float(x.max())
+    return (x - mn) / (mx - mn + eps)
+
+# def upsample_grid(grid, out_hw, interp=cv2.INTER_LINEAR):
+def upsample_grid(grid, out_hw, interp=cv2.INTER_NEAREST):
+    H, W = out_hw
+    return cv2.resize(grid.astype(np.float32), (W, H), interpolation=interp)
+
+def gaussian_blur(x, sigma=1.0):
+    if sigma <= 0:
+        return x
+    ksize = int(6 * sigma + 1)
+    if ksize % 2 == 0:
+        ksize += 1
+    if ksize < 3:
+        ksize = 3
+    return cv2.GaussianBlur(x, (ksize, ksize), sigma, borderType=cv2.BORDER_REFLECT_101)
+
+# sobel operator
+def gradient_mag(gray):
+    gx = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=3)
+    gy = cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=3)
+    dst = np.sqrt(gx * gx + gy * gy)
+    # plt.imshow(dst, cmap='gray')
+    # plt.savefig('/home/xinyi/Project/FD-VQA/test_videos/freq_test_dct_only/Sobel_operator_result.jpg', dpi=300)
+    return dst
+
+# Sobel magnitude -> normalize -> threshold -> dilate edge region
+def edge_sobel(gray, thr=0.20, dilate_px=2):
+    g = gradient_mag(gray).astype(np.float32)
+    g = norm01(g)  # normalize
+    edge = (g >= thr).astype(np.uint8)
+    if dilate_px > 0:
+        k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * dilate_px + 1, 2 * dilate_px + 1))
+        edge = cv2.dilate(edge, k, iterations=1)
+    # plt.imshow(edge, cmap='gray')
+    # plt.savefig('/home/xinyi/Project/FD-VQA/test_videos/freq_test_dct_only/edge_result.jpg', dpi=300)
+    return edge
+
+# per-block fraction of near edge pixels
+def block_edge(edge_mask, block=16):
+    h, w = edge_mask.shape
+    gh, gw = h // block, w // block
+    frac = np.zeros((gh, gw), dtype=np.float32)
+    for by in range(gh):
+        for bx in range(gw):
+            y0, x0 = by * block, bx * block
+            frac[by, bx] = float(edge_mask[y0:y0 + block, x0:x0 + block].mean())
+    # plt.imshow(frac, cmap='gray')
+    # plt.savefig('/home/xinyi/Project/FD-VQA/test_videos/freq_test_dct_only/edge_block.jpg', dpi=300)
+    return frac
+
+# per-block mean
+def block_mean(img, block=16):
+    h, w = img.shape
+    gh, gw = h // block, w // block
+    x = img[:gh*block, :gw*block]
+    x = x.reshape(gh, block, gw, block)      # (block_row, inblock_row, block_col, inblock_col)
+    return x.mean(axis=(1, 3)).astype(np.float32)
+
+# discontinuities across block boundaries: |I[:, x] - I[:, x-1]| & |I[y, :] - I[y-1, :]|
+def blockiness_boundary_map(gray, block_size=16, blur_sigma=1.0):
+    h, w = gray.shape
+    v = np.zeros((h, w), dtype=np.float32)
+    hmap = np.zeros((h, w), dtype=np.float32)
+
+    for x in range(block_size, w, block_size):
+        v[:, x] = np.abs(gray[:, x] - gray[:, x - 1])
+    for y in range(block_size, h, block_size):
+        hmap[y, :] = np.abs(gray[y, :] - gray[y - 1, :])
+    b = v + hmap
+    if blur_sigma > 0:
+        b = gaussian_blur(b, blur_sigma)
+    return b
+
+# ----------------------------
+# DCT 16x16 energies
+# ----------------------------
+def dct_energy_ratios(gray, block=16, low_k=4, mid_k=8, eps=1e-6):
+    """
+    Return per-block energy ratios for frequency bands.
+    low:  [0:low_k, 0:low_k]
+    mid:  [0:mid_k, 0:mid_k] - low
+    high: remaining
+    """
+    H, W = gray.shape
+    gh, gw = H // block, W // block
+    r_low = np.zeros((gh, gw), dtype=np.float32)
+    r_mid = np.zeros((gh, gw), dtype=np.float32)
+    r_high = np.zeros((gh, gw), dtype=np.float32)
+
+    for by in range(gh):
+        for bx in range(gw):
+            y0, x0 = by * block, bx * block
+            patch = gray[y0:y0 + block, x0:x0 + block].astype(np.float32)
+
+            C = cv2.dct(patch)
+            E = C * C
+            E_low = E[:low_k, :low_k].sum()  # 4 x 4
+            E_mid = E[:mid_k, :mid_k].sum() - E_low  # 8×8
+            E_total = E.sum()
+            E_high = max(E_total - (E_low + E_mid), 0.0)
+
+            denom = E_total + eps
+            r_low[by, bx] = E_low / denom
+            r_mid[by, bx] = E_mid / denom
+            r_high[by, bx] = E_high / denom
+    return r_low, r_mid, r_high
+
+# Temporal map via FFT
+def temporal_fft_map(gray_seq, *, block=16, hf_start_bin=2, eps=1e-6):
+    # each frame -> (K, gh, gw) block mean grid
+    grids = [block_mean(g.astype(np.float32), block=block) for g in gray_seq]
+    s = np.stack(grids, axis=0)
+
+    # rFFT along time
+    X = np.fft.rfft(s, axis=0)  # (F, gh, gw)
+    E = (X.real * X.real + X.imag * X.imag).astype(np.float32)  # power spectrum
+
+    # F = K // 2 + 1
+    dc = E[0]
+    if E.shape[0] <= 1:
+        z = np.zeros_like(dc, dtype=np.float32)
+        return z, z
+    P = E[1:]  # drop DC -> (F-1, gh, gw)
+    non_dc = P.sum(axis=0)
+
+    # changes relative to DC
+    motion = non_dc / (dc + eps)
+
+    # flicker: change is in high temporal freqs
+    start = max(int(hf_start_bin) - 1, 0)  # index in P
+    if start >= P.shape[0]:
+        flicker = np.zeros_like(non_dc, dtype=np.float32)
+    else:
+        hi = P[start:].sum(axis=0)
+        flicker = hi / (non_dc + eps)
+    return motion.astype(np.float32), flicker.astype(np.float32)
+
+def fuse_temporal_maps(motion_grid, flicker_grid, *, beta=0.5):
+    m = norm01(motion_grid)
+    f = np.clip(flicker_grid, 0.0, 1.0)
+    # boosts where flicker is high
+    w = m * ((1.0 - beta) + beta * f)
+    return norm01(w)
+
+# ----------------------------
+# DCT -> two stream weights
+# ----------------------------
+def compute_twostream_dct(
+    gray_seq,
+    *,
+    block=16,
+):
+    K = len(gray_seq)
+    gray_anchor = gray_seq[0]
+    H, W = gray_anchor.shape
+
+    r_low_stack, r_mid_stack, r_high_stack = [], [], []
+    for g in gray_seq:
+        r_low, r_mid, r_high = dct_energy_ratios(g, block=block)
+        r_low_stack.append(r_low)  # (gh, gw)
+        r_mid_stack.append(r_mid)
+        r_high_stack.append(r_high)
+    r_low_stack = np.stack(r_low_stack, axis=0)  # (K, gh, gw)
+    r_mid_stack = np.stack(r_mid_stack, axis=0)
+    r_high_stack = np.stack(r_high_stack, axis=0)
+
+    # frequency band (anchor frame)
+    anchor_low_grid = r_low_stack[0]  # (gh, gw)
+    anchor_mid_grid = r_mid_stack[0]
+    anchor_high_grid = r_high_stack[0]
+
+    # Ringing map (anchor)
+    edge_mask = edge_sobel(gray_anchor)     # around edges
+    edge_frac = block_edge(edge_mask, block=block)
+    mh_band = r_mid_stack[0] + r_high_stack[0]  # mid/high frequency energy
+    ring_score = np.maximum(mh_band, 0.0)       # score = mid_high - 0 * r_low # alpha = 0
+    edge_min_frac = 0.05
+    ringing_grid = np.where(edge_frac >= edge_min_frac, edge_frac * ring_score, 0.0).astype(np.float32)
+    s = np.percentile(ringing_grid, 99) + 1e-6
+    ringing_grid01 = np.clip(ringing_grid / s, 0.0, 1.0)
+
+    # Blur map (anchor): like low-pass filtering
+    hf = 0.5 * r_mid_stack[0] + 1.0 * r_high_stack[0]
+    blur_raw = np.clip(1.0 - hf, 0.0, 1.0)
+    sobel_g = gradient_mag(gray_anchor).astype(np.float32)
+    sobel_g_grid = block_mean(sobel_g, block=block)
+    sobel_g_grid = norm01(sobel_g_grid)  # soft structure weight
+    blur_grid = np.clip(blur_raw * sobel_g_grid, 0.0, 1.0)
+
+    # Blockiness map (anchor): boundary discontinuities
+    boundary_pix = blockiness_boundary_map(gray_anchor, block_size=block)
+    blockiness_grid = norm01(block_mean(boundary_pix, block=block))
+
+    # Temporal (window): FFT along time
+    if K >= 4:
+        motion_grid, flick_grid = temporal_fft_map(gray_seq, block=block, hf_start_bin=2)
+        temporal_grid = fuse_temporal_maps(motion_grid, flick_grid, beta=0.5)
+    elif K == 2:
+        E_stack = norm01(r_mid_stack + r_high_stack)
+        temporal_grid = norm01(np.abs(E_stack[1] - E_stack[0]))
+
+    # -------------Combine---------------
+    w_art = norm01(1.0 * ringing_grid01 + 1.0 * blur_grid + 1.0 * blockiness_grid + 1.0 * temporal_grid)
+    w_str = 1.0 - w_art
+
+    debug = {
+        # frequency band (anchor)
+        "dct_low_grid": anchor_low_grid,
+        "dct_mid_grid": anchor_mid_grid,
+        "dct_high_grid": anchor_high_grid,
+        # ringing (anchor)
+        "ringing_grid": ringing_grid01,
+        "edge_px": edge_mask,
+        # blur (anchor)
+        "blur_grid": blur_grid,
+        # blockiness (anchor)
+        "blockiness_grid": blockiness_grid,
+        # temporal (window)
+        "temporal_grid": temporal_grid,
+    }
+    return w_art, w_str, debug
+
+# ----------------------------
+# visualization panel
+# ----------------------------
+def save_panel(out_png, frame_rgb, w_art, w_str, debug):
+    fig = plt.figure(figsize=(16, 9), dpi=160)
+    def add(ax_i, title, img, cmap=None, vmin=0, vmax=1):
+        ax = fig.add_subplot(3, 4, ax_i)
+        ax.set_title(title)
+        if cmap is None:
+            ax.imshow(img)
+        else:
+            ax.imshow(img, cmap=cmap, vmin=vmin, vmax=vmax)
+        ax.axis("off")
+
+    add(1, "Frame_t (anchor)", frame_rgb)
+    add(2, "DCT LOW (grid)",  norm01(debug["dct_low_grid"]),  cmap="viridis")
+    add(3, "DCT MID (grid)",  norm01(debug["dct_mid_grid"]),  cmap="viridis")
+    add(4, "DCT HIGH (grid)", norm01(debug["dct_high_grid"]), cmap="viridis")
+    # ringing
+    add(5, "EDGE (mask)", debug["edge_px"], cmap="viridis")
+    add(6, "RINGING (mid/high, grid)", debug["ringing_grid"], cmap="viridis")
+    # blur
+    add(7, "BLUR (lowpass, grid)", debug["blur_grid"], cmap="viridis")
+    # blockiness
+    add(8, "BLOCKINESS (boundary, grid)", debug["blockiness_grid"], cmap="viridis")
+    # temporal
+    add(9, "TEMPORAL (grid)", debug["temporal_grid"], cmap="viridis")
+    # all map
+    add(10, "W_art", w_art, cmap="viridis")
+    add(11, "W_str", w_str, cmap="viridis")
+    os.makedirs(os.path.dirname(out_png), exist_ok=True)
+    fig.tight_layout()
+    fig.savefig(out_png)
+    plt.close(fig)
+
+
+# ----------------------------
+# Main
+# ----------------------------
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--video", default="/home/xinyi/Project/FD-VQA/test_videos/SDR_Animal_5ngj.mp4")
+    parser.add_argument("--out_dir", default="/home/xinyi/Project/FD-VQA/test_videos/freq_test_dct_only")
+    parser.add_argument("--size", type=int, default=224)
+
+    # fixed-T anchors over whole video (duration-uniform)
+    parser.add_argument("--num_anchors", type=int, default=16)
+    parser.add_argument("--win", type=int, default=6)
+    parser.add_argument("--win_step", type=int, default=1)
+    parser.add_argument("--block", type=int, default=16)
+
+    parser.add_argument("--no_panel", action="store_true")
+    args = parser.parse_args()
+    os.makedirs(args.out_dir, exist_ok=True)
+
+    vr = VideoReader(args.video, ctx=cpu(0))
+    total_frames = len(vr)
+    if total_frames <= 1:
+        raise RuntimeError(f"Video too short / frame count unavailable: total_frames={total_frames}")
+    print("total_frames:", total_frames)
+
+    size = args.size
+    win = args.win
+    win_step = args.win_step
+    num_anchors = args.num_anchors
+
+    anchor_idxs = sample_frames_uniform(total_frames, num_anchors, win=win, win_step=win_step)
+    needed = collect_needed(anchor_idxs, total_frames, win, win_step)
+    print("anchor_idxs:", anchor_idxs)
+    cache = cache_needed_frames(vr, needed, size)
+    print("cached:", len(cache), "needed:", len(needed))
+
+    frame_all, w_art_all, w_str_all = [], [], []
+    anchors_kept = []
+    image_idx = 0
+
+    for anchor in tqdm(anchor_idxs, desc="Processing anchors (DCT)"):
+        out = read_window_from_cache(cache, anchor, total_frames, win, win_step)
+        if out is None:
+            continue
+        anchor_frame, gray_seq, idxs = out
+
+        w_art, w_str, dbg = compute_twostream_dct(
+            gray_seq,
+            block=args.block,
+        )
+
+        frame_all.append(anchor_frame)
+        w_art_all.append(w_art)
+        w_str_all.append(w_str)
+        anchors_kept.append(idxs)
+
+        image_idx += 1
+        if not args.no_panel:
+            save_panel(
+                os.path.join(args.out_dir, f"anchor_{anchor:03d}_{image_idx:02d}.png"),
+                anchor_frame,
+                w_art,
+                w_str,
+                dbg,
+            )
+
+    print(f"Done. Outputs saved to: {args.out_dir}")
+    print(anchors_kept)
+    print(f"total_frames={total_frames}, num_anchors_target={num_anchors}, anchors_produced={len(w_str_all)}")
+
+if __name__ == "__main__":
+    main()

src/module/read_frame_decord.py

@@ -0,0 +1,138 @@
+import cv2
+import numpy as np
+import os
+os.environ["DECORD_DUPLICATE_WARNING_THRESHOLD"] = "1.0"
+from decord import VideoReader, cpu
+
+def sample_frames_uniform(total_frames, num_anchors, win=8, win_step=1):
+    if total_frames <= 0:
+        return [0] * num_anchors
+
+    min_win = (win - 1) * win_step + 1
+    if total_frames < num_anchors + min_win:
+        if total_frames >= num_anchors:
+            anchor_idxs = np.linspace(0, total_frames - 1, num_anchors).round().astype(int)
+            return anchor_idxs.tolist()
+        else:
+            anchor_idxs = list(range(total_frames))
+            anchor_idxs.extend([total_frames - 1] * (num_anchors - total_frames))
+            return anchor_idxs
+
+    max_start = total_frames - min_win
+    anchor_idxs = np.linspace(0, max_start, num_anchors).round().astype(int)
+    return anchor_idxs.tolist()
+
+# window idx for anchor
+def window_indices(anchor, total_frames, win, win_step):
+    last = total_frames - 1
+    idxs = [anchor + k * win_step for k in range(win)]
+    return [i if i < total_frames else last for i in idxs]
+
+# collect idx for all windows
+def collect_needed(anchor_idxs, total_frames, win, win_step):
+    needed = set()
+    for a in anchor_idxs:
+        for idx in window_indices(a, total_frames, win, win_step):
+            needed.add(idx)
+    return needed
+
+# read frame for needed, read video in sequence
+def cache_needed_frames(vr, needed, size):
+    if not needed:
+        return {}
+    max_idx = max(needed)
+    cache = {}
+
+    last_ok = -1
+    for i in range(max_idx + 1):
+        try:
+            # read next frame
+            frame_rgb = vr.next().asnumpy()  # decord frame RGB (H,W,3)
+        except StopIteration:
+            # video ended early, like cap.read() failed
+            print(f"[DEBUG] cap.read() failed at i={i}, last_ok={last_ok}, max_idx={max_idx}")
+            print(f"[DEBUG] max_cached={max(cache.keys()) if cache else None}, last_ok={last_ok}")
+            break
+        except Exception as e:
+            # decode error
+            print(f"[DEBUG] cap.read() failed at i={i}, last_ok={last_ok}, max_idx={max_idx} | {e}")
+            print(f"[DEBUG] max_cached={max(cache.keys()) if cache else None}, last_ok={last_ok}")
+            break
+
+        last_ok = i
+        if i in needed:
+            # frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR) # keep same as OpenCV: BGR
+            frame = cv2.resize(frame_rgb, (size, size), interpolation=cv2.INTER_AREA)
+            gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY).astype(np.float32)
+            cache[i] = (frame, gray)
+    return cache
+
+# each anchor: read window from cache frames
+def read_window_from_cache(cache, anchor, total_frames, win, win_step):
+    idxs = window_indices(anchor, total_frames, win, win_step)
+    # print(f"window idxs: {idxs}")
+    if not cache:
+        return None
+
+    # fallback = last frame
+    last_avail = max(cache.keys())
+    fallback = cache[last_avail]
+
+    anchor_frame = None
+    gray_seq = []
+    for j, idx in enumerate(idxs):
+        if idx in cache:
+            data = cache[idx]
+            fallback = data
+        else:
+            data = fallback
+
+        frame, gray = data
+        if j == 0:              # window [0] is anchor
+            anchor_frame = frame
+        gray_seq.append(gray)
+
+    if len(gray_seq) == 1:
+        gray_seq.append(gray_seq[0])
+        idxs.append(idxs[0])
+    return anchor_frame, gray_seq, idxs
+
+
+if __name__ == "__main__":
+    video_path = "/home/xinyi/Project/FD-VQA/test_videos/3369925072.mp4"
+    print("video:", video_path)
+
+    vr = VideoReader(video_path, ctx=cpu(0))
+    total_frames = len(vr)
+    if total_frames <= 1:
+        raise RuntimeError(f"Video too short / frame count unavailable: total_frames={total_frames}")
+    print("total_frames:", total_frames)
+
+    size = 224
+    win = 8
+    win_step = 1
+    num_anchors = 16
+
+    anchor_idxs = sample_frames_uniform(total_frames, num_anchors, win=win, win_step=win_step)
+    needed = collect_needed(anchor_idxs, total_frames, win, win_step)
+    print("anchor_idxs:", anchor_idxs)
+    print("needed:", needed)
+
+    cache = cache_needed_frames(vr, needed, size)
+    print("cached:", len(cache), "needed:", len(needed))
+
+    frames_list, grays_list, processed_list = [], [], []
+    for a in anchor_idxs:
+        out = read_window_from_cache(cache, a, total_frames, win, win_step)
+        if out is None:
+            continue
+        anchor_frame, gray_seq, idxs = out
+        frames_list.append(anchor_frame)
+        grays_list.append(gray_seq)
+        processed_list.append(idxs)
+
+
+    print("frames:", len(frames_list))
+    print("one frame shape:", frames_list[0].shape)
+    print("each window length:", len(grays_list[0]) if grays_list else 0)
+    print("window idxs:", processed_list if processed_list else None)

src/train.py

@@ -0,0 +1,604 @@
+import os
+os.environ["DECORD_DUPLICATE_WARNING_THRESHOLD"] = "1.0"
+import argparse
+import csv
+from pathlib import Path
+import numpy as np
+from scipy.stats import pearsonr as _pearsonr_scipy, spearmanr as _spearmanr_scipy
+import torch
+import torch.nn.functional as F
+import cv2
+from decord import VideoReader, cpu
+from torch.amp import GradScaler, autocast
+from tqdm import tqdm
+
+from module.compute_weight_map import process_video, compute_weight_map
+from model.qd_model import QD_MODEL
+
+
+# ----------------------------
+# Data utils
+# ----------------------------
+def read_vid_mos_csv(csv_path):
+    rows = []
+    with open(csv_path, "r", encoding="utf-8") as f:
+        reader = csv.DictReader(f)
+        if not reader.fieldnames:
+            raise RuntimeError("CSV has no header")
+        for r in reader:
+            vid = str(r["vid"]).strip()
+            mos = float(r["mos"])
+            rows.append((vid, mos))
+    return rows
+
+def split_rows(rows, seed=42, test_ratio=0.2, val_ratio=0.1):
+    rng = np.random.default_rng(int(seed))
+    idx = np.arange(len(rows))
+    rng.shuffle(idx)
+
+    n = len(rows)
+    n_test = int(round(n * test_ratio))
+    n_train_all = n - n_test  # train+val
+    n_val = int(round(n_train_all * val_ratio))    # val from train_all
+
+    val = [rows[i] for i in idx[:n_val]]
+    train = [rows[i] for i in idx[n_val:n_train_all]]
+    test = [rows[i] for i in idx[n_train_all:]]
+    return train, val, test
+
+def split_train_val(rows, seed=42, val_ratio=0.1):
+    rng = np.random.default_rng(int(seed))
+    idx = np.arange(len(rows))
+    rng.shuffle(idx)
+
+    n = len(rows)
+    n_val = int(round(n * val_ratio))
+    val = [rows[i] for i in idx[:n_val]]
+    train = [rows[i] for i in idx[n_val:]]
+    return train, val
+
+def pearsonr(x, y, eps=1e-12):
+    # PLCC (SciPy): returns a torch scalar tensor so call-site ".item()" still works
+    if hasattr(x, "detach"):
+        x = x.detach().cpu().numpy()
+    if hasattr(y, "detach"):
+        y = y.detach().cpu().numpy()
+    x = np.asarray(x).reshape(-1)
+    y = np.asarray(y).reshape(-1)
+
+    # avoid NaN when constant / too short
+    if x.size < 2 or np.std(x) < eps or np.std(y) < eps:
+        return torch.tensor(0.0)
+    r, _p = _pearsonr_scipy(x, y)
+    if np.isnan(r):
+        r = 0.0
+    return torch.tensor(float(r))
+
+def spearmanr(x, y, eps=1e-12):
+    # SRCC (SciPy): handles ties correctly; returns torch scalar tensor
+    if hasattr(x, "detach"):
+        x = x.detach().cpu().numpy()
+    if hasattr(y, "detach"):
+        y = y.detach().cpu().numpy()
+
+    x = np.asarray(x).reshape(-1)
+    y = np.asarray(y).reshape(-1)
+    if x.size < 2 or np.std(x) < eps or np.std(y) < eps:
+        return torch.tensor(0.0)
+
+    r, _p = _spearmanr_scipy(x, y)
+    if np.isnan(r):
+        r = 0.0
+    return torch.tensor(float(r))
+
+# ----------------------------
+# Train utils
+# ----------------------------
+def build_scheduler(optim, args):
+    warm = int(args.warmup_epochs)
+    total = int(args.epochs)
+    warm = max(0, min(warm, total - 1))
+
+    # warmup
+    warmup = torch.optim.lr_scheduler.LinearLR(
+        optim,
+        start_factor=0.1,
+        total_iters=warm if warm > 0 else 1,
+    )
+    # cosine warmup
+    cosine = torch.optim.lr_scheduler.CosineAnnealingLR(
+        optim,
+        T_max=(total - warm) if (total - warm) > 0 else 1,
+        eta_min=float(args.min_lr),
+    )
+
+    if warm > 0:
+        return torch.optim.lr_scheduler.SequentialLR(
+            optim,
+            schedulers=[warmup, cosine],
+            milestones=[warm],
+        )
+    return cosine
+
+def com_loss(y_pred, y_true, reg_w=0.6, rank_w=1.0, huber_beta=1.0, margin=0.0):
+    # 1) Huber / SmoothL1
+    if huber_beta is None:
+        reg_loss = F.l1_loss(y_pred, y_true, reduction="mean")
+    else:
+        reg_loss = F.smooth_l1_loss(y_pred, y_true, beta=float(huber_beta), reduction="mean")
+    reg_loss = reg_loss * float(reg_w)
+
+    # 2) pairwise hinge rank
+    B = y_true.shape[0]
+    if B < 2 or float(rank_w) == 0.0:
+        rank_loss = y_pred.new_tensor(0.0)
+        return reg_loss + rank_loss, reg_loss, rank_loss
+    pred_diff = y_pred.unsqueeze(1) - y_pred.unsqueeze(0)  # [B,B]
+    true_diff = y_true.unsqueeze(1) - y_true.unsqueeze(0)  # [B,B]
+    s = torch.sign(true_diff)                               # -1,0,+1
+
+    # y_true = pair, ignore
+    mask = (s != 0).float()
+    # hinge: max(0, margin - s*(pred_i - pred_j))
+    rank_mat = F.relu(float(margin) - s * pred_diff) * mask
+    denom = mask.sum().clamp_min(1.0)
+    rank_loss = rank_mat.sum() / denom
+    rank_loss = rank_loss * float(rank_w)
+
+    total = reg_loss + rank_loss
+    return total, reg_loss, rank_loss
+
+# ----------------------------
+# Dataset
+# ----------------------------
+class VQADataset(torch.utils.data.Dataset):
+    """
+    Returns per item:
+      rgb:  [3, T, H, W] float in [0,1]  (RGB)
+      w_art: [1, T, H, W] float in [0,1]
+      w_str: [1, T, H, W] float in [0,1]
+      y:    scalar float (MOS, optional normalized)
+      vid:  str
+    """
+    def __init__(self, rows, db_path, clip_len, size, win, win_step, mos_mean=None, mos_std=None):
+        self.rows = rows
+        self.db_path = str(db_path)
+        self.clip_len = int(clip_len)
+        self.size = int(size)
+        self.win = int(win)
+        self.win_step = int(win_step)
+        self.mos_mean = mos_mean
+        self.mos_std = mos_std
+
+    def __len__(self):
+        return len(self.rows)
+
+    def __getitem__(self, idx):
+        vid, mos = self.rows[int(idx)]
+        num_anchors = self.clip_len
+        size = self.size
+        win = self.win
+        win_step = self.win_step
+
+        # get video path
+        base_path = Path(self.db_path) / vid
+        video_path = None
+        for ext in ("mp4", "avi", "mkv"):
+            p = Path(str(base_path) + f".{ext}")
+            if p.exists():
+                video_path = str(p)
+                break
+        if video_path is None:
+            raise FileNotFoundError(f"Cannot find {vid} video")
+
+        try:
+            # read video
+            vr = VideoReader(video_path, ctx=cpu(0))
+            frame_all, w_art_all, w_str_all, anchors_kept = process_video(
+                vr,
+                size=size, num_anchors=num_anchors, win=win, win_step=win_step,
+            )
+            frames_np, w_art_np, w_str_np = compute_weight_map(frame_all, w_art_all, w_str_all)
+            # print("frames_np:", frames_np.shape, frames_np.dtype)
+            # print("w_art_np:", w_art_np.shape, w_art_np.dtype)
+            # print("w_str_np:", w_str_np.shape, w_str_np.dtype)
+            # print("anchors_kept:", len(anchors_kept), "example:", anchors_kept[0])
+        except Exception as e:
+            print("\n[DATA ERROR]")
+            print("idx:", idx)
+            print("vid:", vid)
+            raise
+        finally:
+            # release decord video reader
+            try:
+                if vr is not None:
+                    del vr
+            except Exception:
+                pass
+
+        # fixed length sampling
+        T = self.clip_len
+
+        # frames_sel = [cv2.cvtColor(frames_np[i], cv2.COLOR_BGR2RGB) for i in range(T)] # frame BGR -> RGB: [T,H,W,3] -> [3,T,H,W]
+        frames_sel = [frames_np[i] for i in range(T)]  # RGB frames_np
+        rgb = torch.from_numpy(np.stack(frames_sel, axis=0)).float()
+        rgb = rgb.permute(3, 0, 1, 2).contiguous() / 255.0
+
+        # W_art / W_str: [T,H,W] -> [1,T,H,W]
+        w_art = torch.from_numpy(np.stack([w_art_np[i] for i in range(T)], axis=0).astype(np.float32)).unsqueeze(0).float()
+        w_str = torch.from_numpy(np.stack([w_str_np[i] for i in range(T)], axis=0).astype(np.float32)).unsqueeze(0).float()
+
+        # MOS
+        y = float(mos)
+        if self.mos_mean is not None and self.mos_std is not None:
+            y = (y - self.mos_mean) / (self.mos_std + 1e-8)
+        y = torch.tensor(y).float()
+        return rgb, w_art, w_str, y, str(vid)
+
+
+# ----------------------------
+# Train
+# ----------------------------
+@torch.no_grad()
+def _gather_cat(xs):
+    if not xs:
+        return torch.empty(0)
+    return torch.cat(xs, dim=0)
+
+def run_epoch(model, loader, device, *, optim=None, amp=True, mos_mean=None, mos_std=None, desc="", show_pbar=True, log_interval=10):
+    is_train = optim is not None
+    model.train(is_train)
+
+    scaler = getattr(run_epoch, "_scaler", None)
+    if scaler is None:
+        device_type = "cuda" if str(device).startswith("cuda") else "cpu"
+        run_epoch._scaler = GradScaler(device_type, enabled=(amp and device_type == "cuda"))
+        scaler = run_epoch._scaler
+
+    losses = []
+    y_all = []
+    yhat_all = []
+
+    # ---- tqdm progress bar ----
+    it = loader
+    if show_pbar:
+        it = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True)
+
+    for step, (rgb, w_art, w_str, y, vid) in enumerate(it, start=1):
+        rgb = rgb.to(device, non_blocking=True)  # [B,3,T,H,W]
+        w_art = w_art.to(device, non_blocking=True)   # [B,1,T,H,W]
+        w_str = w_str.to(device, non_blocking=True)   # [B,1,T,H,W]
+        y = y.to(device, non_blocking=True).float()  # [B]
+
+        if is_train:
+            optim.zero_grad(set_to_none=True)
+
+        device_type = "cuda" if str(device).startswith("cuda") else "cpu"
+        if is_train:
+            with autocast(device_type=device_type, enabled=(amp and device_type == "cuda")):
+                yhat, _aux = model(rgb, w_art, w_str)  # yhat: [B]
+                loss, loss_reg, loss_rank = com_loss(yhat, y)
+
+            scaler.scale(loss).backward()
+            scaler.step(optim)
+            scaler.update()
+        else:
+            with torch.inference_mode():
+                with autocast(device_type=device_type, enabled=(amp and device_type == "cuda")):
+                    yhat, _aux = model(rgb, w_art, w_str)
+                    loss, loss_reg, loss_rank = com_loss(yhat, y)
+
+        loss_cpu = loss.detach().float().cpu()
+        losses.append(loss_cpu)
+        y_all.append(y.detach().float().cpu())
+        yhat_all.append(yhat.detach().float().cpu())
+
+        # ---- update bar every log_interval steps ----
+        if show_pbar and (step % int(log_interval) == 0 or step == len(loader)):
+            avg_loss_so_far = torch.stack(losses).mean().item()
+            lrs = None
+            if is_train and hasattr(optim, "param_groups") and optim.param_groups:
+                lrs = [pg.get("lr", None) for pg in optim.param_groups]
+
+            postfix = {"loss": f"{avg_loss_so_far:.4f}"}
+            if lrs is not None:
+                postfix["lrs"] = ",".join([f"{x:.2e}" for x in lrs if x is not None])
+            it.set_postfix(postfix)
+
+    y_all = _gather_cat(y_all)
+    yhat_all = _gather_cat(yhat_all)
+
+    # ---- 反标准化：在 MOS 原尺度上算相关系数 ----
+    if mos_mean is not None and mos_std is not None:
+        y_all = y_all * mos_std + mos_mean
+        yhat_all = yhat_all * mos_std + mos_mean
+
+    plcc = pearsonr(y_all, yhat_all).item() if y_all.numel() > 1 else 0.0
+    srcc = spearmanr(y_all, yhat_all).item() if y_all.numel() > 1 else 0.0
+    rmse = torch.sqrt(torch.mean((yhat_all - y_all) ** 2)).item() if y_all.numel() > 0 else 0.0
+
+    avg_loss = torch.stack(losses).mean().item() if losses else 0.0
+    return avg_loss, plcc, srcc, rmse
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    # ----- data -----
+    ap.add_argument("--csv_path", default="/home/xinyi/Project/FD-VQA/metadata/LSVQ_TRAIN_metadata.csv")
+    ap.add_argument("--db_path", default="/media/xinyi/server/LSVQ/")
+    ap.add_argument("--split_seed", type=int, default=42)
+    ap.add_argument("--test_ratio", type=float, default=0.2)
+    ap.add_argument("--val_ratio", type=float, default=0.1)  # train 80%
+    # ----- video processing -----
+    ap.add_argument("--clip_len", type=int, default=16)
+    ap.add_argument("--resize", type=int, default=224)
+    ap.add_argument("--win", type=int, default=6)
+    ap.add_argument("--win_step", type=int, default=1)
+    # ----- runtime -----
+    ap.add_argument("--batch_size", type=int, default=8)
+    ap.add_argument("--num_workers", type=int, default=2)
+    ap.add_argument("--device", type=str, default="cuda")
+    ap.add_argument("--no_amp", action="store_true")
+    # ----- hyperparams -----
+    ap.add_argument("--epochs", type=int, default=35)
+    ap.add_argument("--warmup_epochs", type=int, default=3)
+    ap.add_argument("--lr", type=float, default=1e-5)
+    ap.add_argument("--min_lr", type=float, default=1e-6)
+    ap.add_argument("--finetune_lr", type=float, default=5e-5)
+    ap.add_argument("--weight_decay", type=float, default=1e-2)
+    ap.add_argument("--clip_unfreeze_blocks", type=int, default=4)
+    ap.add_argument("--finetune_last_stage", action="store_true")
+    ap.add_argument("--patience", type=int, default=6)
+    # ----- save -----
+    ap.add_argument("--save_dir", type=str, default="checkpoints")
+    ap.add_argument("--save_name", type=str, default="qd_model.pt")
+
+    args = ap.parse_args()
+    torch.manual_seed(args.split_seed)
+    device = torch.device(args.device)
+    amp = not bool(args.no_amp)
+
+    # ----------------------------
+    # Load rows and split
+    # ----------------------------
+    csv_path = Path(args.csv_path)
+    if csv_path.name == "LSVQ_TRAIN_metadata.csv":
+        # LSVQ official split
+        test_csv = csv_path.parent / "LSVQ_TEST_metadata.csv"
+        if not test_csv.exists():
+            raise FileNotFoundError(f"Cannot find LSVQ test csv: {test_csv}")
+        train_all = read_vid_mos_csv(str(csv_path))
+        test_rows = read_vid_mos_csv(str(test_csv))
+        train_rows, val_rows = split_train_val(
+            train_all,
+            seed=args.split_seed,
+            val_ratio=args.val_ratio,
+        )
+        print(f"[LSVQ split] train={len(train_rows)} val={len(val_rows)} test={len(test_rows)}")
+    elif csv_path.name == "KVQ_TRAIN_metadata.csv":
+        # KVQ challenge split
+        val_csv = csv_path.parent / "KVQ_VAL_metadata.csv"
+        test_csv = csv_path.parent / "KVQ_TEST_metadata.csv"
+        train_rows = read_vid_mos_csv(str(csv_path))
+        val_rows = read_vid_mos_csv(str(val_csv))
+        test_rows = read_vid_mos_csv(str(test_csv))
+        print(f"[KVQ split] train={len(train_rows)} val={len(val_rows)} test={len(test_rows)}")
+    else:
+        # default split for other datasets
+        rows = read_vid_mos_csv(str(csv_path))
+        train_rows, val_rows, test_rows = split_rows(
+            rows,
+            seed=args.split_seed,
+            test_ratio=args.test_ratio,
+            val_ratio=args.val_ratio,
+        )
+        # print("sizes:", len(rows), len(train_rows), len(val_rows), len(test_rows))
+        # print("train first 3:", train_rows[:3])
+        # print("val   first 3:", val_rows[:3])
+        # print("test  first 3:", test_rows[:3])
+
+    # MOS normalization stats from train split
+    mos_train = np.array([mos for _vid, mos in train_rows], dtype=np.float32)
+    mos_mean = float(mos_train.mean()) if len(mos_train) else 0.0
+    mos_std = float(mos_train.std()) if len(mos_train) else 1.0
+    if mos_std <= 1e-8:
+        mos_std = 1.0
+
+    # ----------------------------
+    # DB and datasets
+    # ----------------------------
+    ds_train = VQADataset(
+        train_rows, args.db_path,
+        clip_len=args.clip_len,
+        size=args.resize,
+        win=args.win,
+        win_step=args.win_step,
+        mos_mean=mos_mean,
+        mos_std=mos_std,
+    )
+    ds_val = VQADataset(
+        val_rows, args.db_path,
+        clip_len=args.clip_len,
+        size=args.resize,
+        win=args.win,
+        win_step=args.win_step,
+        mos_mean=mos_mean,
+        mos_std=mos_std,
+    )
+    ds_test = VQADataset(
+        test_rows, args.db_path,
+        clip_len=args.clip_len,
+        size=args.resize,
+        win=args.win,
+        win_step=args.win_step,
+        mos_mean=mos_mean,
+        mos_std=mos_std,
+    )
+
+    pin = str(device).startswith("cuda")
+
+    loader_train = torch.utils.data.DataLoader(
+        ds_train,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        pin_memory=pin,
+        # persistent_workers=(args.num_workers > 0),
+        prefetch_factor=4 if args.num_workers > 0 else None,
+        drop_last=False,
+    )
+    loader_val = torch.utils.data.DataLoader(
+        ds_val,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        pin_memory=pin,
+        drop_last=False,
+    )
+    loader_test = torch.utils.data.DataLoader(
+        ds_test,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        pin_memory=pin,
+        drop_last=False,
+    )
+
+    # ----------------------------
+    # Model
+    # ----------------------------
+    model = QD_MODEL(
+        clip_model="openai/clip-vit-base-patch16",
+    ).to(device)
+
+    # Stage A: freeze CLIP
+    model.freeze_clip_all()
+
+    clip_params_all = []
+    other_params_all = []
+    for name, p in model.named_parameters():
+        if name.startswith("encoder."):
+            clip_params_all.append(p)
+        else:
+            other_params_all.append(p)
+
+    param_groups = []
+    if other_params_all:
+        param_groups.append({"params": other_params_all, "lr": float(args.lr)})
+    if clip_params_all:
+        param_groups.append({"params": clip_params_all, "lr": float(args.finetune_lr)})
+
+    optim = torch.optim.AdamW(param_groups, weight_decay=float(args.weight_decay))
+    scheduler = build_scheduler(optim, args)
+    # ----------------------------
+    # Train loop
+    # ----------------------------
+    save_dir = Path(args.save_dir)
+    save_dir.mkdir(parents=True, exist_ok=True)
+    save_path = save_dir / args.save_name
+    best_path = save_dir / (Path(args.save_name).stem + ".best.pt")
+    best_weights_path = save_dir / (Path(args.save_name).stem + ".best_weights.pt")
+
+    best_val_srcc = -1e18
+    bad_epochs = 0
+    did_unfreeze = False
+
+    for epoch in tqdm(range(1, int(args.epochs) + 1), desc="Epochs", dynamic_ncols=True):
+
+        # Stage B: optional CLIP finetune after warmup
+        if (not did_unfreeze) and bool(args.finetune_last_stage) and epoch == (int(args.warmup_epochs) + 1):
+            model.unfreeze_clip_last_blocks(n_blocks=int(args.clip_unfreeze_blocks), also_unfreeze_ln=True)
+            did_unfreeze = True
+
+            # 不重建 optim, 保留 Adam 状态，只更新 lr
+            if hasattr(optim, "param_groups") and len(optim.param_groups) >= 2:
+                optim.param_groups[0]["lr"] = float(args.lr)
+                optim.param_groups[1]["lr"] = float(args.finetune_lr)
+
+            print(
+                f"[Stage B] Unfroze CLIP last {int(args.clip_unfreeze_blocks)} blocks | "
+                f"lr={float(args.lr)} finetune_lr={float(args.finetune_lr)}"
+            )
+
+        tr_loss, tr_plcc, tr_srcc, tr_rmse = run_epoch(
+            model, loader_train, device,
+            optim=optim,
+            amp=amp,
+            mos_mean=mos_mean,
+            mos_std=mos_std,
+            desc=f"Train e{epoch:03d}",
+            show_pbar=True,
+            log_interval=10,
+        )
+        va_loss, va_plcc, va_srcc, va_rmse = run_epoch(
+            model, loader_val, device,
+            optim=None,
+            amp=amp,
+            mos_mean=mos_mean,
+            mos_std=mos_std,
+            desc=f"Val   e{epoch:03d}",
+            show_pbar=True,
+            log_interval=10,
+        )
+        scheduler.step()
+        print(
+            f"epoch {epoch:03d} | "
+            f"train: loss={tr_loss:.4f} plcc={tr_plcc:.4f} srcc={tr_srcc:.4f} rmse={tr_rmse:.4f} | "
+            f"val: loss={va_loss:.4f} plcc={va_plcc:.4f} srcc={va_srcc:.4f} rmse={va_rmse:.4f}"
+        )
+
+        # Save "last" checkpoint every epoch
+        ckpt = {
+            "epoch": epoch,
+            "model": model.state_dict(),
+            "optim": optim.state_dict(),
+            "mos_mean": mos_mean,
+            "mos_std": mos_std,
+            "args": vars(args),
+            "best_val_srcc": best_val_srcc,
+        }
+        torch.save(ckpt, str(save_path))
+
+        # Save best by val SRCC (higher is better)
+        if va_srcc > best_val_srcc:
+            best_val_srcc = va_srcc
+            bad_epochs = 0
+            ckpt["best_val_srcc"] = best_val_srcc
+            torch.save(ckpt, str(best_path))
+            torch.save(model.state_dict(), str(best_weights_path))
+            print(
+                f"  [best] val_srcc={best_val_srcc:.4f} (val_rmse={va_rmse:.4f}) -> saved "
+                f"{best_path} and {best_weights_path}"
+            )
+        else:
+            bad_epochs += 1
+            if bad_epochs >= int(args.patience):
+                print(
+                    f"[EarlyStop] val_srcc did not improve for {bad_epochs} epochs. "
+                    f"Stop at epoch {epoch}."
+                )
+                break
+    # ----------------------------
+    # Test (load best)
+    # ----------------------------
+    if best_weights_path.exists():
+        sd = torch.load(str(best_weights_path), map_location=device, weights_only=True)
+        model.load_state_dict(sd, strict=True)
+        print(f"Loaded best weights: {best_weights_path}")
+    elif best_path.exists():
+        best = torch.load(str(best_path), map_location=device)
+        model.load_state_dict(best["model"], strict=True)
+        print(f"Loaded best checkpoint: {best_path} (val_srcc={best.get('best_val_srcc', None)})")
+
+    te_loss, te_plcc, te_srcc, te_rmse = run_epoch(
+        model, loader_test, device,
+        optim=None,
+        amp=amp,
+        mos_mean=mos_mean,
+        mos_std=mos_std,
+    )
+    print(f"TEST | loss={te_loss:.4f} plcc={te_plcc:.4f} srcc={te_srcc:.4f} rmse={te_rmse:.4f}")
+
+
+if __name__ == "__main__":
+    main()

src/transfer.py

@@ -0,0 +1,291 @@
+import argparse
+from pathlib import Path
+import numpy as np
+import torch
+
+from train import (
+    read_vid_mos_csv,
+    split_rows,
+    VQADataset,
+    run_epoch,
+    build_scheduler,
+)
+from model.qd_model import QD_MODEL
+
+
+def load_pretrained_weights(model, pretrained_path, device):
+    p = Path(pretrained_path)
+    obj = torch.load(str(p), map_location=device, weights_only=True)
+    # *.pt: dict checkpoint
+    if isinstance(obj, dict) and "model" in obj:
+        model.load_state_dict(obj["model"], strict=True)
+        return obj
+    # best_weights.pt: state_dict
+    else:
+        model.load_state_dict(obj, strict=True)
+        return None
+
+def make_loaders(rows_train, rows_val, rows_test, args, mos_mean, mos_std, device):
+    ds_train = VQADataset(
+        rows_train, args.db_path,
+        clip_len=args.clip_len, size=args.resize, win=args.win, win_step=args.win_step,
+        mos_mean=mos_mean, mos_std=mos_std,
+    )
+    ds_val = VQADataset(
+        rows_val, args.db_path,
+        clip_len=args.clip_len, size=args.resize, win=args.win, win_step=args.win_step,
+        mos_mean=mos_mean, mos_std=mos_std,
+    )
+    ds_test = VQADataset(
+        rows_test, args.db_path,
+        clip_len=args.clip_len, size=args.resize, win=args.win, win_step=args.win_step,
+        mos_mean=mos_mean, mos_std=mos_std,
+    )
+
+    pin = str(device).startswith("cuda")
+    loader_train = torch.utils.data.DataLoader(
+        ds_train, batch_size=args.batch_size, shuffle=True,
+        num_workers=args.num_workers, pin_memory=pin,
+        persistent_workers=(args.num_workers > 0),
+        prefetch_factor=4 if args.num_workers > 0 else None,
+        drop_last=False,
+    )
+    loader_val = torch.utils.data.DataLoader(
+        ds_val, batch_size=args.batch_size, shuffle=False,
+        num_workers=args.num_workers, pin_memory=pin, drop_last=False,
+    )
+    loader_test = torch.utils.data.DataLoader(
+        ds_test, batch_size=args.batch_size, shuffle=False,
+        num_workers=args.num_workers, pin_memory=pin, drop_last=False,
+    )
+    return loader_train, loader_val, loader_test
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    # ----- mode -----
+    ap.add_argument("--mode", choices=["finetune", "test_only"], required=True)
+    ap.add_argument("--pretrained", default="/home/xinyi/Project/FD-VQA/src/checkpoints/kvq/qd_model.best.pt", help="pretrain model path")
+    # ----- data -----
+    ap.add_argument("--csv_path", default="/home/xinyi/Project/FD-VQA/metadata/SHORTS-SDR-DATASET_metadata.csv")
+    ap.add_argument("--db_path", default="/media/xinyi/server/video_dataset/shorts-hdr-dataset/sdr/")
+    ap.add_argument("--split_seed", type=int, default=0)
+    ap.add_argument("--test_ratio", type=float, default=0.2)
+    ap.add_argument("--val_ratio", type=float, default=0.1)  # train 80%
+    # ----- video processing -----
+    ap.add_argument("--clip_len", type=int, default=16)
+    ap.add_argument("--resize", type=int, default=224)
+    ap.add_argument("--win", type=int, default=6)
+    ap.add_argument("--win_step", type=int, default=1)
+    # ----- runtime -----
+    ap.add_argument("--batch_size", type=int, default=8)
+    ap.add_argument("--num_workers", type=int, default=4)
+    ap.add_argument("--device", type=str, default="cuda")
+    ap.add_argument("--no_amp", action="store_true")
+    # ----- finetune hyperparams -----
+    ap.add_argument("--epochs", type=int, default=10)
+    ap.add_argument("--warmup_epochs", type=int, default=1)
+    ap.add_argument("--lr", type=float, default=1e-5)
+    ap.add_argument("--min_lr", type=float, default=1e-6)
+    ap.add_argument("--finetune_lr", type=float, default=5e-5)
+    ap.add_argument("--weight_decay", type=float, default=1e-2)
+    ap.add_argument("--clip_unfreeze_blocks", type=int, default=4)
+    ap.add_argument("--finetune_last_stage", action="store_true")
+    ap.add_argument("--patience", type=int, default=6)
+    # ----- save -----
+    ap.add_argument("--save_dir", type=str, default="checkpoints_transfer")
+    ap.add_argument("--save_name", type=str, default="transfer.pt")
+
+    # ----- MOS normalization control -----
+    # finetune: use target train mean/std
+    ap.add_argument("--test_only_norm", choices=["none", "use_source_ckpt"], default="use_source_ckpt")
+
+    args = ap.parse_args()
+    torch.manual_seed(args.split_seed)
+    device = torch.device(args.device)
+    amp = not bool(args.no_amp)
+
+    # ----------------------------
+    # Load target rows and split
+    # ----------------------------
+    rows = read_vid_mos_csv(args.csv_path)
+    if args.mode == "finetune":
+        csv_path = Path(args.csv_path)
+        if csv_path.name == "KVQ_TRAIN_metadata.csv":
+            # KVQ challenge split
+            val_csv = csv_path.parent / "KVQ_VAL_metadata.csv"
+            test_csv = csv_path.parent / "KVQ_TEST_metadata.csv"
+            train_rows = read_vid_mos_csv(str(csv_path))
+            val_rows = read_vid_mos_csv(str(val_csv))
+            test_rows = read_vid_mos_csv(str(test_csv))
+            print(f"[KVQ split] train={len(train_rows)} val={len(val_rows)} test={len(test_rows)}")
+        else:
+            train_rows, val_rows, test_rows = split_rows(
+                rows, seed=args.split_seed, test_ratio=args.test_ratio, val_ratio=args.val_ratio
+            )
+        mos_train = np.array([mos for _vid, mos in train_rows], dtype=np.float32)
+        mos_mean = float(mos_train.mean()) if len(mos_train) else 0.0
+        mos_std = float(mos_train.std()) if len(mos_train) else 1.0
+        if mos_std <= 1e-8:
+            mos_std = 1.0
+    else:
+        # test_only
+        train_rows, val_rows, test_rows = [], [], rows
+        mos_mean, mos_std = None, None
+
+    # ----------------------------
+    # Build model, load pretrained
+    # ----------------------------
+    model = QD_MODEL(
+        clip_model="openai/clip-vit-base-patch16",
+    ).to(device)
+
+    ckpt = load_pretrained_weights(model, args.pretrained, device)
+    print(f"Loaded pretrained: {args.pretrained}")
+
+    # ----------------------------
+    # Decide normalization for test_only
+    # ----------------------------
+    if args.mode == "test_only":
+        if args.test_only_norm == "use_source_ckpt" and isinstance(ckpt, dict):
+            # 注：这样对 SRCC/PLCC 没影响
+            mos_mean = ckpt.get("mos_mean", None)
+            mos_std = ckpt.get("mos_std", None)
+            if mos_mean is None or mos_std is None:
+                mos_mean, mos_std = None, None
+                print("[warn] pretrained ckpt has no mos_mean/std, fallback to no normalization.")
+            else:
+                print(f"test_only uses source mos_mean/std from ckpt: mean={mos_mean:.4f}, std={mos_std:.4f}")
+        else:
+            print("test_only uses no MOS normalization.")
+
+    # ----------------------------
+    # Mode: test_only
+    # ----------------------------
+    if args.mode == "test_only":
+        ds_test = VQADataset(
+            test_rows, args.db_path,
+            clip_len=args.clip_len, size=args.resize, win=args.win, win_step=args.win_step,
+            mos_mean=mos_mean, mos_std=mos_std,
+        )
+        pin = str(device).startswith("cuda")
+        loader_test = torch.utils.data.DataLoader(
+            ds_test, batch_size=args.batch_size, shuffle=False,
+            num_workers=args.num_workers, pin_memory=pin, drop_last=False,
+        )
+        print("num_test_rows =", len(test_rows))
+        print("len(ds_test)  =", len(ds_test))
+        print("len(loader_test) =", len(loader_test), "batch_size =", args.batch_size)
+
+        te_loss, te_plcc, te_srcc, te_rmse = run_epoch(
+            model, loader_test, device,
+            optim=None, amp=amp, mos_mean=mos_mean, mos_std=mos_std,
+            desc="TestOnly", show_pbar=True
+        )
+        print(f"TEST_ONLY | loss={te_loss:.4f} plcc={te_plcc:.4f} srcc={te_srcc:.4f} rmse={te_rmse:.4f}")
+        return
+
+    # ----------------------------
+    # DataLoaders
+    # ----------------------------
+    loader_train, loader_val, loader_test = make_loaders(
+        train_rows, val_rows, test_rows, args,
+        mos_mean=mos_mean, mos_std=mos_std, device=device
+    )
+
+    # ----------------------------
+    # Mode: finetune
+    # ----------------------------
+    model.freeze_clip_all()
+    did_unfreeze = False
+
+    clip_params_all = []
+    other_params_all = []
+    for name, p in model.named_parameters():
+        if name.startswith("encoder."):
+            clip_params_all.append(p)
+        else:
+            other_params_all.append(p)
+
+    param_groups = []
+    if other_params_all:
+        param_groups.append({"params": other_params_all, "lr": float(args.lr)})
+    if clip_params_all:
+        param_groups.append({"params": clip_params_all, "lr": float(args.finetune_lr)})
+
+    optim = torch.optim.AdamW(param_groups, weight_decay=float(args.weight_decay))
+    scheduler = build_scheduler(optim, args)
+
+    save_dir = Path(args.save_dir)
+    save_dir.mkdir(parents=True, exist_ok=True)
+    last_path = save_dir / args.save_name
+    best_path = save_dir / (Path(args.save_name).stem + ".best.pt")
+    best_weights_path = save_dir / (Path(args.save_name).stem + ".best_weights.pt")
+
+    best_val_srcc = -1e18
+    bad_epochs = 0
+
+    for epoch in range(1, int(args.epochs) + 1):
+        if (not did_unfreeze) and bool(args.finetune_last_stage) and epoch == (int(args.warmup_epochs) + 1):
+            model.unfreeze_clip_last_blocks(n_blocks=int(args.clip_unfreeze_blocks), also_unfreeze_ln=True)
+            did_unfreeze = True
+            print(f"[Finetune] Unfroze CLIP last {int(args.clip_unfreeze_blocks)} blocks")
+
+        tr_loss, tr_plcc, tr_srcc, tr_rmse = run_epoch(
+            model, loader_train, device,
+            optim=optim, amp=amp, mos_mean=mos_mean, mos_std=mos_std,
+            desc=f"FT Train e{epoch:03d}", show_pbar=True
+        )
+        va_loss, va_plcc, va_srcc, va_rmse = run_epoch(
+            model, loader_val, device,
+            optim=None, amp=amp, mos_mean=mos_mean, mos_std=mos_std,
+            desc=f"FT Val   e{epoch:03d}", show_pbar=True
+        )
+        scheduler.step()
+
+        print(
+            f"epoch {epoch:03d} | "
+            f"train: loss={tr_loss:.4f} plcc={tr_plcc:.4f} srcc={tr_srcc:.4f} rmse={tr_rmse:.4f} | "
+            f"val: loss={va_loss:.4f} plcc={va_plcc:.4f} srcc={va_srcc:.4f} rmse={va_rmse:.4f}"
+        )
+
+        ckpt_out = {
+            "epoch": epoch,
+            "model": model.state_dict(),
+            "optim": optim.state_dict(),
+            "mos_mean": mos_mean,
+            "mos_std": mos_std,
+            "args": vars(args),
+            "best_val_srcc": best_val_srcc,
+        }
+        torch.save(ckpt_out, str(last_path))
+
+        if va_srcc > best_val_srcc:
+            best_val_srcc = va_srcc
+            bad_epochs = 0
+            ckpt_out["best_val_srcc"] = best_val_srcc
+            torch.save(ckpt_out, str(best_path))
+            torch.save(model.state_dict(), str(best_weights_path))
+            print(f"  [best] val_srcc={best_val_srcc:.4f} -> saved {best_weights_path}")
+        else:
+            bad_epochs += 1
+            if bad_epochs >= int(args.patience):
+                print(f"[EarlyStop] val_srcc not improved for {bad_epochs} epochs. Stop.")
+                break
+
+    # load best and test
+    if best_weights_path.exists():
+        sd = torch.load(str(best_weights_path), map_location=device, weights_only=True)
+        model.load_state_dict(sd, strict=True)
+        print(f"Loaded best weights: {best_weights_path}")
+
+    te_loss, te_plcc, te_srcc, te_rmse = run_epoch(
+        model, loader_test, device,
+        optim=None, amp=amp, mos_mean=mos_mean, mos_std=mos_std,
+        desc="FT Test", show_pbar=True
+    )
+    print(f"FINETUNE TEST | loss={te_loss:.4f} plcc={te_plcc:.4f} srcc={te_srcc:.4f} rmse={te_rmse:.4f}")
+
+
+if __name__ == "__main__":
+    main()

src/transfer_test_only.py

@@ -0,0 +1,284 @@
+import os
+os.environ["DECORD_DUPLICATE_WARNING_THRESHOLD"] = "1.0"
+import argparse
+import csv
+from pathlib import Path
+import torch
+from torch.amp import autocast
+from tqdm import tqdm
+
+from train import VQADataset, com_loss, pearsonr, read_vid_mos_csv, spearmanr
+from model.qd_model import QD_MODEL
+
+
+def load_checkpoint(ckpt_path, device):
+    ckpt = torch.load(str(ckpt_path), map_location=device, weights_only=True)
+    if isinstance(ckpt, dict) and "model" in ckpt:
+        return {
+            "state_dict": ckpt["model"],
+            "train_mos_mean": ckpt.get("mos_mean"),
+            "train_mos_std": ckpt.get("mos_std"),
+            "train_args": ckpt.get("args", {}),
+            "is_full_checkpoint": True,
+        }
+    if isinstance(ckpt, dict):
+        return {
+            "state_dict": ckpt,
+            "train_mos_mean": None,
+            "train_mos_std": None,
+            "train_args": {},
+            "is_full_checkpoint": False,
+        }
+    raise TypeError(f"Unsupported checkpoint type: {type(ckpt)!r}")
+
+def infer_test_scale(rows):
+    mos_values = [float(mos) for _vid, mos in rows]
+    if not mos_values:
+        raise ValueError("Cannot infer test scale from empty rows")
+
+    lo = min(mos_values)
+    hi = max(mos_values)
+
+    if 0.0 <= lo and hi <= 1.0:
+        return 0.0, 1.0
+    if 1.0 <= lo and hi <= 5.0:
+        return 1.0, 5.0
+    if 0.0 <= lo and hi <= 5.0:
+        return 0.0, 5.0
+    return 0.0, 100.0
+
+def linear_remap(x, src_min, src_max, dst_min, dst_max):
+    src_min = float(src_min)
+    src_max = float(src_max)
+    dst_min = float(dst_min)
+    dst_max = float(dst_max)
+
+    if abs(src_max - src_min) <= 1e-12:
+        raise ValueError("Source scale range must be non-zero")
+
+    return (x - src_min) / (src_max - src_min) * (dst_max - dst_min) + dst_min
+
+def save_predictions_csv(save_path, vids, y_true_raw, pred_train_scale, pred_eval_scale):
+    save_path = Path(save_path)
+    save_path.parent.mkdir(parents=True, exist_ok=True)
+
+    with open(save_path, "w", newline="", encoding="utf-8") as f:
+        writer = csv.writer(f)
+        writer.writerow(["vid", "y_true_raw", "pred_train_scale", "pred_eval_scale"])
+        for vid, y_true, pred_train, pred_eval in zip(
+            vids,
+            y_true_raw.tolist(),
+            pred_train_scale.tolist(),
+            pred_eval_scale.tolist(),
+            strict=False,
+        ):
+            writer.writerow([vid, float(y_true), float(pred_train), float(pred_eval)])
+
+    return save_path
+
+@torch.no_grad()
+def evaluate_and_collect(
+    model,
+    loader,
+    device,
+    *,
+    amp=True,
+    train_mos_mean,
+    train_mos_std,
+    train_scale_min,
+    train_scale_max,
+    test_scale_min,
+    test_scale_max,
+    desc="",
+    show_pbar=True,
+    log_interval=10,
+):
+    model.eval()
+
+    losses = []
+    y_all = []
+    yhat_all = []
+    vids_all = []
+
+    it = loader
+    if show_pbar:
+        it = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True)
+
+    for step, (rgb, w_art, w_str, y, vid) in enumerate(it, start=1):
+        rgb = rgb.to(device, non_blocking=True)
+        w_art = w_art.to(device, non_blocking=True)
+        w_str = w_str.to(device, non_blocking=True)
+        y = y.to(device, non_blocking=True).float()
+
+        device_type = "cuda" if str(device).startswith("cuda") else "cpu"
+        with autocast(device_type=device_type, enabled=(amp and device_type == "cuda")):
+            yhat, _aux = model(rgb, w_art, w_str)
+            loss, _loss_reg, _loss_rank = com_loss(yhat, y)
+
+        losses.append(loss.detach().float().cpu())
+        y_all.append(y.detach().float().cpu())
+        yhat_all.append(yhat.detach().float().cpu())
+        vids_all.extend(list(vid))
+
+        if show_pbar and (step % int(log_interval) == 0 or step == len(loader)):
+            avg_loss_so_far = torch.stack(losses).mean().item()
+            it.set_postfix({"loss": f"{avg_loss_so_far:.4f}"})
+
+    if y_all:
+        y_all = torch.cat(y_all, dim=0)
+        yhat_all = torch.cat(yhat_all, dim=0)
+    else:
+        y_all = torch.empty(0)
+        yhat_all = torch.empty(0)
+
+    y_true_raw = y_all * float(train_mos_std) + float(train_mos_mean)
+    pred_train_scale = yhat_all * float(train_mos_std) + float(train_mos_mean)
+    pred_eval_scale = linear_remap(
+        pred_train_scale,
+        src_min=float(train_scale_min),
+        src_max=float(train_scale_max),
+        dst_min=float(test_scale_min),
+        dst_max=float(test_scale_max),
+    )
+
+    plcc = pearsonr(y_true_raw, pred_eval_scale).item() if y_true_raw.numel() > 1 else 0.0
+    srcc = spearmanr(y_true_raw, pred_eval_scale).item() if y_true_raw.numel() > 1 else 0.0
+    rmse = (
+        torch.sqrt(torch.mean((pred_eval_scale - y_true_raw) ** 2)).item()
+        if y_true_raw.numel() > 0
+        else 0.0
+    )
+    avg_loss = torch.stack(losses).mean().item() if losses else 0.0
+
+    return {
+        "loss": avg_loss,
+        "plcc": plcc,
+        "srcc": srcc,
+        "rmse": rmse,
+        "vids": vids_all,
+        "y_true_raw": y_true_raw,
+        "pred_train_scale": pred_train_scale,
+        "pred_eval_scale": pred_eval_scale,
+    }
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt_path", type=str, default="/home/xinyi/Project/FD-VQA/src/checkpoints/lsvq/qd_model.best.pt")
+    ap.add_argument("--csv_path", type=str, default="/home/xinyi/Project/FD-VQA/metadata/KVQ_metadata.csv")
+    ap.add_argument("--db_path", type=str, default="/media/xinyi/server/video_dataset/KVQ")
+
+    ap.add_argument("--clip_len", type=int, default=16)
+    ap.add_argument("--resize", type=int, default=224)
+    ap.add_argument("--win", type=int, default=6)
+    ap.add_argument("--win_step", type=int, default=1)
+
+    ap.add_argument("--batch_size", type=int, default=8)
+    ap.add_argument("--num_workers", type=int, default=2)
+    ap.add_argument("--device", type=str, default="cuda")
+    ap.add_argument("--no_amp", action="store_true")
+
+    ap.add_argument("--train_scale_min", type=float, default=0.0)
+    ap.add_argument("--train_scale_max", type=float, default=100.0)
+    ap.add_argument("--test_scale_min", type=float, default=1.0)
+    ap.add_argument("--test_scale_max", type=float, default=5.0)
+
+    ap.add_argument("--save_pred_csv", type=str, default="/home/xinyi/Project/FD-VQA/src/transfer_test/transfer_test_only_konvid_1k.csv")
+    args = ap.parse_args()
+
+    device = torch.device(args.device)
+    amp = not bool(args.no_amp)
+    ckpt_info = load_checkpoint(Path(args.ckpt_path), device)
+
+    train_mos_mean = ckpt_info["train_mos_mean"]
+    train_mos_std = ckpt_info["train_mos_std"]
+    if train_mos_mean is None or train_mos_std is None:
+        raise ValueError(
+            "Prefer loading *.best.pt / *.pt, or pass --train_mos_mean and --train_mos_std manually."
+        )
+    if float(train_mos_std) <= 1e-8:
+        raise ValueError("train_mos_std must be > 0")
+
+    rows = read_vid_mos_csv(args.csv_path)
+    if not rows:
+        raise ValueError(f"No rows found in csv: {args.csv_path}")
+
+    if args.test_scale_min is None or args.test_scale_max is None:
+        inferred_test_scale_min, inferred_test_scale_max = infer_test_scale(rows)
+        test_scale_min = inferred_test_scale_min
+        test_scale_max = inferred_test_scale_max
+    else:
+        test_scale_min = float(args.test_scale_min)
+        test_scale_max = float(args.test_scale_max)
+
+    dataset = VQADataset(
+        rows,
+        args.db_path,
+        clip_len=args.clip_len,
+        size=args.resize,
+        win=args.win,
+        win_step=args.win_step,
+        mos_mean=float(train_mos_mean),
+        mos_std=float(train_mos_std),
+    )
+
+    pin = str(device).startswith("cuda")
+    loader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=int(args.batch_size),
+        shuffle=False,
+        num_workers=int(args.num_workers),
+        pin_memory=pin,
+        drop_last=False,
+        prefetch_factor=4 if int(args.num_workers) > 0 else None,
+    )
+
+    model = QD_MODEL(
+        clip_model="openai/clip-vit-base-patch16",
+    ).to(device)
+    model.load_state_dict(ckpt_info["state_dict"], strict=True)
+
+    print(f"Loaded checkpoint: {args.ckpt_path}")
+    print(f"Training normalization: mean={float(train_mos_mean):.6f}, std={float(train_mos_std):.6f}")
+    print(
+        f"Scale mapping: train=[{float(args.train_scale_min):.3f}, {float(args.train_scale_max):.3f}] -> "
+        f"test=[{float(test_scale_min):.3f}, {float(test_scale_max):.3f}]"
+    )
+    print(f"Test rows: {len(rows)}")
+
+    metrics = evaluate_and_collect(
+        model,
+        loader,
+        device,
+        amp=amp,
+        train_mos_mean=float(train_mos_mean),
+        train_mos_std=float(train_mos_std),
+        train_scale_min=float(args.train_scale_min),
+        train_scale_max=float(args.train_scale_max),
+        test_scale_min=float(test_scale_min),
+        test_scale_max=float(test_scale_max),
+        desc="Cross-dataset test",
+        show_pbar=True,
+        log_interval=10,
+    )
+
+    print(
+        "TEST | "
+        f"loss={metrics['loss']:.4f} "
+        f"plcc={metrics['plcc']:.4f} "
+        f"srcc={metrics['srcc']:.4f} "
+        f"rmse={metrics['rmse']:.4f}"
+    )
+
+    if args.save_pred_csv:
+        save_path = save_predictions_csv(
+            args.save_pred_csv,
+            metrics["vids"],
+            metrics["y_true_raw"],
+            metrics["pred_train_scale"],
+            metrics["pred_eval_scale"],
+        )
+        print(f"Saved predictions to: {save_path}")
+
+
+if __name__ == "__main__":
+    main()