Light-VQA+

This work is the extension of "Light-VQA: A Multi-Dimensional Quality Assessment Model for Low-Light Video Enhancement", which can be found at Light-VQA

Usgae

Datasets And Pretrained Weights

The VEC-QA dataset used in this work can be downloaded at:

JBOX

BaiduNetdisk (tqw0)

Our model trained on VEC-QA can be downloaded at:

JBOX

BaiduNetdisk (ui9v).

To use our pretrained model, please follow the instructions below.

Install Requirements

Python Version: 3.8

  pip install openai-clip
  pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2
  pip install pandas
  pip install scipy
  pip install opencv-python
  pip install pytorchvideo

Please run the code in your shell to build the environment for our model.

Training

After downloading our VEC-QA dataset or building your own dataset under the same directory, run the following codes in your shell:

step 1

download the swin-transformer weights into the directory at here.

step 2

get the key_frames of all the videos.

  python extract_key_frames.py

You may need to reset the 'videos_dir' in codes first so as to match your own dataset. (line 52)

step 3

get the brightness and noise features and brightness consistency features.

  python extract_BNF.py
  python extract_BCF.py

You may need to reset the 'videos_dir' in codes first so as to match your own dataset. (line 32 for BNF and line 161 for BCF)

step 4

get the temporal features.

  python extract_temporal_features.py

You may need to reset the 'videos_dir' and 'datainfo' in codes first so as to match your own dataset. (line 85-86)

The process in extracting these features is carried out on train-val-test splits separately, so you may need to prepare the csv file first. Just refer to the csv files in data.

step 5

start training.

  python train.py > final.txt

The parameters can be set at line 195-222. Remember to change the data path in line 61-64 first.

You should keep track of the output during the training process for it contains the attention weights information, which can be found at the final.txt output files.

Testing

  python test.py

Before starting testing, you need to set the attention weights in line 22 and the ckpts path in line 25.