feat: 📝 add maniqa-koniq, kadid

chaofengc · Apr 3, 2023 · fe95923 · fe95923
1 parent ac70d7a
commit fe95923
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/pyiqa/archs/maniqa_arch.py b/pyiqa/archs/maniqa_arch.py
@@ -10,7 +10,6 @@
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 import timm
 
 from timm.models.vision_transformer import Block
@@ -20,11 +19,12 @@
 from einops import rearrange
 
 from pyiqa.utils.registry import ARCH_REGISTRY
-from .func_util import extract_2d_patches
 from pyiqa.archs.arch_util import load_pretrained_network
 
 default_model_urls = {
-    'pipal': 'https://github.com/chaofengc/IQA-PyTorch/releases/download/v0.1-weights/MANIQA_PIPAL-ae6d356b.pth'
+    'pipal': 'https://github.com/chaofengc/IQA-PyTorch/releases/download/v0.1-weights/MANIQA_PIPAL-ae6d356b.pth',
+    'koniq': 'https://github.com/IIGROUP/MANIQA/releases/download/Koniq10k/ckpt_koniq10k.pt',
+    'kadid': 'https://github.com/IIGROUP/MANIQA/releases/download/Kadid10k/ckpt_kadid10k.pt',
 }
 
 
@@ -83,6 +83,7 @@ def __init__(self, embed_dim=768, num_outputs=1, patch_size=8, drop=0.1,
                  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):
@@ -155,7 +156,7 @@ def __init__(self, embed_dim=768, num_outputs=1, patch_size=8, drop=0.1,
             load_pretrained_network(self, pretrained_model_path, True, weight_keys='params')
             # load_pretrained_network(self, pretrained_model_path, True, )
         elif pretrained:
-            load_pretrained_network(self, default_model_urls['pipal'], True)
+            load_pretrained_network(self, default_model_urls[train_dataset], True)
 
     def extract_feature(self, save_output):
         x6 = save_output.outputs[6][:, 1:]

diff --git a/pyiqa/default_model_configs.py b/pyiqa/default_model_configs.py
@@ -261,6 +261,20 @@
         },
         'metric_mode': 'NR',
     },
+    'maniqa-koniq': {
+        'metric_opts': {
+            'type': 'MANIQA',
+            'train_dataset': 'koniq',
+        },
+        'metric_mode': 'NR',
+    },
+    'maniqa-kadid': {
+        'metric_opts': {
+            'type': 'MANIQA',
+            'train_dataset': 'kadid',
+        },
+        'metric_mode': 'NR',
+    },
     'clipiqa': {
         'metric_opts': {
             'type': 'CLIPIQA',