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[source]
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)[source]

Bases: torch.nn.Module

Neural IMage Assessment (NIMA) 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 (str) – Backbone model name supported by timm.

  • train_dataset (str) – Dataset tag used to select default checkpoint.

  • num_classes (int) – Number of output classes for score distribution.

  • dropout_rate (float) – Dropout ratio in classifier head.

  • pretrained (bool) – Whether to load pretrained NIMA weights.

  • pretrained_model_path (str | None) – Optional local checkpoint path.

Notes

Typical input sizes are 224 (VGG/MobileNet) or 299 (Inception-family backbones).

preprocess(x)[source]

Resize/crop (eval mode) and normalize input image tensor.

forward(x, return_mos=True, return_dist=False)[source]

Compute image quality using NIMA.

Parameters:

  • x (torch.Tensor) – Input tensor with shape (N, 3, H, W).

  • return_mos (bool) – Whether to return MOS prediction.

  • return_dist (bool) – Whether to return score distribution.

Returns:

One tensor when a single output is requested, otherwise [mos, dist].

Return type:

torch.Tensor | list[torch.Tensor]