avinabsaha/reiqa

Official implementation for CVPR2023 Paper "Re-IQA : Unsupervised Learning for Image Quality Assessment in the Wild"

Re-IQA: Unsupervised Learning for Image Quality Assessment in the Wild

Official PyTorch implementation of Re-IQA, a No-Reference Image Quality Assessment algorithm proposed in IEEE/CVF CVPR 2023. Re-IQA achieves SoTA performance across popular NR-IQA databases : KonIQ, FLIVE, SPAQ, CLIVE, LIVE-IQA, CSIQ-IQA, TID-2013 and KADID.

Useful Links : Preprint OpenAccess Supplementary YouTube Video Slides Poster

Usage

The code has been tested on Linux systems with python 3.9, pytorch 1.13.1, torchvision 0.14.1 and CUDA 11.7. Other dependencies can be installed via requirements.txt.

Training Re-IQA

Quality Aware Module

Download Training Data

We follow the data training pipeline provided in CONTRIQUE. Create a directory mkdir training_data to store images used for training RE-IQA-Quality-Aware Module.

1. KADIS-700k : Download KADIS-700k dataset. Store this data in the training_data/kadis700k directory. Note we do not use the synthetically distorted images as in CONTRIQUE.
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. We only use the images labelled as "naturally-blurred" from the training (220) and evaluation (411) sets. We have renumbered the images in the provided csv file.
5. VOC : VOC images are stored in the training_data/UGC_images/VOC2012 directory.

Training Command

We use DDP Training across 6 nodes.

python main_contrast.py --method MoCov2 --cosine --head mlp --multiprocessing-distributed --csv_path ./csv_files/moco_train.csv --model_path ./expt0  --optimizer LARS --tb_path ./expt0 -j 28 --batch_size 630 --learning_rate 12 --n_aug 11  --epochs 40 --n_scale 2 --n_distortions 1 --patch_size 160 --world-size 6 --warm --swap_crops 1 --dist-url tcp://[your first node address]:[specified port] --rank 0

On the 5 other nodes, run the same command with --rank 1-5

Clarifications : In the paper submitted to CVPR, we had reported training for 25 epochs. However, we observed more stabilized IQA results when trained for 40 epochs. Thus, we recommend training for 40 epochs. Also, the earlier implementation of the paper, a learning rate of 0.6 was used with batch size scaling. The learning rate scaling formula based on batch size used was ((batch_size * 2 scales) / 64 * 0.6 ≈ 12). However, in the current code, we have simplified it to directly take the scaled learning rate as an input rather than calculating it using the scaling formula.

Obtaining Quality Aware Features

Download the quality-aware trained model from Google Drive and store in a folder named re-iqa_ckpts. Finally, to obtain Re-IQA Quality-Aware features, run

python demo_quality_aware_feats.py --head mlp

Content Aware Module

We utilized the vanilla MoCo-v2 training code using ResNet-50 architecture and ImageNet database provided in the PyContrast repository to train our Content Aware Module using the default settings.

Obtaining Content Aware Features

Download the content-aware trained model from Google Drive and store in a folder named re-iqa_ckpts. Finally, to obtain Re-IQA Content-Aware features, run

python demo_content_aware_feats.py --head mlp

Training Linear Regressor

We used Sklearn's Linear Regression Models with Regularization for training the final IQA model. We recommend using either one of Ridge or Elastic Net models and performing extensive hyper-parameter search for each database to extract maximum performance.

Citation

If you use this code for your research, please cite the following paper:

A. Saha, S. Mishra, and A. C. Bovik, “Re-IQA : Unsupervised Learning for Image Quality Assessment in the Wild,” Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2023, https://doi.org/10.48550/arXiv.2304.00451.

@InProceedings{Saha_2023_CVPR,
    author    = {Saha, Avinab and Mishra, Sandeep and Bovik, Alan C.},
    title     = {Re-IQA: Unsupervised Learning for Image Quality Assessment in the Wild},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2023},
    pages     = {5846-5855}
}

Acknowledgements

Our Code is based on PyContrast. We thank the Yonglong Tian for making the code available.

Contacts

• Avinab Saha ( [email protected] ) -- Graduate student, LIVE, Dept. of ECE, UT Austin.
• Sandeep Mishra ( [email protected] ) -- Graduate student, LIVE, Dept. of ECE, UT Austin.

reiqa's Issues

Why Linear Regressor？

Good work！I'm interested in your work！and I'd like to ask a question that I hope will be answered. My understanding is that the relationship between the two features and the quality score should be non-linear. Why is a linear regression recommended? And can you please share your regressor code?

how can i get the final score？

I'm highly positive about your work and very interested. I would like to ask some questions and hope you can help me! I would like to obtain a final quantitative metric score.

VOC dataset

Could you please provide more detailed information on how you obtained the VOC dataset? It seems like you downloaded the VOC2012 dataset from the link you shared. You appear to have used the data inside the 'JPEGImages' folder. However, that folder contains only 17,125 images, while you seem to have used a total of 33,260 images. How did you obtain the remaining images? For example, files like '2008_000001.jpg' cannot be found in any folder. It would be greatly appreciated if you could clarify this.

How to get the final score?

Hi

thank you for your work. I wonder if there is a model to get the final scores instead of the features?

Thank you!

Fix Dependency Name in requirements.txt

Thanks for sharing your very interesting work!

I have encountered an installation issue from requirements.txt:

ERROR: No matching distribution found for sklearn==1.1.3

Changing sklearn==1.1.3 to scikit-learn==1.1.3, solved it.
So, I am suggesting updating requirements.txt to prevent this error for future installations.

Thanks again for the great work!
Kind Regards,
Joao

Is it truly unsupervised or more of pretraining for quality estimation?

The paper seems to use linear regression to map the pre-trained embeddings to the visual quality. I am wondering if it is right to call it unsupervised learning for image quality estimation?

could you provide an example of test code with a trained model?

Hi, thanks for your sharing of the code.
I found the test_options.py in the current code,
but I cannot find a test code with a trained model
(only training code provided)

How to train on a single server with multiple GPUs

Hi, thanks for sharing this beautiful work! I don't have different nodes, could you help to offer a script to train model on a single server with multiple GPUs? :)

Suggestions on the final dataset size for Linear Regressor

Hi! Thanks for sharing this awesome work! Wondering do you guys have some idea on how much data do we needed to train the final linear regressor?