load and save of ui settings

.gitignore

@@ -68,3 +68,4 @@ db.sqlite3-journal
 # Flask stuff:
 instance/
 .webassets-cache
+stream3r/__pycache__/stream_session.cpython-311.pyc

app.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
+import json
 import os
 import cv2
 import torch
@@ -15,6 +16,7 @@ import glob
 import gc
 import time
 import zipfile
+from typing import Any, Dict, Optional
 from stream3r.models.stream3r import STream3R
 from stream3r.stream_session import StreamSession
 from stream3r.models.components.utils.load_fn import load_and_preprocess_images
@@ -140,6 +142,33 @@ def _copy_with_unique_name(src_path: str, dst_dir: str) -> str:
     return dest_path
 
 
+def load_session_settings(target_dir: str) -> Dict[str, Any]:
+    settings_path = os.path.join(target_dir, "session_settings.json")
+    if not os.path.exists(settings_path):
+        return {}
+    try:
+        with open(settings_path, "r", encoding="utf-8") as handle:
+            data = json.load(handle)
+            if isinstance(data, dict):
+                return data
+    except (json.JSONDecodeError, OSError) as exc:
+        print(f"Failed to load session settings from {settings_path}: {exc}")
+    return {}
+
+
+def build_frame_filter_choices(target_dir_images: str) -> list[str]:
+    if not os.path.isdir(target_dir_images):
+        return ["All"]
+    files = sorted(os.listdir(target_dir_images))
+    return ["All"] + [f"{idx}: {name}" for idx, name in enumerate(files)]
+
+
+def sanitize_frame_filter_label(label: Optional[str]) -> str:
+    if not label:
+        return "All"
+    return label.replace('.', '_').replace(':', '').replace(' ', '_')
+
+
 # -------------------------------------------------------------------------
 # 1) Core model inference
 # -------------------------------------------------------------------------
