52Hz commited on
Commit
a38e5c8
·
1 Parent(s): 4edb95b

Update main_test_SRMNet.py

Browse files
Files changed (1) hide show
  1. main_test_SRMNet.py +3 -2
main_test_SRMNet.py CHANGED
@@ -31,7 +31,8 @@ def main():
31
 
32
  # Prepare testing data
33
  files = natsorted(glob(os.path.join(args.input_dir, '*.JPG')) + glob(os.path.join(args.input_dir, '*.PNG')))
34
-
 
35
  os.makedirs(args.result_dir, exist_ok=True)
36
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
37
  # Build model
@@ -63,7 +64,7 @@ def main():
63
 
64
  f = os.path.splitext(os.path.split(file_)[-1])[0]
65
  save_img((os.path.join(args.result_dir, f + '.png')), restored)
66
- print(os.path.join(args.result_dir, f + '.png'))
67
  print('finish !')
68
 
69
 
 
31
 
32
  # Prepare testing data
33
  files = natsorted(glob(os.path.join(args.input_dir, '*.JPG')) + glob(os.path.join(args.input_dir, '*.PNG')))
34
+ if len(files) == 0:
35
+ raise Exception(f"No files found at {args.inp_dir}")
36
  os.makedirs(args.result_dir, exist_ok=True)
37
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
38
  # Build model
 
64
 
65
  f = os.path.splitext(os.path.split(file_)[-1])[0]
66
  save_img((os.path.join(args.result_dir, f + '.png')), restored)
67
+ print('{}'.format(os.path.join(args.result_dir, f + '.png')))
68
  print('finish !')
69
 
70