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

   :param base_model_name: Backbone model name supported by :mod:`timm`.
   :type base_model_name: str
   :param train_dataset: Dataset tag used to select default checkpoint.
   :type train_dataset: str
   :param num_classes: Number of output classes for score distribution.
   :type num_classes: int
   :param dropout_rate: Dropout ratio in classifier head.
   :type dropout_rate: float
   :param pretrained: Whether to load pretrained NIMA weights.
   :type pretrained: bool
   :param pretrained_model_path: Optional local checkpoint path.
   :type pretrained_model_path: str | None

   .. rubric:: Notes

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


   .. py:method:: preprocess(x)

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



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

      Compute image quality using NIMA.

      :param x: Input tensor with shape ``(N, 3, H, W)``.
      :type x: torch.Tensor
      :param return_mos: Whether to return MOS prediction.
      :type return_mos: bool
      :param return_dist: Whether to return score distribution.
      :type return_dist: bool

      :returns: One tensor when a single output is
                requested, otherwise ``[mos, dist]``.
      :rtype: torch.Tensor | list[torch.Tensor]