Update inference_diffir.py

IFF-0303 · Oct 28, 2023 · 7d3c3cc · 7d3c3cc
1 parent 47653c2
commit 7d3c3cc
Showing 1 changed file with 15 additions and 15 deletions.
diff --git a/DiffIR-RealSR/inference_diffir.py b/DiffIR-RealSR/inference_diffir.py
@@ -51,18 +51,18 @@ def pad_test(lq,scale):
     im_list.sort()
     im_list = [name for name in im_list if name.endswith('.png')]
 
-    with torch.no_grad():
-        for name in im_list:
-            path = os.path.join(args.im_path, name)
-            im = cv2.imread(path)
-            im = img2tensor(im)
-            im = im.unsqueeze(0).cuda(0)/255.
-            lq,mod_pad_h,mod_pad_w= pad_test(im,args.scale)
-            with torch.no_grad():
-                sr = model(lq)
-            _, _, h, w = sr.size()
-            sr = sr[:, :, 0:h - mod_pad_h * args.scale, 0:w - mod_pad_w * args.scale]
-            im_sr = tensor2img(sr, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1))
-            save_path = os.path.join(args.res_path, name.split('.')[0]+'_out.png')
-            cv2.imwrite(save_path, im_sr)
-            print(save_path)
+
+    for name in im_list:
+        path = os.path.join(args.im_path, name)
+        im = cv2.imread(path)
+        im = img2tensor(im)
+        im = im.unsqueeze(0).cuda(0)/255.
+        lq,mod_pad_h,mod_pad_w= pad_test(im,args.scale)
+        with torch.no_grad():
+            sr = model(lq)
+        _, _, h, w = sr.size()
+        sr = sr[:, :, 0:h - mod_pad_h * args.scale, 0:w - mod_pad_w * args.scale]
+        im_sr = tensor2img(sr, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1))
+        save_path = os.path.join(args.res_path, name.split('.')[0]+'_out.png')
+        cv2.imwrite(save_path, im_sr)
+        print(save_path)