feat: 🚩 add input validation check

Author: chaofengc

pyiqa/models/inference_model.py

@@ -63,6 +63,13 @@ def __init__(
     def load_weights(self, weights_path, weight_keys='params'):
         self.net = load_pretrained_network(self.net, weights_path, weight_keys=weight_keys)
 
+    def is_valid_input(self, x):
+        if x is not None:
+            assert isinstance(x, torch.Tensor), 'Input must be a torch.Tensor'
+            assert x.dim() == 4, 'Input must be 4D tensor (B, C, H, W)'
+            assert x.shape[1] in [1, 3], 'Input must be RGB or gray image'
+            assert x.min() >= 0 and x.max() <= 1, f'Input must be normalized to [0, 1], but got min={x.min():.4f}, max={x.max():.4f}'
+    
     def forward(self, target, ref=None, **kwargs):
         device = self.dummy_param.device
 
@@ -80,6 +87,9 @@ def forward(self, target, ref=None, **kwargs):
                         assert ref is not None, 'Please specify reference image for Full Reference metric'
                         ref = imread2tensor(ref, rgb=True)
                         ref = ref.unsqueeze(0)
+                
+                self.is_valid_input(target)
+                self.is_valid_input(ref)
 
                 if self.metric_mode == 'FR':
                     assert ref is not None, 'Please specify reference image for Full Reference metric'