19

models/SRFlow/srflow.py

@@ -26,6 +26,7 @@ def return_SRFlow_result(lr, divide=True, conf_path='models/SRFlow/code/confs/SR
     if not divide:
         lr = np.round(255 * lr).astype(np.uint8) 
 
+
     scale = opt['scale']
     pad_factor = 2
 
@@ -67,20 +68,20 @@ def return_SRFlow_result_from_tensor(lr_tensor, divide):
     sr_tensor = torch.cat(sr_list, dim=0)
 
     if not divide:
-        sr_tensor = sr_tensor / 255
+        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)
-
+    lr_tensor = PILToTensor()(lr).unsqueeze(0) / 255
+    print(lr_tensor.shape)
     random_tensor = torch.randn(8, 3, 64, 64)
-    sr = return_SRFlow_result_from_tensor(random_tensor, divide=False)
+    sr = return_SRFlow_result_from_tensor(lr_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()