pavancm/contrique

Official implementation for "Image Quality Assessment using Contrastive Learning"

Image Quality Assessment using Contrastive Learning

Pavan C. Madhusudana, Neil Birkbeck, Yilin Wang, Balu Adsumilli and Alan C. Bovik

This is the official repository of the paper Image Quality Assessment using Contrastive Learning

Usage

The code has been tested on Linux systems with python 3.6. Please refer to requirements.txt for installing dependent packages.

Running CONTRIQUE

In order to obtain quality score using CONTRIQUE model, checkpoint needs to be downloaded. The following command can be used to download the checkpoint.

wget -L https://utexas.box.com/shared/static/rhpa8nkcfzpvdguo97n2d5dbn4qb03z8.tar -O models/CONTRIQUE_checkpoint25.tar -q --show-progress

Alternatively, the checkpoint can also be downloaded using this link.

Google drive link for the checkpoint link.

Obtaining Quality Scores

We provide trained regressor models in models directory which can be used for predicting image quality using features obtained from CONTRIQUE model. For demonstration purposes, some sample images provided in the sample_images folder.

For blind quality prediction, the following commands can be used.

python3 demo_score.py --im_path sample_images/60.bmp --model_path models/CONTRIQUE_checkpoint25.tar --linear_regressor_path models/CLIVE.save
python3 demo_score.py --im_path sample_images/img66.bmp --model_path models/CONTRIQUE_checkpoint25.tar --linear_regressor_path models/LIVE.save

For Full-reference quality assessment, the folllowing command can be employed.

python3 demos_score_FR.py --ref_path sample_images/churchandcapitol.bmp --dist_path sample_images/img66.bmp --model_path models/CONTRIQUE_checkpoint25.tar --linear_regressor_path models/CSIQ_FR.save

Obtaining CONTRIQUE Features

For calculating CONTRIQUE features, the following commands can be used. The features are saved in '.npy' format.

python3 demo_feat.py --im_path sample_images/60.bmp --model_path models/CONTRIQUE_checkpoint25.tar --feature_save_path features.npy
python3 demo_feat.py --im_path sample_images/img66.bmp --model_path models/CONTRIQUE_checkpoint25.tar --feature_save_path features.npy

Training CONTRIQUE

Download Training Data

Create a directory mkdir training_data to store images used for training CONTRIQUE.

1. KADIS-700k : Download KADIS-700k dataset and execute the supllied codes to generate synthetically distorted images. Store this data in the training_data/kadis700k directory.
2. AVA : Download AVA dataset and store in the training_data/UGC_images/AVA_Dataset directory.
3. COCO : COCO dataset contains 330k images spread across multiple competitions. We used 4 folders training_data/UGC_images/test2015, training_data/UGC_images/train2017, training_data/UGC_images/val2017, training_data/UGC_images/unlabeled2017 for training.
4. CERTH-Blur : Blur dataset images are stored in the training_data/UGC_images/blur_image directory.
5. VOC : VOC images are stored in the training_data/UGC_images/VOC2012 directory.

Training Model

Download csv files containing path to images and corresponding distortion classes.

wget -L https://utexas.box.com/shared/static/124n9sfb27chgt59o8mpxl7tomgvn2lo.csv -O csv_files/file_names_ugc.csv -q --show-progress
wget -L https://utexas.box.com/shared/static/jh5cmu63347auyza37773as5o9zxctby.csv -O csv_files/file_names_syn.csv -q --show-progress

The above files can also be downloaded manually using these links link1, link2 Google drive links link1, link2

For training with a single GPU the following command can be used

python3 train.py --batch_size 256 --lr 0.6 --epochs 25

Training with multiple GPUs using Distributed training (Recommended)

Run the following commands on different terminals concurrently

CUDA_VISIBLE_DEVICES=0 python3 train.py --nodes 4 --nr 0 --batch_size 64 --lr 0.6 --epochs 25
CUDA_VISIBLE_DEVICES=1 python3 train.py --nodes 4 --nr 1 --batch_size 64 --lr 0.6 --epochs 25
CUDA_VISIBLE_DEVICES=2 python3 train.py --nodes 4 --nr 2 --batch_size 64 --lr 0.6 --epochs 25
CUDA_VISIBLE_DEVICES=3 python3 train.py --nodes 4 --nr 3 --batch_size 64 --lr 0.6 --epochs 25

Note that in distributed training, batch_size value will be the number of images to be loaded on each GPU. During CONTRIQUE training equal number of images will be loaded from both synthetic and authentic distortions. Thus in the above example code, 128 images will be loaded on each GPU.

Training Linear Regressor

