pyiqa.archs.nima_arch
=====================

.. py:module:: pyiqa.archs.nima_arch

.. autoapi-nested-parse::

   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
---------------

.. py:data:: default_model_urls

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

   Bases: :py:obj:`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.

   :param 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
   :param default input shape:
                               - vgg and mobilenet: (N, 3, 224, 224)
                               - inception: (N, 3, 299, 299)


   .. py:method:: preprocess(x)


   .. py:method:: forward(x, return_mos=True, return_dist=False)

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