Add test methods to RCAN and SRGAN models

models/RCAN/rcan.py

@@ -118,6 +118,15 @@ class RCAN(nn.Module):
             x = self.forward(x)
             x = Image.fromarray((x.squeeze(0).permute(1, 2, 0).detach().numpy() * 255).astype('uint8'))
         return x
+    
+    def test(self, x):
+        """
+        x is a tensor
+        """
+        self.eval()
+        with torch.no_grad():
+            x = self.forward(x)
+        return x
 
 if __name__ == '__main__':
     current_dir = os.path.dirname(os.path.realpath(__file__))

models/SRFlow/srflow.py

@@ -39,7 +39,41 @@ def return_SRFlow_result(lr, conf_path='models/SRFlow/code/confs/SRFlow_DF2K_4X.
     sr = Image.fromarray((sr).astype('uint8'))
     return sr
 
+def test(lr, conf_path='models/SRFlow/code/confs/SRFlow_DF2K_4X.yml', heat=0.7):
+    """
+    Apply Super-Resolution using SRFlow model to the input LR (low-resolution) image.
+
+    Args:
+    - lr: tensor
+    - conf_path (str): Configuration file path for the SRFlow model. Default is SRFlow_DF2K_4X.yml.
+    - heat (float): Heat parameter for the SRFlow model. Default is 0.6.
+
+    Returns:
+    - sr: tensor
+    """
+    model, opt = load_model(conf_path)
+
+    scale = opt['scale']
+    pad_factor = 2
+
+    h, w, c = lr.shape
+    lr = impad(lr, bottom=int(np.ceil(h / pad_factor) * pad_factor - h),
+                        right=int(np.ceil(w / pad_factor) * pad_factor - w))
+    
+    lr_t = t(lr)
+    heat = opt['heat']
+    
+    sr_t = model.get_sr(lq=lr_t, heat=heat)
+
+    sr = rgb(torch.clamp(sr_t, 0, 1))
+    sr = sr[:h * scale, :w * scale]
+
+    sr = Image.fromarray((sr).astype('uint8'))
+    return sr
+
 if __name__ == '__main__':
     ip = Image.open('images/demo.png')
-    sr = return_SRFlow_result(ip)
+    lr = PILToTensor()(ip).permute(1, 2, 0).numpy()
+    sr = test(lr)
+    # sr = return_SRFlow_result(ip)
     print(sr.size)

models/SRGAN/srgan.py

@@ -64,6 +64,15 @@ class GeneratorResnet(nn.Module):
             x = Image.fromarray((x.squeeze(0).permute(1, 2, 0).detach().numpy() * 255).astype('uint8'))
         return TF.adjust_brightness(x, 1.1)
     
+    def test(self, x):
+        """
+        x is a tensor
+        """
+        self.eval()
+        with torch.no_grad():
+            x = self.forward(x)
+        return x
+    
 if __name__ == '__main__':
     current_dir = os.path.dirname(os.path.realpath(__file__))
 
@@ -72,4 +81,6 @@ if __name__ == '__main__':
     model.eval()
     with torch.no_grad():
         input_image = Image.open('images/demo.png')
-        output_image = model.inference(input_image)
+        input_image = ToTensor()(input_image).unsqueeze(0)
+        output_image = model.test(input_image)
+        print(output_image.max())