Update app.py

app.py

@@ -129,13 +129,46 @@ def get_reconstructed_scene(outdir, model, device):
     torch.cuda.empty_cache()
     return splatfile, video, depth_colored
 
+def extract_images(input_video, session_id):
+
+    start_time = time.time()
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    base_dir = os.path.join(os.environ["ANYSPLAT_PROCESSED"], session_id)
+    target_dir = base_dir
+    target_dir_images = os.path.join(target_dir, "images")
+
+    if os.path.exists(target_dir):
+        shutil.rmtree(target_dir)
+        
+    os.makedirs(target_dir)
+    os.makedirs(target_dir_images)
+
+    image_paths = []
+
+    if input_images is not None:
+        for file_data in input_images:
+            if isinstance(file_data, dict) and "name" in file_data:
+                file_path = file_data["name"]
+            else:
+                file_path = file_data
+            dst_path = os.path.join(target_dir_images, os.path.basename(file_path))
+            shutil.copy(file_path, dst_path)
+            image_paths.append(dst_path)
+
+    # Sort final images for gallery
+    image_paths = sorted(image_paths)
+
+    end_time = time.time()
+    print(
+        f"Files copied to {target_dir_images}; took {end_time - start_time:.3f} seconds"
+    )
+    return target_dir, image_paths
+    
 
-# 2) Handle uploaded video/images --> produce target_dir + images
 def extract_frames(input_video, session_id):
-    """
-    Create a new 'target_dir' + 'images' subfolder, and place user-uploaded
-    images or extracted frames from video into it. Return (target_dir, image_paths).
-    """
+
     start_time = time.time()
     gc.collect()
     torch.cuda.empty_cache()
@@ -144,9 +177,9 @@ def extract_frames(input_video, session_id):
     target_dir = base_dir
     target_dir_images = os.path.join(target_dir, "images")
 
-    # Clean up if somehow that folder already exists
     if os.path.exists(target_dir):
         shutil.rmtree(target_dir)
+        
     os.makedirs(target_dir)
     os.makedirs(target_dir_images)
 
@@ -187,18 +220,22 @@ def extract_frames(input_video, session_id):
     return target_dir, image_paths
 
 
-def update_gallery_on_upload(input_video, session_id):
-    """
-    Whenever user uploads or changes files, immediately handle them
-    and show in the gallery. Return (target_dir, image_paths).
-    If nothing is uploaded, returns "None" and empty list.
-    """
-    if not input_video and not input_images:
+def update_gallery_on_video_upload(input_video, session_id):
+
+    if not input_video:
         return None, None, None
         
     target_dir, image_paths = extract_frames(input_video, session_id)
     return None, target_dir, image_paths
 
+def update_gallery_on_images_upload(input_images, session_id):
+
+    if not input_images:
+        return None, None, None
+        
+    target_dir, image_paths = extract_images(input_images, session_id)
+    return None, target_dir, image_paths
+
 @spaces.GPU()
 def generate_splats_from_video(video_path, session_id=None):
     
@@ -239,35 +276,14 @@ def generate_splats_from_images(images_folder, session_id=None):
     return plyfile, video, depth_colored
 
 def cleanup(request: gr.Request):
-    """
-    Clean up session-specific directories and temporary files when the user session ends.
-    
-    This function is triggered when the Gradio demo is unloaded (e.g., when the user
-    closes the browser tab or navigates away). It removes all temporary files and
-    directories created during the user's session to free up storage space.
-    
-    Args:
-        request (gr.Request): Gradio request object containing session information
-    """
+
     sid = request.session_hash
     if sid:
         d1 = os.path.join(os.environ["ANYSPLAT_PROCESSED"], sid)
         shutil.rmtree(d1, ignore_errors=True)
         
 def start_session(request: gr.Request):
-    """
-    Initialize a new user session and return the session identifier.
-    
-    This function is triggered when the Gradio demo loads and creates a unique
-    session hash that will be used to organize outputs and temporary files
-    for this specific user session.
-    
-    Args:
-        request (gr.Request): Gradio request object containing session information
-        
-    Returns:
-        str: Unique session hash identifier
-    """
+
     return request.session_hash
 
 
@@ -322,7 +338,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(label="Upload Files", 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"
@@ -397,11 +413,17 @@ if __name__ == "__main__":
             outputs=[reconstruction_output, rgb_video, depth_video])
 
         input_video.upload(
-            fn=update_gallery_on_upload,
+            fn=update_gallery_on_video_upload,
             inputs=[input_video, session_state],
             outputs=[reconstruction_output, target_dir_output, image_gallery],
         )
 
+        input_images.upload(
+            fn=update_gallery_on_images_upload,
+            inputs=[input_images, session_state],
+            outputs=[reconstruction_output, target_dir_output, image_gallery],
+        )
+
         demo.unload(cleanup)
         demo.queue()
         demo.launch(show_error=True, share=True)