fix1830

models/SRFlow/srflow.py

@@ -41,9 +41,8 @@ def return_SRFlow_result(lr, divide, conf_path='models/SRFlow/code/confs/SRFlow_
     sr = rgb(torch.clamp(sr_t, 0, 1))
     sr = sr[:h * scale, :w * scale]
 
-
-    sr = Image.fromarray((sr).astype('uint8'))
-    return sr
+    sr_img = Image.fromarray((sr).astype('uint8'))
+    return sr_img
 
 def return_SRFlow_result_from_tensor(lr_tensor, divide=True):
     """
@@ -68,8 +67,8 @@ def return_SRFlow_result_from_tensor(lr_tensor, divide=True):
     sr_tensor = torch.cat(sr_list, dim=0)
 
     if not divide:
-        sr_tensor /= 255
-  
+        sr_tensor = sr_tensor / 255
+        
     return sr_tensor 
 
 if __name__ == '__main__':
@@ -77,10 +76,11 @@ if __name__ == '__main__':
     
     lr_tensor = PILToTensor()(lr).unsqueeze(0)
 
-    sr = return_SRFlow_result_from_tensor(lr_tensor)
-    print(sr.shape)
+    random_tensor = torch.randn(8, 3, 64, 64)
+    sr = return_SRFlow_result_from_tensor(random_tensor, divide=False)
+    print(sr)
 
-    # Show SR image of the first one in the batch
-    plt.imshow(np.transpose(sr[0].cpu().detach().numpy(), (1, 2, 0)))
-    # plt.axis('off')
-    plt.show()
+    # # Show SR image of the first one in the batch
+    # plt.imshow(np.transpose(sr[0].cpu().detach().numpy(), (1, 2, 0)))
+    # # plt.axis('off')
+    # plt.show()