Change path logic

Vembala · Jun 3, 2021 · f7d455e · f7d455e
1 parent 572d482
commit f7d455e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/run.py b/run.py
@@ -6,6 +6,7 @@
 import glob
 from tqdm import tqdm
 import jittor as jt
+from pathlib import Path
 
 jt.flags.use_cuda = 1
 
@@ -43,7 +44,7 @@
 
 combine_model.predict_shadow(mat_img)
 
-output_file = 'ori'+ filename[-3:] +'.jpg'
+output_file = 'ori'+ Path(fileName).stem +'.jpg'
 print('Output file:',output_file)
 cv2.imwrite(output_file,cv2.cvtColor(combine_model.generated, cv2.COLOR_BGR2RGB))
 i = i + 1