pyiqa.archs.maniqa_arch
=======================

.. py:module:: pyiqa.archs.maniqa_arch

.. autoapi-nested-parse::

   MANIQA proposed by

       MANIQA: Multi-dimension Attention Network for No-Reference Image Quality Assessment
       Sidi Yang, Tianhe Wu, Shuwei Shi, Shanshan Lao, Yuan Gong, Mingdeng Cao, Jiahao Wang and Yujiu Yang.
       CVPR Workshop 2022, winner of NTIRE2022 NRIQA challenge

   Reference:
       - Official github: https://github.com/IIGROUP/MANIQA







Module Contents
---------------

.. py:data:: default_model_urls

.. py:class:: TABlock(dim, drop=0.1)

   Bases: :py:obj:`torch.nn.Module`


   Token-attention block used in MANIQA stages.


   .. py:method:: forward(x)


.. py:class:: SaveOutput

   Forward-hook collector for intermediate ViT block outputs.


   .. py:method:: clear()


.. py:class:: MANIQA(embed_dim=768, num_outputs=1, patch_size=8, drop=0.1, depths=[2, 2], window_size=4, dim_mlp=768, num_heads=[4, 4], img_size=224, num_tab=2, scale=0.13, test_sample=20, pretrained=True, pretrained_model_path=None, train_dataset='pipal', default_mean=None, default_std=None, **kwargs)

   Bases: :py:obj:`torch.nn.Module`


   MANIQA no-reference IQA model.

   :param embed_dim: Embedding dimension.
   :type embed_dim: int
   :param num_outputs: Number of output channels.
   :type num_outputs: int
   :param patch_size: Patch size used by ViT backbone.
   :type patch_size: int
   :param drop: Dropout ratio for prediction heads.
   :type drop: float
   :param depths: Depths of Swin blocks.
   :type depths: list[int]
   :param window_size: Swin attention window size.
   :type window_size: int
   :param dim_mlp: MLP dimension used in Swin blocks.
   :type dim_mlp: int
   :param num_heads: Number of attention heads in Swin blocks.
   :type num_heads: list[int]
   :param img_size: Input crop size.
   :type img_size: int
   :param num_tab: Number of token-attention blocks per stage.
   :type num_tab: int
   :param scale: Swin scaling factor.
   :type scale: float
   :param test_sample: Number of evaluation crops.
   :type test_sample: int
   :param pretrained: Whether to load pretrained model weights.
   :type pretrained: bool
   :param pretrained_model_path: Optional local checkpoint path.
   :type pretrained_model_path: str | None
   :param train_dataset: Checkpoint key for pretrained loading.
   :type train_dataset: str
   :param default_mean: Optional custom normalization mean.
   :type default_mean: torch.Tensor | None
   :param default_std: Optional custom normalization std.
   :type default_std: torch.Tensor | None
   :param \*\*kwargs: Reserved compatibility arguments.


   .. py:method:: extract_feature(save_output)

      Concatenate selected ViT block tokens into MANIQA feature tensor.



   .. py:method:: forward(x)

      Predict image quality score.

      :param x: Input tensor with shape ``(N, 3, H, W)``.
      :type x: torch.Tensor

      :returns: Predicted score tensor with shape ``(N, 1)``.
      :rtype: torch.Tensor