j

models/SRFlow/srflow.py

@@ -41,8 +41,6 @@ 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]
 
-    if not divide:
-            sr /= 255
 
     sr = Image.fromarray((sr).astype('uint8'))
     return sr
@@ -64,10 +62,14 @@ def return_SRFlow_result_from_tensor(lr_tensor, divide=True):
         lr_image = ToPILImage()(lr_tensor[b])
         sr_image = return_SRFlow_result(lr_image, divide)
         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 /= 255
+        
     return sr_tensor
 
 if __name__ == '__main__':