After CONTRIQUE model training is complete, a linear regressor is trained using CONTRIQUE features and corresponding ground truth quality scores using the following command.

python3 train_regressor.py --feat_path feat.npy --ground_truth_path scores.npy --alpha 0.1

Contact

Please contact Pavan ([email protected]) if you have any questions, suggestions or corrections to the above implementation.

Citation

@article{madhusudana2021st,
  title={Image Quality Assessment using Contrastive Learning},
  author={Madhusudana, Pavan C and Birkbeck, Neil and Wang, Yilin and Adsumilli, Balu and Bovik, Alan C},
  journal={arXiv:2110.13266},
  year={2021}
}

contrique's Issues

About distributed training

Is the distributed training set up by you single-machine multi-card or multi-machine multi-card? I ran into NCCL errors training on two Gpus on the same machine

Data Organization for Model Training: Structure within Dataset Folders

When training models, if the dataset is organized in a designated folder structure such as training_data/UGC_images/blur_image, would the blur_image folder only contain the images themselves or should it mirror the structure of a downloaded dataset, as seen in something like CERTH_ImageBlurDataset/EvaluationSet/DigitalBlurSet/ where additional subfolders or metadata files might also be present?

The link of KADIS-700k dataset is dead

How can I get the KADIS-700k dataset?

interpret the score using the provided weight file

Hi,

I would like to use your code to assess my images for my own research. I followed the README file and download the weight file named CONTRIQUE_checkpoint25.tar and got the scores so far. I'm just curious about how to interpret these scores. is the higher the better?

thank you!

How to train custom dataset

@pavancm Hi all,

I am looking to train a model using custom dataset where I have only image. Can you please help me out with any explanation, idea for training process.
Please tell me step by step for training with my data

Thanks!!

About MS-colorspace

I have tried training for different color space,and get similar result in RGB、Graysacle.But I failed to get a good result in MS-colorspace.Is there are something special I should take into consider when training through MS-colorspace?And is large amount of data for pretraining essential?

[Recommendation] Recommendation to provide link of synthesized image distortion scripts.

In case of custom dataset training requirements, and to avoid downloading KaDID-30K dataset, it's recommended to provide the link of synthesizing image distortions Matlab scripts. Because the dataset is so large(40+G)...

KADIS-700 link is not working

http://database.mmsp-kn.de/kadid-10k-database.html links are not working and alternative link ref_img folder is empty!
this data set is very important in your work, i don't know how to train the network without it.

[Recommendation] Recommendation to provide link of synthesized image distortion scripts.

In case of custom dataset training requirements, and to avoid download KaDID-30K dataset, it's recommended to provide the link of synthesizing iamge distortions Matlab scripts. Because the dataset is so large(40+G)...

How to generate score.npy file

Using the provided datasets i am facing an issue regarding the regressor training. I have trained the model but i don't understand how the scores.npy file is generated

performance of supervised finetune using CONTRIQUE？

Hello，I wonder the performance of supervised finetune using CONTRIQUE encoder compared to imagenet pretrained model, but I can't find such exp in paper.

can you share the results if you have done such exp, Thanks.

岭回归

我想问一下你们是怎么实现，分批次图像质量评估数据集来训练岭回归网络，最终得到一个训练模型的，我只能复现一个批次一个训练模型，怎么把多个批次的模型整合成一个训练模型呢？你们有相关代码提供吗？

[link time out error]

Are there links to google drive or baidu drive provided? Because the link of uteas box is not reachable from China...

How to interpret score values?

Hi, thanks for this work.
I want to use this metric to evaluate the quality of some image for my research, in both FR and NR mode.
I was wondering what is the score range for FR and the score range for NR. How should we interpret this values?
For examples, the NR score for churchandcapitol.bmp is 16.765732, but on image66.bmp which is distorted, the NR score is 56.64854 while using the same linear regressor LIVE for both. On the other hand, with CLIVE I obtain NR score 84.52191 for churchandcapitol.bmp and 23.473944 for image66.bmp.

Higher is better? Lower is better? I've noticed that using different linear regressor the output values are in different ranges.

I also tried to compute the FR between two identical images (churchandcapitol.bmp) and the results is -0.04409191 with CSIQ_FR linear regressor, and 1.8522606 with linear regressor LIVE_FR. Also, when I compare image 33.bmp with itself I obtain the same value of -0.04409191 with CSIQ_FR and 1.8522606 with LIVE_FR.

Could you please provide some insights on the interpretation of the scores?

Thank you for your time.
Have a great day.

How to generate distorted images from kadis700k

I see that the data provided by the official source is in MATLAB format. Do I need to run the generate_kadis700k.m file using MATLAB? How long approximately will it take?

