Update app.py

app.py

@@ -82,14 +82,8 @@ def save_splat_file(splat_data, output_path):
     with open(output_path, "wb") as f:
         f.write(splat_data)
 
-def get_reconstructed_scene(outdir, model, device):
+def get_reconstructed_scene(outdir, image_files, model, device):
 
-    image_files = sorted(
-        [
-            os.path.join(outdir, "images", f)
-            for f in os.listdir(os.path.join(outdir, "images"))
-        ]
-    )
     images = [process_image(img_path) for img_path in image_files]
     images = torch.stack(images, dim=0).unsqueeze(0).to(device)  # [1, K, 3, 448, 448]
     b, v, c, h, w = images.shape
@@ -245,7 +239,7 @@ def generate_splats_from_video(video_path, session_id=None):
     return plyfile, rgb_vid, depth_vid, image_paths
     
 @spaces.GPU()
-def generate_splats_from_images(images_folder, session_id=None):
+def generate_splats_from_images(image_paths, session_id=None):
 
     if session_id is None:
         session_id = uuid.uuid4().hex
@@ -256,16 +250,11 @@ def generate_splats_from_images(images_folder, session_id=None):
 
     base_dir = os.path.join(os.environ["ANYSPLAT_PROCESSED"], session_id)
 
-    all_files = (
-        sorted(os.listdir(images_folder))
-        if os.path.isdir(images_folder)
-        else []
-    )
-    all_files = [f"{i}: {filename}" for i, filename in enumerate(all_files)]
+    all_files = [os.path.basename(p) for p in image_paths]
 
     print("Running run_model...")
     with torch.no_grad():
-        plyfile, video, depth_colored = get_reconstructed_scene(base_dir, model, device)
+        plyfile, video, depth_colored = get_reconstructed_scene(base_dir, all_files, model, device)
 
     end_time = time.time()
     print(f"Total time: {end_time - start_time:.2f} seconds (including IO)")
@@ -413,7 +402,7 @@ if __name__ == "__main__":
                
         submit_btn.click(
             fn=generate_splats_from_images,
-            inputs=[target_dir_output, session_state],
+            inputs=[image_gallery, session_state],
             outputs=[reconstruction_output, rgb_video, depth_video])
 
         input_video.upload(