fix

app.py

@@ -46,7 +46,7 @@ def generate_vstack_from_image(image: Image.Image, input_focal_position: int, nu
     pipe.to(args.device)
     batch = convert_to_batch(image, input_focal_position=input_focal_position)
     output_frames, focal_stack_num = inference_on_image(args, batch, pipe, device)
-    save_dir = os.path.join(OUTPUT_DIR, batch['name'])
+    save_dir = os.path.join(OUTPUT_DIR)
 
     write_output(save_dir, output_frames, focal_stack_num, batch['icc_profile'])
 

simple_inference.py

@@ -186,7 +186,7 @@ def write_output(output_dir, frames, focal_stack_num, icc_profile):
     os.makedirs(output_dir, exist_ok=True)
 
     print("Frames shape: ", frames.shape)
-    export_to_video(frames.permute(0,2,3,1).cpu().numpy(), os.path.join(output_dir, "stack.mp4"), fps=20)
+    export_to_video(frames.permute(0,2,3,1).cpu().numpy(), os.path.join(output_dir, "stack.mp4"), fps=5)
 
     #save images
     for i in range(9):
@@ -243,7 +243,8 @@ def main():
         img = Image.open(args.image_path)
         batch = convert_to_batch(img, input_focal_position=6)
         output_frames, focal_stack_num = inference_on_image(args, batch, pipeline, device)
-        val_save_dir = os.path.join(args.output_dir, "validation_images", batch['name'])
+        name = os.path.splitext(os.path.basename(args.image_path))[0]
+        val_save_dir = os.path.join(args.output_dir, "validation_images", name)
         write_output(val_save_dir, output_frames, focal_stack_num, batch['icc_profile'])