why we need to concat syn and ugc together

Please tell me why we need to concat syn and ugc together
syn_i1 = syn_i1.cuda(non_blocking=True)
ugc_i1 = ugc_i1.cuda(non_blocking=True)
x_i1 = torch.cat((syn_i1,ugc_i1),dim=0)

    #image 2
    syn_i2 = syn_i2.cuda(non_blocking=True)
    ugc_i2 = ugc_i2.cuda(non_blocking=True)
    x_i2 = torch.cat((syn_i2,ugc_i2),dim=0)

How to generate feat.npy for the score regressor training?

I have made custom data set and labels with Matlab code and the below python scripts:

def gen_csv(dir_list,
            csv_path,
            ext=".bmp",
            data_mode="syn",
            write_mode="w"):
    """
    Generate CSV file from dir list
    :param data_mode: syn | ref | ugc
    :param write_mode: w | a+
    """
    if len(dir_list) == 0:
        print("[Err]: empty dir list.")
        return

    print("Data mode: ", data_mode)
    print("Write mode: ", write_mode)

    N_DISTORTIONS = 25
    N_CLS = N_DISTORTIONS * 5 + 1
    print("N_CLS: {:d}".format(N_CLS))
    cnt = 0
    with open(csv_path, write_mode, encoding="utf-8") as f:
        if write_mode != "a+":
            f.write(",Unnamed: 0,File_names,labels\n")
        for dir_path in dir_list:
            print("Processing {:s}...".format(dir_path))
            img_names = [x for x in os.listdir(dir_path) if x.endswith(ext)]
            for img_name in img_names:
                img_path = dir_path + "/" + img_name
                if not os.path.isfile(img_path):
                    continue

                print("processing {:s}".format(img_path))

                items = img_name.split(".")[0].split("_")
                # print(items)

                if len(items) > 3 and data_mode == "syn":
                    distort_type = int(items[-2])
                    distort_level = int(items[-1])

                f.write('{:d},{:d},{:s},"['.format(cnt, cnt, img_path))

                for i in range(N_CLS):
                    if data_mode == "syn":
                        if i == (distort_type - 1) * 5 + (distort_level - 1) + 1:
                            if i == N_CLS - 1:
                                f.write("1.0")
                            else:
                                f.write("1.0, ")
                        elif 0 <= i < N_CLS - 1:
                            f.write("0.0, ")
                        elif i == N_CLS - 1:
                            f.write("0.0")
                    elif data_mode == "ref":
                        if i == 0:
                            f.write("1.0, ")
                        elif 0 <= i < N_CLS - 1:
                            f.write("0.0, ")
                        elif i == N_CLS - 1:
                            f.write("0.0")
                    elif data_mode == "ugc":
                        if 0 <= i < N_CLS - 1:
                            f.write("0.0, ")
                        elif i == N_CLS - 1:
                            f.write("0.0")
                    else:
                        print("[Err]: invalid data mode, data mode should be:"
                              "syn | ref | ugc")
                f.write(']"\n')

                cnt += 1

the loss reduced from 11.7 to 3.4 now after 29 epoches...

How to train the seond stage regressor?

Nan loss occured in the second epoch

[How to train on custom dataset]

Thank you for your excellent work and contribution! I wonder how to train CONTRIQUE on my own datasets, not considering the linear projector training.

About score.pny

Can you share your score.npy file in this repository? So that we can measure the SROCC and reproduce the results of your paper. Thank you very much!

Error when calculating features

When i try to calculate the features using the demo_feat.py for some images of the Koniq dataset , in order to train the regressor i run into this problem
RuntimeError: Given groups=1, weight of size [64, 3, 7, 7], expected input[1, 1, 224, 224] to have 3 channels, but got 1 channels instead
Is there a way to fix it?

Finetuning on another dataset

Hi
Found the paper very useful and trying to Finetune on my own dataset. I wanted to know what the loss values looked like for your dataset. According to the experiments I am running, the loss doesn't converge and is stuck around 3.9-4 even after extensive training.

• my dataset contains around 2.5 L image (100k synthetic + 150k UGC).
• Using a learning rate of 0.001 for later epochs.
  Any Ideas on how to improve the training regimen so that the model converges?

About dist_label in train.py

In your code train.py , you defined dist_label twice in line 46 and 51.
Therefore, i am a bit confused about the definition of dist_label..
Thanks for your great work !

The difference between trainning from image-classification task pretrained model and training from the start

When training, I find training with pretrained resnet50 model takes less trainning time.I want to know why the article choose to training from the start without pretrained?

The csv files are empty