7

models/SRFlow/srflow.py

@@ -7,7 +7,7 @@ from PIL import Image
 import matplotlib.pyplot as plt
 from torchvision.transforms import PILToTensor, ToPILImage
 
-def return_SRFlow_result(lr, divide=True, conf_path='models/SRFlow/code/confs/SRFlow_DF2K_4X.yml'):
+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.
 
@@ -23,10 +23,6 @@ def return_SRFlow_result(lr, divide=True, conf_path='models/SRFlow/code/confs/SR
 
     lr = PILToTensor()(lr).permute(1, 2, 0).numpy()
 
-    if not divide:
-        lr = np.round(255 * lr).astype(np.uint8) 
-
-
     scale = opt['scale']
     pad_factor = 2
 
@@ -45,7 +41,7 @@ def return_SRFlow_result(lr, divide=True, conf_path='models/SRFlow/code/confs/SR
     sr_img = Image.fromarray((sr).astype('uint8'))
     return sr_img
 
-def return_SRFlow_result_from_tensor(lr_tensor, divide):
+def return_SRFlow_result_from_tensor(lr_tensor):
     """
     Apply Super-Resolution using SRFlow model to the input batched BCHW tensor.
 
@@ -60,25 +56,22 @@ def return_SRFlow_result_from_tensor(lr_tensor, divide):
 
     for b in range(batch_size):
         lr_image = ToPILImage()(lr_tensor[b])
-        sr_image = return_SRFlow_result(lr_image, divide)
+        sr_image = return_SRFlow_result(lr_image)
         sr_tensor = PILToTensor()(sr_image).unsqueeze(0)
         
         sr_list.append(sr_tensor)
 
     sr_tensor = torch.cat(sr_list, dim=0)
 
-    if not divide:
-        sr_tensor = torch.ones_like(sr_tensor) - sr_tensor / 255
-
     return sr_tensor 
 
 if __name__ == '__main__':
     lr = Image.open('images/demo.png')
     
-    lr_tensor = PILToTensor()(lr).unsqueeze(0) / 255
+    lr_tensor = PILToTensor()(lr).unsqueeze(0) 
     print(lr_tensor.shape)
     random_tensor = torch.randn(8, 3, 64, 64)
-    sr = return_SRFlow_result_from_tensor(lr_tensor, divide=False)
+    sr = return_SRFlow_result_from_tensor(lr_tensor)
     print(sr)
 
     # Show SR image of the first one in the batch