Spaces:
Running
Running
Update main_test_SRMNet.py
Browse files- 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 |
|