Update app.py

app.py

@@ -82,8 +82,14 @@ def save_splat_file(splat_data, output_path):
     with open(output_path, "wb") as f:
         f.write(splat_data)
 
-def get_reconstructed_scene(outdir, image_files, model, device):
+def get_reconstructed_scene(outdir, 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
@@ -234,12 +240,12 @@ def generate_splats_from_video(video_path, session_id=None):
         session_id = uuid.uuid4().hex
 
     images_folder, image_paths = extract_frames(video_path, session_id)
-    plyfile, rgb_vid, depth_vid = generate_splats_from_images(image_paths, session_id)
+    plyfile, rgb_vid, depth_vid = generate_splats_from_images(images_folder, session_id)
 
     return plyfile, rgb_vid, depth_vid, image_paths
     
 @spaces.GPU()
-def generate_splats_from_images(image_paths, session_id=None):
+def generate_splats_from_images(images_folder, session_id=None):
 
     if session_id is None:
         session_id = uuid.uuid4().hex
@@ -250,9 +256,16 @@ def generate_splats_from_images(image_paths, 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)]
+
     print("Running run_model...")
     with torch.no_grad():
-        plyfile, video, depth_colored = get_reconstructed_scene(base_dir, image_paths, model, device)
+        plyfile, video, depth_colored = get_reconstructed_scene(base_dir, model, device)
 
     end_time = time.time()
     print(f"Total time: {end_time - start_time:.2f} seconds (including IO)")
@@ -322,7 +335,7 @@ if __name__ == "__main__":
                     with gr.Tab("Video"):
                         input_video = gr.Video(label="Upload Video", sources=["upload"], interactive=True, height=512)
                     with gr.Tab("Images"):
-                        input_images = gr.File(file_count="multiple", label="Upload Files", type="filepath", height=512)
+                        input_images = gr.File(file_count="multiple", label="Upload Files", height=512)
                         
                     submit_btn = gr.Button(
                         "Generate Gaussian Splat", scale=1, variant="primary"
@@ -400,7 +413,7 @@ if __name__ == "__main__":
                
         submit_btn.click(
             fn=generate_splats_from_images,
-            inputs=[image_gallery, session_state],
+            inputs=[target_dir_output, session_state],
             outputs=[reconstruction_output, rgb_video, depth_video])
 
         input_video.upload(
@@ -417,4 +430,4 @@ if __name__ == "__main__":
 
         demo.unload(cleanup)
         demo.queue()
-        demo.launch(show_error=True, share=True)
+        demo.launch(show_error=True, share=True)