add new func

models/SRFlow/srflow.py

@@ -4,9 +4,10 @@ import sys
 sys.path.append('models')
 from SRFlow.code import imread, impad, load_model, t, rgb
 from PIL import Image
-from torchvision.transforms import PILToTensor
+import matplotlib.pyplot as plt
+from torchvision.transforms import PILToTensor, ToPILImage
 
-def return_SRFlow_result(lr, conf_path='models/SRFlow/code/confs/SRFlow_DF2K_4X.yml', heat=0.7):
+def return_SRFlow_result(lr, conf_path='models/SRFlow/code/confs/SRFlow_DF2K_4X.yml'):
     """
     Apply Super-Resolution using SRFlow model to the input LR (low-resolution) image.
 
@@ -39,43 +40,37 @@ def return_SRFlow_result(lr, conf_path='models/SRFlow/code/confs/SRFlow_DF2K_4X.
     sr = Image.fromarray((sr).astype('uint8'))
     return sr
 
-def test_srflow(lr, conf_path='models/SRFlow/code/confs/SRFlow_DF2K_4X.yml', heat=0.7):
+def return_SRFlow_result_from_tensor(lr_tensor):
     """
-    Apply Super-Resolution using SRFlow model to the input LR (low-resolution) image.
+    Apply Super-Resolution using SRFlow model to the input batched BCHW tensor.
 
     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.
+    - lr_tensor: Batched BCHW tensor
 
     Returns:
-    - sr: tensor
+    - sr_tensor: Processed batched BCHW tensor
     """
-    model, opt = load_model(conf_path)
+    batch_size = lr_tensor.shape[0]
+    sr_list = []
 
-    scale = opt['scale']
-    pad_factor = 2
-    
-    lr = lr.squeeze(0).permute(0, 1, 2).numpy()
-    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)
+    for b in range(batch_size):
+        lr_image = ToPILImage()(lr_tensor[b])
+        sr_image = return_SRFlow_result(lr_image)
+        sr_tensor = PILToTensor()(sr_image).unsqueeze(0)
+        sr_list.append(sr_tensor)
 
-    sr = rgb(torch.clamp(sr_t, 0, 1))
-    sr = sr[:h * scale, :w * scale]
-    sr = sr.unsqueeze(0).permute(0, 3, 1, 2)
-    
-    return sr
+    sr_tensor = torch.cat(sr_list, dim=0)
+    return sr_tensor
 
 if __name__ == '__main__':
-    ip = Image.open('images/demo.png')
-    lr = PILToTensor()(ip).permute(1, 2, 0).numpy()
-    print(lr.shape)
-    sr = test_srflow(lr)
-    # sr = return_SRFlow_result(ip)
-    print(sr.size)
+    lr = Image.open('images/demo.png')
+    
+    lr_tensor = PILToTensor()(lr).unsqueeze(0)
+
+    sr = return_SRFlow_result_from_tensor(lr_tensor)
+    print(sr.shape)
+    
+    plt.imshow(np.transpose(sr[0].cpu().detach().numpy(), (1, 2, 0)))
+    plt.axis('off')
+    plt.title('Super-Resolved Image')
+    plt.show()