pyiqa.archs.nima_arch

NIMA model. Reference:

Talebi, Hossein, and Peyman Milanfar. “NIMA: Neural image assessment.” IEEE transactions on image processing 27, no. 8 (2018): 3998-4011.

Created by: https://github.com/yunxiaoshi/Neural-IMage-Assessment/blob/master/model/model.py

Modified by: Chaofeng Chen (https://github.com/chaofengc)

Module Contents

pyiqa.archs.nima_arch.default_model_urls

class pyiqa.archs.nima_arch.NIMA(base_model_name='vgg16', train_dataset='ava', num_classes=10, dropout_rate=0.0, pretrained=True, pretrained_model_path=None)

Bases: torch.nn.Module

Neural IMage Assessment model.

Modification:

• for simplicity, we use global average pool for all models

• we remove the dropout, because parameters with avg pool is much less.

Parameters:

• base_model_name – pretrained model to extract features, can be any models supported by timm. Models used in the paper: vgg16, inception_resnet_v2, mobilenetv2_100

• shape (default input) –

  • vgg and mobilenet: (N, 3, 224, 224)

  • inception: (N, 3, 299, 299)

preprocess(x)

forward(x, return_mos=True, return_dist=False)

Computation image quality using NIMA. :param x: An input tensor. Shape \((N, C, H, W)\). :param return_mos: Whether to return mos_score. :param retuen_dist: Whether to return dist_score.