Refactor function to improve performance and readability

models/SRFlow/srflow.py

@@ -55,7 +55,8 @@ def test_srflow(lr, conf_path='models/SRFlow/code/confs/SRFlow_DF2K_4X.yml', hea
 
     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))
@@ -67,12 +68,14 @@ def test_srflow(lr, conf_path='models/SRFlow/code/confs/SRFlow_DF2K_4X.yml', hea
 
     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
 
 if __name__ == '__main__':
     ip = Image.open('images/demo.png')
     lr = PILToTensor()(ip).permute(1, 2, 0).numpy()
-    sr = test(lr)
+    print(lr.shape)
+    sr = test_srflow(lr)
     # sr = return_SRFlow_result(ip)
     print(sr.size)

models/SRGAN/m.py

@@ -0,0 +1,9 @@
+from PIL import Image
+from torchvision.transforms import PILToTensor
+ip = Image.open('images/demo.png')
+lr = PILToTensor()(ip).unsqueeze(0).permute(0, 2, 3, 1)
+print(lr.shape)
+lr = lr.squeeze(0).permute(0, 1, 2).numpy()
+# lr = PILToTensor()(ip).permute(1, 2, 0)
+# lr = lr.unsqueeze(0).permute(0, 3, 1, 2).numpy()
+print(lr.shape)