fix

app.py

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

simple_inference.py

@@ -124,29 +124,26 @@ def find_scale(height, width):
         # Reduce the scale slightly
         scale -= 0.01
 
-def convert_to_batch(image, input_focal_position, sample_frames=9):
-    scene, focal_stack_num = image, input_focal_position
-    from PIL import Image
-    with Image.open(scene) as img:
-
-        icc_profile = img.info.get("icc_profile")
-        if icc_profile is None:
-            icc_profile = "none"
-        original_pixels = torch.from_numpy(np.array(img)).float().permute(2,0,1)
-        original_pixels = original_pixels / 255
-        width, height = img.size
-        scaled_width, scaled_height = find_scale(width, height)
-
-        img_resized = img.resize((scaled_width, scaled_height))
-        img_tensor = torch.from_numpy(np.array(img_resized)).float()
-        img_normalized = img_tensor / 127.5 - 1
-        img_normalized = img_normalized.permute(2, 0, 1)
-
-        pixels = torch.zeros((1, sample_frames, 3, scaled_height, scaled_width))
-        pixels[0, focal_stack_num] = img_normalized
-        
-        name = os.path.splitext(os.path.basename(scene))[0]
-        return {"pixel_values": pixels, "focal_stack_num": focal_stack_num, "original_pixel_values": original_pixels, 'icc_profile': icc_profile, "name": name}
+
+def convert_to_batch(img, input_focal_position, sample_frames=9):
+    focal_stack_num = input_focal_position
+    icc_profile = img.info.get("icc_profile")
+    if icc_profile is None:
+        icc_profile = "none"
+    original_pixels = torch.from_numpy(np.array(img)).float().permute(2,0,1)
+    original_pixels = original_pixels / 255
+    width, height = img.size
+    scaled_width, scaled_height = find_scale(width, height)
+
+    img_resized = img.resize((scaled_width, scaled_height))
+    img_tensor = torch.from_numpy(np.array(img_resized)).float()
+    img_normalized = img_tensor / 127.5 - 1
+    img_normalized = img_normalized.permute(2, 0, 1)
+
+    pixels = torch.zeros((1, sample_frames, 3, scaled_height, scaled_width))
+    pixels[0, focal_stack_num] = img_normalized
+    
+    return {"pixel_values": pixels, "focal_stack_num": focal_stack_num, "original_pixel_values": original_pixels, 'icc_profile': icc_profile}
 
 
 def inference_on_image(args, batch, pipeline, device):
@@ -240,12 +237,11 @@ def main():
     if args.output_dir is not None:
         os.makedirs(args.output_dir, exist_ok=True)
 
-
     pipeline, device = load_model(args)
 
-    batch = convert_to_batch(args.image_path, input_focal_position=6)
-
     with torch.no_grad():
+        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'])
         write_output(val_save_dir, output_frames, focal_stack_num, batch['icc_profile'])