@@ -341,7 +370,8 @@ def update_gallery_on_upload(input_video, input_images, input_zip, session_state
     Handle any new uploads (video, images, or zip) and render preview.
     """
     if not input_video and not input_images and not input_zip and not session_state:
-        return None, current_target_dir, None, None, None
+        default_updates = [gr.update()] * 9
+        return (None, current_target_dir, None, None, None, *default_updates)
 
     target_dir, image_paths, session_loaded = handle_uploads(
         input_video,
@@ -356,7 +386,90 @@ def update_gallery_on_upload(input_video, input_images, input_zip, session_state
     else:
         message = "Upload complete. Click 'Reconstruct' to begin 3D processing."
 
-    return None, target_dir, image_paths, message, None
+    target_dir_images = os.path.join(target_dir, "images")
+    frame_filter_choices = build_frame_filter_choices(target_dir_images)
+    frame_value = "All"
+    frame_update = gr.update(choices=frame_filter_choices, value=frame_value)
+
+    streaming_update = gr.update()
+    mode_update = gr.update()
+    conf_update = gr.update()
+    mask_black_update = gr.update()
+    mask_white_update = gr.update()
+    show_cam_update = gr.update()
+    mask_sky_update = gr.update()
+    prediction_mode_update = gr.update()
+
+    reconstruction_value = None
+
+    if session_loaded:
+        settings = load_session_settings(target_dir)
+
+        if settings:
+            if "frame_filter" in settings:
+                potential_value = settings.get("frame_filter", "All")
+                if potential_value in frame_filter_choices:
+                    frame_value = potential_value
+                frame_update = gr.update(choices=frame_filter_choices, value=frame_value)
+
+            if "streaming" in settings:
+                streaming_update = gr.update(value=bool(settings.get("streaming", True)))
+
+            if settings.get("mode") in {"causal", "window", "full"}:
+                mode_update = gr.update(value=settings["mode"])
+
+            if "conf_thres" in settings:
+                try:
+                    conf_update = gr.update(value=float(settings["conf_thres"]))
+                except (TypeError, ValueError):
+                    pass
+
+            if "mask_black_bg" in settings:
+                mask_black_update = gr.update(value=bool(settings.get("mask_black_bg", False)))
+            if "mask_white_bg" in settings:
+                mask_white_update = gr.update(value=bool(settings.get("mask_white_bg", False)))
+            if "show_cam" in settings:
+                show_cam_update = gr.update(value=bool(settings.get("show_cam", True)))
+            if "mask_sky" in settings:
+                mask_sky_update = gr.update(value=bool(settings.get("mask_sky", False)))
+
+            pred_mode_value = settings.get("prediction_mode")
+            if pred_mode_value in {"Depthmap and Camera Branch", "Pointmap Branch"}:
+                prediction_mode_update = gr.update(value=pred_mode_value)
+
+            try:
+                conf_val = settings["conf_thres"]
+                mode_val = settings["mode"]
+                pred_mode_val = settings["prediction_mode"]
+                mask_black_val = bool(settings.get("mask_black_bg", False))
+                mask_white_val = bool(settings.get("mask_white_bg", False))
+                show_cam_val = bool(settings.get("show_cam", True))
+                mask_sky_val = bool(settings.get("mask_sky", False))
+                glb_candidate = os.path.join(
+                    target_dir,
+                    f"glbscene_{conf_val}_{sanitize_frame_filter_label(frame_value)}_maskb{mask_black_val}_maskw{mask_white_val}_cam{show_cam_val}_sky{mask_sky_val}_pred{pred_mode_val.replace(' ', '_')}_mode{mode_val}.glb",
+                )
+                if os.path.exists(glb_candidate):
+                    reconstruction_value = glb_candidate
+            except (KeyError, AttributeError):
+                pass
+
+    return (
+        reconstruction_value,
+        target_dir,
+        image_paths,
+        message,
+        None,
+        streaming_update,
+        mode_update,
+        conf_update,
+        frame_update,
+        mask_black_update,
+        mask_white_update,
+        show_cam_update,
+        mask_sky_update,
+        prediction_mode_update,
+    )
 
 
 def update_gallery_without_session(input_video, input_images, input_zip, current_target_dir):
@@ -391,9 +504,7 @@ def gradio_demo(
 
     # Prepare frame_filter dropdown
     target_dir_images = os.path.join(target_dir, "images")
-    all_files = sorted(os.listdir(target_dir_images)) if os.path.isdir(target_dir_images) else []
-    all_files = [f"{i}: {filename}" for i, filename in enumerate(all_files)]
-    frame_filter_choices = ["All"] + all_files
+    frame_filter_choices = build_frame_filter_choices(target_dir_images)
 
     print("Running run_model...")
     with torch.no_grad():
@@ -403,6 +514,25 @@ def gradio_demo(
     prediction_save_path = os.path.join(target_dir, "predictions.npz")
     np.savez(prediction_save_path, **predictions)
 
+    frame_filter_value = frame_filter if frame_filter is not None else "All"
+
+    session_settings = {
+        "streaming": bool(streaming),
+        "mode": mode,
+        "conf_thres": float(conf_thres),
+        "frame_filter": frame_filter_value,
+        "mask_black_bg": bool(mask_black_bg),
+        "mask_white_bg": bool(mask_white_bg),
+        "show_cam": bool(show_cam),
+        "mask_sky": bool(mask_sky),
+        "prediction_mode": prediction_mode,
+    }
+    try:
+        with open(os.path.join(target_dir, "session_settings.json"), "w", encoding="utf-8") as handle:
+            json.dump(session_settings, handle, indent=2)
+    except OSError as exc:
+        print(f"Failed to write session settings: {exc}")
+
     session_state_file = None
     if streaming:
         if session_cache_path is None:
@@ -417,7 +547,7 @@ def gradio_demo(
     # Build a GLB file name
     glbfile = os.path.join(
         target_dir,
-        f"glbscene_{conf_thres}_{frame_filter.replace('.', '_').replace(':', '').replace(' ', '_')}_maskb{mask_black_bg}_maskw{mask_white_bg}_cam{show_cam}_sky{mask_sky}_pred{prediction_mode.replace(' ', '_')}_mode{mode}.glb",
+        f"glbscene_{conf_thres}_{sanitize_frame_filter_label(frame_filter)}_maskb{mask_black_bg}_maskw{mask_white_bg}_cam{show_cam}_sky{mask_sky}_pred{prediction_mode.replace(' ', '_')}_mode{mode}.glb",
     )
 
     # Convert predictions to GLB
@@ -441,7 +571,8 @@ def gradio_demo(
 
     end_time = time.time()
     print(f"Total time: {end_time - start_time:.2f} seconds (including IO)")
-    log_msg = f"Reconstruction Success ({len(all_files)} frames). Waiting for visualization."
+    num_frames = max(0, len(frame_filter_choices) - 1)
+    log_msg = f"Reconstruction Success ({num_frames} frames). Waiting for visualization."
 
     return (
         glbfile,
@@ -511,7 +642,7 @@ def update_visualization(
     loaded = np.load(predictions_path)
     predictions = {key: np.array(loaded[key]) for key in key_list}
 
-    sanitized_frame = frame_filter.replace('.', '_').replace(':', '').replace(' ', '_') if frame_filter else "All"
+    sanitized_frame = sanitize_frame_filter_label(frame_filter)
     glbfile = os.path.join(
         target_dir,
         f"glbscene_{conf_thres}_{sanitized_frame}_maskb{mask_black_bg}_maskw{mask_white_bg}_cam{show_cam}_sky{mask_sky}_pred{prediction_mode.replace(' ', '_')}_mode{mode_value}.glb",
@@ -702,7 +833,7 @@ with gr.Blocks(
                 streaming = gr.Radio(
                     [('stream', True), ('batch', False)],
                     label="Streaming or Batch Mode",
-                    value=False,
+                    value=True,
                     scale=1,
                 )
 
@@ -710,7 +841,7 @@ with gr.Blocks(
                 mode = gr.Radio(
                     ["causal", "window", "full"],
                     label="Select Processing Mode",
-                    value="causal",
+                    value="window",
                     scale=1,
                 )
 
@@ -801,7 +932,22 @@ with gr.Blocks(
             mode,
             False,
         )
-        return glbfile, log_msg, target_dir, dropdown, image_paths, session_file
+        return (
+            glbfile,
+            log_msg,
+            target_dir,
+            dropdown,
+            image_paths,
+            session_file,
+            False,
+            mode,
+            conf_thres,
+            mask_black_bg,
+            mask_white_bg,
+            show_cam,
+            mask_sky,
+            prediction_mode,
+        )
 
     gr.Markdown("Click any row to load an example.", elem_classes=["example-log"])
 
@@ -828,6 +974,14 @@ with gr.Blocks(
             frame_filter,
             image_gallery,
             session_state_output,
+            streaming,
+            mode,
+            conf_thres,
+            mask_black_bg,
+            mask_white_bg,
+            show_cam,
+            mask_sky,
+            prediction_mode,
         ],
         fn=example_pipeline,
         cache_examples=False,
@@ -981,7 +1135,22 @@ with gr.Blocks(
     # -------------------------------------------------------------------------
     # Auto-update gallery whenever user uploads or changes their files
     # -------------------------------------------------------------------------
-    upload_outputs = [reconstruction_output, target_dir_output, image_gallery, log_output, session_state_output]
+    upload_outputs = [
+        reconstruction_output,
+        target_dir_output,
+        image_gallery,
+        log_output,
+        session_state_output,
+        streaming,
+        mode,
+        conf_thres,
+        frame_filter,
+        mask_black_bg,
+        mask_white_bg,
+        show_cam,
+        mask_sky,
+        prediction_mode,
+    ]
     no_session_inputs = [input_video, input_images, input_zip, target_dir_output]
 
     input_video.change(fn=update_gallery_without_session, inputs=no_session_inputs, outputs=upload_outputs)

requirements.txt

@@ -41,6 +41,7 @@ scipy
 seaborn
 pyglet<2
 huggingface-hub[torch]>=0.22
+spaces
 
 # --------- eval --------- #
 accelerate