Update app.py

app.py

@@ -13,6 +13,7 @@ import tempfile
 import time
 from datetime import datetime
 from pathlib import Path
+import uuid
 
 import cv2
 import gradio as gr
@@ -30,9 +31,9 @@ from src.utils.image import process_image
 os.environ["ANYSPLAT_PROCESSED"] = f"{os.getcwd()}/proprocess_results"
 
 
-# 1) Core model inference
+
 def get_reconstructed_scene(outdir, model, device):
-    # Load Images
+
     image_files = sorted(
         [
             os.path.join(outdir, "images", f)
@@ -45,10 +46,9 @@ def get_reconstructed_scene(outdir, model, device):
 
     assert c == 3, "Images must have 3 channels"
 
-    # Run Inference
+
     gaussians, pred_context_pose = model.inference((images + 1) * 0.5)
 
-    # Save the results
     pred_all_extrinsic = pred_context_pose["extrinsic"]
     pred_all_intrinsic = pred_context_pose["intrinsic"]
     video, depth_colored = save_interpolated_video(
@@ -79,7 +79,7 @@ def get_reconstructed_scene(outdir, model, device):
 
 
 # 2) Handle uploaded video/images --> produce target_dir + images
-def handle_uploads(input_video, input_images, session_id):
+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).
@@ -100,18 +100,6 @@ def handle_uploads(input_video, input_images, session_id):
 
     image_paths = []
 
-    # --- Handle images ---
-    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)
-
-    # --- Handle video ---
     if input_video is not None:
         if isinstance(input_video, dict) and "name" in input_video:
             video_path = input_video["name"]
@@ -147,8 +135,7 @@ def handle_uploads(input_video, input_images, session_id):
     return target_dir, image_paths
 
 
-# 3) Update gallery on upload
-def update_gallery_on_upload(input_video, input_images, session_id):
+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).
@@ -157,13 +144,25 @@ def update_gallery_on_upload(input_video, input_images, session_id):
     if not input_video and not input_images:
         return None, None, None
         
-    target_dir, image_paths = handle_uploads(input_video, input_images, session_id)
+    target_dir, image_paths = extract_frames(input_video, session_id)
     return None, target_dir, image_paths
 
+@spaces.GPU()
+def generate_splat_from_video(video_path, session_id=None):
+    
+    if session_id is None:
+        session_id = uuid.uuid4().hex
 
+    target_dir, image_paths = extract_frames(input_video, session_id)
+    
+    return generate_splat_from_images(images_folder, session_id), image_paths
+    
 @spaces.GPU()
-def generate_splat(images_folder, session_id=None):
+def generate_splat_from_images(images_folder, session_id=None):
 
+    if session_id is None:
+        session_id = uuid.uuid4().hex
+    
     start_time = time.time()
     gc.collect()
     torch.cuda.empty_cache()
@@ -186,7 +185,22 @@ def generate_splat(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.
@@ -227,43 +241,6 @@ if __name__ == "__main__":
             margin: 0 auto;
             max-width: 1024px;
         }
-
-        .custom-log * {
-            font-style: italic;
-            font-size: 22px !important;
-            background-image: linear-gradient(120deg, #0ea5e9 0%, #6ee7b7 60%, #34d399 100%);
-            -webkit-background-clip: text;
-            background-clip: text;
-            font-weight: bold !important;
-            color: transparent !important;
-            text-align: center !important;
-        }
-        
-        .example-log * {
-            font-style: italic;
-            font-size: 16px !important;
-            background-image: linear-gradient(120deg, #0ea5e9 0%, #6ee7b7 60%, #34d399 100%);
-            -webkit-background-clip: text;
-            background-clip: text;
-            color: transparent !important;
-        }
-        
-        #my_radio .wrap {
-            display: flex;
-            flex-wrap: nowrap;
-            justify-content: center;
-            align-items: center;
-        }
-
-        #my_radio .wrap label {
-            display: flex;
-            width: 50%;
-            justify-content: center;
-            align-items: center;
-            margin: 0;
-            padding: 10px 0;
-            box-sizing: border-box;
-        }
         """
     with gr.Blocks(css=css, title="AnySplat Demo", theme=theme) as demo:
         session_state = gr.State()
@@ -287,12 +264,7 @@ if __name__ == "__main__":
             with gr.Row():
                 with gr.Column():
                     input_video = gr.Video(label="Upload Video", interactive=True, height=512)
-                    input_images = gr.File(
-                        file_count="multiple",
-                        label="Upload Images",
-                        interactive=True,
-                        visible=False
-                    )
+
                     submit_btn = gr.Button(
                         "Reconstruct", scale=1, variant="primary"
                     )
@@ -306,7 +278,6 @@ if __name__ == "__main__":
                         preview=True,
                     )
     
-    
                 with gr.Column():
                     with gr.Column():
                         reconstruction_output = gr.Model3D(
@@ -373,18 +344,16 @@ if __name__ == "__main__":
             # )
             
         submit_btn.click(
-            fn=generate_splat,
+            fn=generate_splat_from_images,
             inputs=[target_dir_output, session_state],
             outputs=[reconstruction_output, rgb_video, depth_video])
 
         input_video.change(
             fn=update_gallery_on_upload,
-            inputs=[input_video, input_images, session_state],
+            inputs=[input_video, session_state],
             outputs=[reconstruction_output, target_dir_output, image_gallery],
         )
-        input_images.change(
-            fn=update_gallery_on_upload,
-            inputs=[input_video, input_images, session_state],
-            outputs=[reconstruction_output, target_dir_output, image_gallery],
-        )
-        demo.queue().launch(show_error=True, share=True)
+
+        demo.unload(cleanup)
+        demo.queue()
+        demo.launch(show_error=True, share=True)