Output path

inference_img.py

@@ -5,6 +5,8 @@ import argparse
 from torch.nn import functional as F
 import warnings
 
+OUTPUT_PATH = "output/"
+
 warnings.filterwarnings("ignore")
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -92,25 +94,27 @@ else:
         tmp.append(img1)
         img_list = tmp
 
-if not os.path.exists('output'):
-    os.mkdir('output')
+if not os.path.exists('/output'):
+    os.mkdir('/output')
 
 print("Checkpoint reached! output folder ok")
 
 for i in range(len(img_list)):
+    filename_exr = os.path.join(OUTPUT_PATH, f"img{i}.exr")
+    filename_png = os.path.join(OUTPUT_PATH, f"img{i}.png")
+
     if args.img[0].endswith('.exr') and args.img[1].endswith('.exr'):
-        cv2.imwrite('output/img{}.exr'.format(i), (img_list[i][0]).cpu().numpy().transpose(1, 2, 0)[:h, :w], [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_HALF])
-        save_path = 'output/img{}.png'.format(i)
-        success = cv2.imwrite(save_path, (img_list[i][0] * 255).byte().cpu().numpy().transpose(1, 2, 0)[:h, :w])
-        print(f"Saving to {save_path} → success: {success}")
-        print("Saving to:", os.path.abspath(save_path))
+        cv2.imwrite(filename_exr, (img_list[i][0]).cpu().numpy().transpose(1, 2, 0)[:h, :w], [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_HALF])
+        
+        success = cv2.imwrite(filename_png, (img_list[i][0] * 255).byte().cpu().numpy().transpose(1, 2, 0)[:h, :w])
+        print(f"Saving to {filename_png} → success: {success}")
+        print("Saving to:", os.path.abspath(filename_png))
 
     else:
-        cv2.imwrite('output/img{}.png'.format(i), (img_list[i][0] * 255).byte().cpu().numpy().transpose(1, 2, 0)[:h, :w])
-        save_path = 'output/img{}.png'.format(i)
-        success = cv2.imwrite(save_path, (img_list[i][0] * 255).byte().cpu().numpy().transpose(1, 2, 0)[:h, :w])
-        print(f"Saving to {save_path} → success: {success}")
-        print("Saving to:", os.path.abspath(save_path))
+        success = cv2.imwrite(filename_png, (img_list[i][0] * 255).byte().cpu().numpy().transpose(1, 2, 0)[:h, :w])
+        print(f"Saving to {filename_png} → success: {success}")
+        print("Saving to:", os.path.abspath(filename_png))
+
 
 
 print("Checkpoint reached!")