Update main_test_CMFNet.py

main_test_CMFNet.py

@@ -55,14 +55,14 @@ def main():
         input_ = F.pad(input_, (0, padw, 0, padh), 'reflect')
         with torch.no_grad():
             restored = model(input_)
-            restored_ = restored[0]
-            restored_= torch.clamp(restored_, 0, 1)
-            restored_= restored_[:, :, :h, :w]
-            restored_= restored_.permute(0, 2, 3, 1).cpu().detach().numpy()
-            restored_= img_as_ubyte(restored_[0])
+            
+        restored = torch.clamp(restored[1], 0, 1)
+        restored = restored[:, :, :h, :w]
+        restored = restored.permute(0, 2, 3, 1).cpu().detach().numpy()
+        restored = img_as_ubyte(restored[0])
 
         f = os.path.splitext(os.path.split(file_)[-1])[0]
-        save_img((os.path.join(out_dir, 'out.png')), restored_)
+        save_img((os.path.join(out_dir, f + '.png')), restored_)
 
 
 
@@ -82,11 +82,6 @@ def load_checkpoint(model, weights):
             new_state_dict[name] = v
         model.load_state_dict(new_state_dict)
 
-def setup(args):
-    save_dir = 'result/'
-    folder = 'test/'
-
-    return folder, save_dir
 
 if __name__ == '__main__':
     main()