Merge pull request #3 from Vembala/master

Master
Vembala · Jun 3, 2021 · 7cbf635 · 7cbf635
2 parents d94ca31 + f7d455e
commit 7cbf635
Showing 1 changed file with 3 additions and 2 deletions.
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
 
@@ -26,7 +27,7 @@
 
 i = 0
 
-filename = ""
+fileName = ""
 
 print('Input file:',fileName)
 mat_img = cv2.imread(fileName)
@@ -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