SCA-IQA Implementation

TF-Keras implementation of sca_iqa_net as described in [Explore Spatial and Channel Attention In Image Quality Assessment] by Junyong You and Jie Yan.

Installation

Clone this repository.
Install required Python packages. The code is developed by PyCharm in Python 3.7. The requirements.txt document is generated by PyCharm.
The models have actually been trained on RTX3090 by Tensorflow 2.5.1.

Training a model

Examples of training PHIQnet on KonIQ-10k and SPAQ databases can be found in image_quality/bin/train_pipeline.py. Argparser should be used, but the authors prefer to use dictionary with parameters being defined. It is easy to convert to take arguments. In principle, the following parameters can be defined:

args = {}
args['multi_gpu'] = 1  # Set if multi-GPU training should be used or not, 1: multi-GPU, 0: single GPU
args['gpu'] = 0  # If single GPU is used, choose which GPU on multi-GPU PC

# Choose between 'resnet50', 'densnet121', 'vgg16', 'efficientnetb0', 'efficientnetb4'
args['backbone'] = 'resnet50'
# args['backbone'] = 'densenet121'
# args['backbone'] = 'vgg16'
# args['backbone'] = 'efficientnetb0'
# args['backbone'] = 'efficientnetb4'

# Depending on which backbone is used, choose the corresponding ImageNet pretrained weights file, set to None is no pretrained weights to be used.
# args['weights'] = r'..\pretrained_weights\vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5'
# args['weights'] = r'..\pretrained_weights\efficientnetb0_notop.h5'
# args['weights'] = r'..\pretrained_weights\efficientnetb4_notop.h5'
# args['weights'] = r'..\pretrained_weights\resnet50_weights_tf_dim_ordering_tf_kernels_notop.h5'
# args['weights'] = r'..\pretrained_weights\densenet121_weights_tf_dim_ordering_tf_kernels_notop.h5'
# args['weights'] = None

database = 'koniq' # 'spaq'
if database == 'koniq':
    # Image and score must be provided
    args['images_scores_file'] = r'..\databases\train_val_test_koniq.pkl'
    args['image_folder'] = r'..\databases\koniq_all'
    args['n_quality_levels'] = 5 # 1 Choose 1 for KonIQ-MOS or 5 for KonIQ 

elif database == 'spaq':
    # Image and score must be provided
    args['images_scores_file'] = r'..\databases\train_val_test_spaq.pkl'
    args['image_folder'] = r'..\databases\koniq_all'
    args['n_quality_levels'] = 1

args['result_folder'] = r'..\results\{}_koniq_sca_{}'.format(
    database, args['backbone'])

args['initial_epoch'] = 0

args['lr_base'] = 1e-4  # 1e-4/2  # change base learning rate dependeing on the batch size
args['lr_schedule'] = True
args['batch_size'] = 32
args['epochs'] = 100

args['image_aug'] = True

args['do_finetune'] = True  # Flag to do fine tune or not

Predict image quality using the trained model

After SCA-IQA has been trained, and the weights have been stored in h5 file, it can be used to predict image quality with arbitrary sizes,

    args = {}
    args['n_quality_levels'] = 5
    args['backbone'] = 'resnet50'
    args['weights'] = 'sca.h5'
    model = scaq_iqa_net(n_quality_levels=args['n_quality_levels'],
                     backbone=args['backbone'])
    model.load_weights(args['weights'])

Please see the scripts in model_evaluation.

Prepare datasets for model training

This work uses two publicly available databases: KonIQ-10k KonIQ-10k: An ecologically valid database for deep learning of blind image quality assessment by V. Hosu, H. Lin, T. Sziranyi, and D. Saupe; and SPAQ Perceptual quality assessment of smartphone photography by Y. Fang, H. Zhu, Y. Zeng, K. Ma, and Z. Wang.

Please see README in databases for details.

Report bugs

Please check the paths of image, weights, etc. first if encountering any problems. Please report any bugs in Issues.

FAQs

To be added

junyongyou / sca_iqa