Clean some codes.

chaofengc · Feb 18, 2022 · b2b5ad1 · b2b5ad1
1 parent 6377322
commit b2b5ad1
Show file tree

Hide file tree

Showing 12 changed files with 23 additions and 844 deletions.
diff --git a/pyiqa/archs/arch_util.py b/pyiqa/archs/arch_util.py
@@ -1,16 +1,11 @@
-import collections.abc
 import math
+
 import torch
-import torchvision
-import warnings
-from distutils.version import LooseVersion
-from itertools import repeat
 from torch import nn as nn
 from torch.nn import functional as F
 from torch.nn import init as init
 from torch.nn.modules.batchnorm import _BatchNorm
 
-from pyiqa.utils import get_root_logger 
 from pyiqa.utils.download_util import load_file_from_url
 
 # --------------------------------------------

diff --git a/pyiqa/archs/brisque_arch.py b/pyiqa/archs/brisque_arch.py
@@ -16,8 +16,8 @@
 import torch
 import torch.nn.functional as F
 from pyiqa.utils.color_util import to_y_channel
-from pyiqa.utils.matlab_functions  import fspecial_gauss, imresize
-from .func_util import estimate_ggd_param, estimate_aggd_param, safe_sqrt, normalize_img_with_guass
+from pyiqa.utils.matlab_functions  import imresize
+from .func_util import estimate_ggd_param, estimate_aggd_param, normalize_img_with_guass
 from pyiqa.utils.download_util import load_file_from_url
 from pyiqa.utils.registry import ARCH_REGISTRY
 

diff --git a/pyiqa/archs/ckdn_arch.py b/pyiqa/archs/ckdn_arch.py
@@ -14,7 +14,6 @@
 import torchvision as tv
 from pyiqa.utils.registry import ARCH_REGISTRY
 from pyiqa.archs.arch_util import load_pretrained_network
-from pyiqa.utils.download_util import load_file_from_url
 
 try:
     from torch.hub import load_state_dict_from_url

diff --git a/pyiqa/archs/dbcnn_arch.py b/pyiqa/archs/dbcnn_arch.py
@@ -6,8 +6,6 @@
 
 """
 
-import os
-
 import torch
 import torchvision
 import torch.nn as nn

diff --git a/pyiqa/archs/func_util.py b/pyiqa/archs/func_util.py
@@ -8,8 +8,8 @@
 
 
 def torch_cov(tensor, rowvar=True, bias=False):
-    """Estimate a covariance matrix (np.cov)
-    https://gist.github.com/ModarTensai/5ab449acba9df1a26c12060240773110
+    r"""Estimate a covariance matrix (np.cov)
+    Ref: https://gist.github.com/ModarTensai/5ab449acba9df1a26c12060240773110
     """
     tensor = tensor if rowvar else tensor.transpose(-1, -2)
     tensor = tensor - tensor.mean(dim=-1, keepdim=True)
@@ -18,6 +18,11 @@ def torch_cov(tensor, rowvar=True, bias=False):
 
 
 def safe_sqrt(x: torch.Tensor) -> torch.Tensor:
+    r"""Safe sqrt with EPS to ensure numeric stability.
+
+    Args:
+        x (torch.Tensor): should be non-negative 
+    """
     EPS = torch.finfo(x.dtype).eps
     return torch.sqrt(x + EPS)
 

diff --git a/pyiqa/archs/inception.py b/pyiqa/archs/inception.py
diff --git a/pyiqa/archs/nima_arch.py b/pyiqa/archs/nima_arch.py
@@ -1,4 +1,7 @@
 r"""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
 
@@ -28,9 +31,6 @@ class NIMA(nn.Module):
         default input shape:
             - vgg and mobilenet: (N, 3, 224, 224)
             - inception: (N, 3, 299, 299)
-    Reference:
-        Talebi, Hossein, and Peyman Milanfar. "NIMA: Neural image assessment." 
-        IEEE transactions on image processing 27, no. 8 (2018): 3998-4011.
     """
     def __init__(self, 
                  base_model_name='vgg16',