init

README.md

@@ -1,14 +1,40 @@
 ---
-title: LongStream
-emoji: 📊
-colorFrom: red
-colorTo: purple
+title: LongStream Demo
 sdk: gradio
-sdk_version: 6.9.0
+sdk_version: 5.44.0
 app_file: app.py
-pinned: false
-license: mit
-short_description: Demo of LongStream
+python_version: "3.10"
+startup_duration_timeout: 1h
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# LongStream Demo
+
+This repository is the Hugging Face Space package for LongStream.
+
+Project page: `https://3dagentworld.github.io/longstream/`
+
+## Space Settings
+
+Set these variables in the Space settings before the first run:
+
+- `LONGSTREAM_HF_REPO=NicolasCC/LongStream`
+- `LONGSTREAM_HF_FILE=50_longstream.pt`
+- `LONGSTREAM_HF_LOCAL_DIR=checkpoints`
+
+Optional:
+
+- `LONGSTREAM_HF_REVISION=v0.1.0`
+- `HF_TOKEN=<token>` if the model repo is private
+
+## Entrypoints
+
+- `app.py`: stable demo
+
+
+## Included Files
+
+- `demo_gradio.py`
+- `demo_gradio_interactive.py`
+- `longstream/`
+- `configs/longstream_infer.yaml`
+

app.py

@@ -0,0 +1,5 @@
+from demo_gradio import main
+
+
+if __name__ == "__main__":
+    main()

configs/longstream_infer.yaml

@@ -0,0 +1,84 @@
+device: cuda
+
+model:
+  checkpoint: checkpoints/50_longstream.pt
+  strict_load: false
+  hf:
+    repo_id: null
+    filename: null
+    revision: null
+    local_dir: checkpoints
+  longstream_cfg:
+    img_size: 518
+    patch_size: 14
+    embed_dim: 1024
+    window_size: 48
+    use_role_embedding: false
+    enable_scale_token: true
+    disable_keyframe_distinction: true
+    use_segment_mask: false
+    enable_camera_head: false
+    freeze: none
+    use_rel_pose_head: true
+    rel_pose_head_cfg:
+      enabled: true
+      keyframe_mode: fixed
+      keyframe_stride: 8
+      reference_source: pred
+      detach_reference: false
+      trunk_depth: 4
+      pose_mode: SE3
+      num_heads: 16
+      mlp_ratio: 4
+      init_values: 0.01
+      trans_act: linear
+      quat_act: linear
+      use_pair_cross_attn: false
+      xattn_temperature: 1.0
+      use_precat: false
+      use_kf_role_embed: false
+      kf_role_embed_init_std: 0.02
+      fl_act: relu
+      use_global_scale: false
+      reinit_camera_head: false
+
+inference:
+  mode: batch_refresh
+  streaming_mode: causal
+  window_size: 48
+  keyframe_mode: fixed
+  keyframe_stride: 8
+  refresh: 4
+  rel_pose_head_cfg:
+    num_iterations: 4
+
+data:
+  format: generalizable
+  data_roots_file: data_roots.txt
+  camera: null
+  img_path: "path/to/your/image/directory"
+  stride: 1
+  max_frames: null
+  size: 518
+  crop: false
+  patch_size: 14
+
+output:
+  root: outputs
+  save_videos: true
+  save_points: true
+  save_frame_points: true
+  save_depth: true
+  save_images: true
+  mask_sky: true
+  max_full_pointcloud_points: 2000000
+  max_frame_pointcloud_points: 200000
+  skyseg_path: skyseg.onnx
+
+evaluation:
+  align_scale: true
+  depth_rel_delta_threshold: 1.25
+  point_f1_threshold: 0.25
+  point_eval_max_points: 100000
+  point_eval_voxel_size: null
+  point_eval_oversample_factor: 4

demo_gradio.py

@@ -0,0 +1,332 @@
+import os
+
+import gradio as gr
+
+from longstream.demo import BRANCH_OPTIONS, create_demo_session, load_metadata
+from longstream.demo.backend import load_frame_previews
+from longstream.demo.export import export_glb
+from longstream.demo.viewer import build_interactive_figure
+
+
+DEFAULT_KEYFRAME_STRIDE = 8
+DEFAULT_REFRESH = 3
+DEFAULT_WINDOW_SIZE = 48
+DEFAULT_CHECKPOINT = os.getenv("LONGSTREAM_CHECKPOINT", "checkpoints/50_longstream.pt")
+
+
+def _run_stable_demo(
+    image_dir,
+    uploaded_files,
+    uploaded_video,
+    checkpoint,
+    device,
+    mode,
+    streaming_mode,
+    refresh,
+    window_size,
+    compute_sky,
+    branch_label,
+    show_cameras,
+    mask_sky,
+    camera_scale,
+    point_size,
+    opacity,
+    preview_max_points,
+    glb_max_points,
+):
+    if not image_dir and not uploaded_files and not uploaded_video:
+        raise gr.Error("Provide an image folder, upload images, or upload a video.")
+    session_dir = create_demo_session(
+        image_dir=image_dir or "",
+        uploaded_files=uploaded_files,
+        uploaded_video=uploaded_video,
+        checkpoint=checkpoint,
+        device=device,
+        mode=mode,
+        streaming_mode=streaming_mode,
+        keyframe_stride=DEFAULT_KEYFRAME_STRIDE,
+        refresh=int(refresh),
+        window_size=int(window_size),
+        compute_sky=bool(compute_sky),
+    )
+    fig = build_interactive_figure(
+        session_dir=session_dir,
+        branch=branch_label,
+        display_mode="All Frames",
+        frame_index=0,
+        point_size=float(point_size),
+        opacity=float(opacity),
+        preview_max_points=int(preview_max_points),
+        show_cameras=bool(show_cameras),
+        camera_scale=float(camera_scale),
+        mask_sky=bool(mask_sky),
+    )
+    glb_path = export_glb(
+        session_dir=session_dir,
+        branch=branch_label,
+        display_mode="All Frames",
+        frame_index=0,
+        mask_sky=bool(mask_sky),
+        show_cameras=bool(show_cameras),
+        camera_scale=float(camera_scale),
+        max_points=int(glb_max_points),
+    )
+    rgb, depth, frame_label = load_frame_previews(session_dir, 0)
+    meta = load_metadata(session_dir)
+    slider = gr.update(
+        minimum=0,
+        maximum=max(meta["num_frames"] - 1, 0),
+        value=0,
+        step=1,
+        interactive=meta["num_frames"] > 1,
+    )
+    sky_msg = ""
+    if meta.get("has_sky_masks"):
+        removed = float(meta.get("sky_removed_ratio") or 0.0) * 100.0
+        sky_msg = f" | sky_removed={removed:.1f}%"
+    status = f"Ready: {meta['num_frames']} frames | branch={branch_label}{sky_msg}"
+    return (
+        fig,
+        glb_path,
+        session_dir,
+        rgb,
+        depth,
+        frame_label,
+        slider,
+        status,
+    )
+
+
+def _update_stable_scene(
+    session_dir,
+    branch_label,
+    show_cameras,
+    mask_sky,
+    camera_scale,
+    point_size,
+    opacity,
+    preview_max_points,
+    glb_max_points,
+):
+    if not session_dir or not os.path.isdir(session_dir):
+        return None, None, "Run reconstruction first."
+    fig = build_interactive_figure(
+        session_dir=session_dir,
+        branch=branch_label,
+        display_mode="All Frames",
+        frame_index=0,
+        point_size=float(point_size),
+        opacity=float(opacity),
+        preview_max_points=int(preview_max_points),
+        show_cameras=bool(show_cameras),
+        camera_scale=float(camera_scale),
+        mask_sky=bool(mask_sky),
+    )
+    glb_path = export_glb(
+        session_dir=session_dir,
+        branch=branch_label,
+        display_mode="All Frames",
+        frame_index=0,
+        mask_sky=bool(mask_sky),
+        show_cameras=bool(show_cameras),
+        camera_scale=float(camera_scale),
+        max_points=int(glb_max_points),
+    )
+    meta = load_metadata(session_dir)
+    sky_msg = ""
+    if meta.get("has_sky_masks"):
+        removed = float(meta.get("sky_removed_ratio") or 0.0) * 100.0
+        sky_msg = f" | sky_removed={removed:.1f}%"
+    return fig, glb_path, f"Updated preview: {branch_label}{sky_msg}"
+
+
+def _update_frame_preview(session_dir, frame_index):
+    if not session_dir or not os.path.isdir(session_dir):
+        return None, None, ""
+    rgb, depth, label = load_frame_previews(session_dir, int(frame_index))
+    return rgb, depth, label
+
+
+def main():
+    with gr.Blocks(title="LongStream Demo") as demo:
+        session_dir = gr.Textbox(visible=False)
+
+        gr.Markdown("# LongStream Demo")
+
+        with gr.Row():
+            image_dir = gr.Textbox(
+                label="Image Folder", placeholder="/path/to/sequence"
+            )
+            uploaded_files = gr.File(
+                label="Upload Images", file_count="multiple", file_types=["image"]
+            )
+            uploaded_video = gr.File(
+                label="Upload Video", file_count="single", file_types=["video"]
+            )
+
+        with gr.Row():
+            checkpoint = gr.Textbox(label="Checkpoint", value=DEFAULT_CHECKPOINT)
+            device = gr.Dropdown(label="Device", choices=["cuda", "cpu"], value="cuda")
+
+        with gr.Accordion("Inference", open=False):
+            with gr.Row():
+                mode = gr.Dropdown(
+                    label="Mode",
+                    choices=["streaming_refresh", "batch_refresh"],
+                    value="batch_refresh",
+                )
+                streaming_mode = gr.Dropdown(
+                    label="Streaming Mode", choices=["causal", "window"], value="causal"
+                )
+            with gr.Row():
+                refresh = gr.Slider(
+                    label="Refresh", minimum=2, maximum=9, step=1, value=DEFAULT_REFRESH
+                )
+                window_size = gr.Slider(
+                    label="Window Size",
+                    minimum=1,
+                    maximum=64,
+                    step=1,
+                    value=DEFAULT_WINDOW_SIZE,
+                )
+                compute_sky = gr.Checkbox(label="Compute Sky Masks", value=True)
+
+        with gr.Accordion("GLB Settings", open=True):
+            with gr.Row():
+                branch_label = gr.Dropdown(
+                    label="Point Cloud Branch",
+                    choices=BRANCH_OPTIONS,
+                    value="Point Head + Pose",
+                )
+                show_cameras = gr.Checkbox(label="Show Cameras", value=True)
+                mask_sky = gr.Checkbox(label="Mask Sky", value=True)
+            with gr.Row():
+                point_size = gr.Slider(
+                    label="Point Size",
+                    minimum=0.05,
+                    maximum=2.0,
+                    step=0.05,
+                    value=0.3,
+                )
+                opacity = gr.Slider(
+                    label="Opacity",
+                    minimum=0.1,
+                    maximum=1.0,
+                    step=0.05,
+                    value=0.75,
+                )
+                preview_max_points = gr.Slider(
+                    label="Preview Max Points",
+                    minimum=5000,
+                    maximum=1000000,
+                    step=10000,
+                    value=100000,
+                )
+            with gr.Row():
+                camera_scale = gr.Slider(
+                    label="Camera Scale",
+                    minimum=0.001,
+                    maximum=0.05,
+                    step=0.001,
+                    value=0.01,
+                )
+                glb_max_points = gr.Slider(
+                    label="GLB Max Points",
+                    minimum=20000,
+                    maximum=1000000,
+                    step=10000,
+                    value=400000,
+                )
+
+        run_btn = gr.Button("Run Stable Demo", variant="primary")
+        status = gr.Markdown("Provide input images, then run reconstruction.")
+
+        plot = gr.Plot(label="Scene Preview")
+
+        glb_file = gr.File(label="Download GLB")
+
+        with gr.Row():
+            frame_slider = gr.Slider(
+                label="Preview Frame",
+                minimum=0,
+                maximum=0,
+                step=1,
+                value=0,
+                interactive=False,
+            )
+            frame_label = gr.Textbox(label="Frame")
+        with gr.Row():
+            rgb_preview = gr.Image(label="RGB", type="numpy")
+            depth_preview = gr.Image(label="Depth Plasma", type="numpy")
+
+        run_btn.click(
+            _run_stable_demo,
+            inputs=[
+                image_dir,
+                uploaded_files,
+                uploaded_video,
+                checkpoint,
+                device,
+                mode,
+                streaming_mode,
+                refresh,
+                window_size,
+                compute_sky,
+                branch_label,
+                show_cameras,
+                mask_sky,
+                camera_scale,
+                point_size,
+                opacity,
+                preview_max_points,
+                glb_max_points,
+            ],
+            outputs=[
+                plot,
+                glb_file,
+                session_dir,
+                rgb_preview,
+                depth_preview,
+                frame_label,
+                frame_slider,
+                status,
+            ],
+        )
+
+        for component in [
+            branch_label,
+            show_cameras,
+            mask_sky,
+            camera_scale,
+            point_size,
+            opacity,
+            preview_max_points,
+            glb_max_points,
+        ]:
+            component.change(
+                _update_stable_scene,
+                inputs=[
+                    session_dir,
+                    branch_label,
+                    show_cameras,
+                    mask_sky,
+                    camera_scale,
+                    point_size,
+                    opacity,
+                    preview_max_points,
+                    glb_max_points,
+                ],
+                outputs=[plot, glb_file, status],
+            )
+
+        frame_slider.change(
+            _update_frame_preview,
+            inputs=[session_dir, frame_slider],
+            outputs=[rgb_preview, depth_preview, frame_label],
+        )
+
+    demo.launch()
+
+
+if __name__ == "__main__":
+    main()

longstream/__init__.py

@@ -0,0 +1 @@
+__all__ = []

longstream/core/cli.py

@@ -0,0 +1,213 @@
+import argparse
+import os
+import sys
+
+import yaml
+
+
+def default_config_path() -> str:
+    return os.path.join(
+        os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
+        "configs",
+        "longstream_infer.yaml",
+    )
+
+
+def add_runtime_arguments(parser):
+    parser.add_argument(
+        "--config",
+        default=default_config_path(),
+        help="Path to longstream config yaml.",
+    )
+    parser.add_argument(
+        "--dataset",
+        default=None,
+        help="Optional dataset hint. Generic format works without it.",
+    )
+    parser.add_argument("--img-path", default=None)
+    parser.add_argument(
+        "--seq-list",
+        default=None,
+        help="Comma-separated sequence names. Default: auto-detect all sequences.",
+    )
+    parser.add_argument("--format", default=None, help="generalizable")
+    parser.add_argument("--data-roots-file", default=None)
+    parser.add_argument("--camera", default=None)
+    parser.add_argument("--output-root", default=None)
+    parser.add_argument("--device", default=None)
+    parser.add_argument("--checkpoint", default=None)
+    parser.add_argument("--hf-repo", default=None)
+    parser.add_argument("--hf-file", default=None)
+    parser.add_argument(
+        "--mode", default=None, help="batch_refresh | streaming_refresh"
+    )
+    parser.add_argument("--streaming-mode", default=None, help="causal | window")
+    parser.add_argument("--window-size", type=int, default=None)
+    parser.add_argument("--keyframe-stride", type=int, default=None)
+    parser.add_argument(
+        "--refresh",
+        type=int,
+        default=None,
+        help="Number of keyframes per refresh span, inclusive of both ends and including the segment start keyframe.",
+    )
+    parser.add_argument(
+        "--keyframes-per-batch",
+        dest="keyframes_per_batch_legacy",
+        type=int,
+        default=None,
+        help=argparse.SUPPRESS,
+    )
+    parser.add_argument("--max-frames", type=int, default=None)
+    parser.add_argument("--depth-rel-delta-threshold", type=float, default=None)
+    parser.add_argument("--point-f1-threshold", type=float, default=None)
+    parser.add_argument("--eval-max-points", type=int, default=None)
+    parser.add_argument("--eval-voxel-size", type=float, default=None)
+    parser.add_argument("--max-full-pointcloud-points", type=int, default=None)
+    parser.add_argument("--max-frame-pointcloud-points", type=int, default=None)
+    parser.add_argument("--save-frame-points", action="store_true")
+    parser.add_argument("--no-save-frame-points", action="store_true")
+    parser.add_argument("--no-align-scale", action="store_true")
+    parser.add_argument("--mask-sky", action="store_true")
+    parser.add_argument("--no-mask-sky", action="store_true")
+    return parser
+
+
+def parse_runtime_args(parser):
+    argv = [arg for arg in sys.argv[1:] if arg.strip()]
+    return parser.parse_args(argv)
+
+
+def load_config_with_overrides(args):
+    with open(args.config, "r") as f:
+        cfg = yaml.safe_load(f) or {}
+    cfg.setdefault("model", {})
+
+    if args.device is not None:
+        cfg["device"] = args.device
+
+    if args.output_root is not None:
+        cfg.setdefault("output", {})
+        cfg["output"]["root"] = args.output_root
+
+    if args.dataset is not None:
+        cfg.setdefault("data", {})
+        cfg["data"]["dataset"] = args.dataset
+
+    if args.img_path is not None:
+        cfg.setdefault("data", {})
+        cfg["data"]["img_path"] = args.img_path
+
+    if args.seq_list is not None:
+        seqs = [s.strip() for s in args.seq_list.split(",") if s.strip()]
+        cfg.setdefault("data", {})
+        cfg["data"]["seq_list"] = seqs
+
+    if args.format is not None:
+        cfg.setdefault("data", {})
+        cfg["data"]["format"] = args.format
+
+    if args.data_roots_file is not None:
+        cfg.setdefault("data", {})
+        cfg["data"]["data_roots_file"] = args.data_roots_file
+
+    if args.camera is not None:
+        cfg.setdefault("data", {})
+        cfg["data"]["camera"] = args.camera
+
+    if args.max_frames is not None:
+        cfg.setdefault("data", {})
+        cfg["data"]["max_frames"] = args.max_frames
+
+    if args.checkpoint is not None:
+        cfg.setdefault("model", {})
+        cfg["model"]["checkpoint"] = args.checkpoint
+
+    if args.hf_repo is not None or args.hf_file is not None:
+        cfg.setdefault("model", {})
+        cfg["model"].setdefault("hf", {})
+        if args.hf_repo is not None:
+            cfg["model"]["hf"]["repo_id"] = args.hf_repo
+        if args.hf_file is not None:
+            cfg["model"]["hf"]["filename"] = args.hf_file
+        if cfg["model"].get("checkpoint") is None:
+            cfg["model"]["checkpoint"] = None
+
+    if args.mode is not None:
+        cfg.setdefault("inference", {})
+        cfg["inference"]["mode"] = args.mode
+
+    if args.streaming_mode is not None:
+        cfg.setdefault("inference", {})
+        cfg["inference"]["streaming_mode"] = args.streaming_mode
+
+    if args.window_size is not None:
+        cfg.setdefault("inference", {})
+        cfg["inference"]["window_size"] = args.window_size
+        cfg["model"].setdefault("longstream_cfg", {})
+        cfg["model"]["longstream_cfg"]["window_size"] = args.window_size
+
+    if args.keyframe_stride is not None:
+        cfg.setdefault("inference", {})
+        cfg["inference"]["keyframe_stride"] = args.keyframe_stride
+        cfg["model"].setdefault("longstream_cfg", {})
+        cfg["model"]["longstream_cfg"].setdefault("rel_pose_head_cfg", {})
+        cfg["model"]["longstream_cfg"]["rel_pose_head_cfg"][
+            "keyframe_stride"
+        ] = args.keyframe_stride
+
+    refresh = args.refresh
+    if refresh is None and args.keyframes_per_batch_legacy is not None:
+        refresh = args.keyframes_per_batch_legacy + 1
+    if refresh is not None:
+        cfg.setdefault("inference", {})
+        cfg["inference"]["refresh"] = refresh
+
+    if args.depth_rel_delta_threshold is not None:
+        cfg.setdefault("evaluation", {})
+        cfg["evaluation"]["depth_rel_delta_threshold"] = args.depth_rel_delta_threshold
+
+    if args.point_f1_threshold is not None:
+        cfg.setdefault("evaluation", {})
+        cfg["evaluation"]["point_f1_threshold"] = args.point_f1_threshold
+
+    if args.eval_max_points is not None:
+        cfg.setdefault("evaluation", {})
+        cfg["evaluation"]["point_eval_max_points"] = args.eval_max_points
+
+    if args.eval_voxel_size is not None:
+        cfg.setdefault("evaluation", {})
+        cfg["evaluation"]["point_eval_voxel_size"] = args.eval_voxel_size
+
+    if args.max_full_pointcloud_points is not None:
+        cfg.setdefault("output", {})
+        cfg["output"]["max_full_pointcloud_points"] = args.max_full_pointcloud_points
+
+    if args.max_frame_pointcloud_points is not None:
+        cfg.setdefault("output", {})
+        cfg["output"]["max_frame_pointcloud_points"] = args.max_frame_pointcloud_points
+
+    if args.save_frame_points:
+        cfg.setdefault("output", {})
+        cfg["output"]["save_frame_points"] = True
+    if args.no_save_frame_points:
+        cfg.setdefault("output", {})
+        cfg["output"]["save_frame_points"] = False
+
+    if args.no_align_scale:
+        cfg.setdefault("evaluation", {})
+        cfg["evaluation"]["align_scale"] = False
+
+    if args.mask_sky:
+        cfg.setdefault("output", {})
+        cfg["output"]["mask_sky"] = True
+    if args.no_mask_sky:
+        cfg.setdefault("output", {})
+        cfg["output"]["mask_sky"] = False
+
+    infer_cfg = cfg.setdefault("inference", {})
+    if "refresh" not in infer_cfg and "keyframes_per_batch" in infer_cfg:
+        infer_cfg["refresh"] = int(infer_cfg["keyframes_per_batch"]) + 1
+
+    cfg.setdefault("data", {})
+    cfg["data"]["format"] = "generalizable"
+    return cfg

longstream/core/infer.py

@@ -0,0 +1,451 @@
+import argparse
+import os
+import yaml
+import cv2
+import numpy as np
+import torch
+from PIL import Image
+
+from longstream.core.model import LongStreamModel
+from longstream.data.dataloader import LongStreamDataLoader
+from longstream.streaming.keyframe_selector import KeyframeSelector
+from longstream.streaming.refresh import run_batch_refresh, run_streaming_refresh
+from longstream.utils.vendor.models.components.utils.pose_enc import (
+    pose_encoding_to_extri_intri,
+)
+from longstream.utils.camera import compose_abs_from_rel
+from longstream.utils.depth import colorize_depth, unproject_depth_to_points
+from longstream.utils.sky_mask import compute_sky_mask
+from longstream.io.save_points import save_pointcloud
+from longstream.io.save_poses_txt import save_w2c_txt, save_intri_txt, save_rel_pose_txt
+from longstream.io.save_images import save_image_sequence, save_video
+
+
+def _to_uint8_rgb(images):
+    imgs = images.detach().cpu().numpy()
+    imgs = np.clip(imgs, 0.0, 1.0)
+    imgs = (imgs * 255.0).astype(np.uint8)
+    return imgs
+
+
+def _ensure_dir(path):
+    os.makedirs(path, exist_ok=True)
+
+
+def _apply_sky_mask(depth, mask):
+    if mask is None:
+        return depth
+    m = (mask > 0).astype(np.float32)
+    return depth * m
+
+
+def _camera_points_to_world(points, extri):
+    pts = np.asarray(points, dtype=np.float64).reshape(-1, 3)
+    R = np.asarray(extri[:3, :3], dtype=np.float64)
+    t = np.asarray(extri[:3, 3], dtype=np.float64)
+    world = (R.T @ (pts.T - t[:, None])).T
+    return world.astype(np.float32, copy=False)
+
+
+def _mask_points_and_colors(points, colors, mask):
+    pts = points.reshape(-1, 3)
+    cols = None if colors is None else colors.reshape(-1, 3)
+    if mask is None:
+        return pts, cols
+    valid = mask.reshape(-1) > 0
+    pts = pts[valid]
+    if cols is not None:
+        cols = cols[valid]
+    return pts, cols
+
+
+def _resize_long_edge(arr, long_edge_size, interpolation):
+    h, w = arr.shape[:2]
+    scale = float(long_edge_size) / float(max(h, w))
+    new_w = int(round(w * scale))
+    new_h = int(round(h * scale))
+    return cv2.resize(arr, (new_w, new_h), interpolation=interpolation)
+
+
+def _prepare_mask_for_model(
+    mask, size, crop, patch_size, target_shape, square_ok=False
+):
+    if mask is None:
+        return None
+    long_edge = (
+        round(size * max(mask.shape[1] / mask.shape[0], mask.shape[0] / mask.shape[1]))
+        if size == 224
+        else size
+    )
+    mask = _resize_long_edge(mask, long_edge, cv2.INTER_NEAREST)
+
+    h, w = mask.shape[:2]
+    cx, cy = w // 2, h // 2
+    if size == 224:
+        half = min(cx, cy)
+        target_w = 2 * half
+        target_h = 2 * half
+        if crop:
+            mask = mask[cy - half : cy + half, cx - half : cx + half]
+        else:
+            mask = cv2.resize(
+                mask, (target_w, target_h), interpolation=cv2.INTER_NEAREST
+            )
+    else:
+        halfw = ((2 * cx) // patch_size) * (patch_size // 2)
+        halfh = ((2 * cy) // patch_size) * (patch_size // 2)
+        if not square_ok and w == h:
+            halfh = int(3 * halfw / 4)
+        target_w = 2 * halfw
+        target_h = 2 * halfh
+        if crop:
+            mask = mask[cy - halfh : cy + halfh, cx - halfw : cx + halfw]
+        else:
+            mask = cv2.resize(
+                mask, (target_w, target_h), interpolation=cv2.INTER_NEAREST
+            )
+
+    if mask.shape[:2] != tuple(target_shape):
+        mask = cv2.resize(
+            mask, (target_shape[1], target_shape[0]), interpolation=cv2.INTER_NEAREST
+        )
+    return mask
+
+
+def _save_full_pointcloud(path, point_chunks, color_chunks, max_points=None, seed=0):
+    if not point_chunks:
+        return
+    points = np.concatenate(point_chunks, axis=0)
+    colors = None
+    if color_chunks and len(color_chunks) == len(point_chunks):
+        colors = np.concatenate(color_chunks, axis=0)
+    if max_points is not None and len(points) > max_points:
+        rng = np.random.default_rng(seed)
+        keep = rng.choice(len(points), size=max_points, replace=False)
+        points = points[keep]
+        if colors is not None:
+            colors = colors[keep]
+    np.save(os.path.splitext(path)[0] + ".npy", points.astype(np.float32, copy=False))
+    save_pointcloud(path, points, colors=colors, max_points=None, seed=seed)
+
+
+def run_inference_cfg(cfg: dict):
+    device = cfg.get("device", "cuda" if torch.cuda.is_available() else "cpu")
+    device_type = torch.device(device).type
+    model_cfg = cfg.get("model", {})
+    data_cfg = cfg.get("data", {})
+    infer_cfg = cfg.get("inference", {})
+    output_cfg = cfg.get("output", {})
+
+    print(f"[longstream] device={device}", flush=True)
+    model = LongStreamModel(model_cfg).to(device)
+    model.eval()
+    print("[longstream] model ready", flush=True)
+
+    loader = LongStreamDataLoader(data_cfg)
+
+    keyframe_stride = int(infer_cfg.get("keyframe_stride", 8))
+    keyframe_mode = infer_cfg.get("keyframe_mode", "fixed")
+    refresh = int(
+        infer_cfg.get("refresh", int(infer_cfg.get("keyframes_per_batch", 3)) + 1)
+    )
+    if refresh < 2:
+        raise ValueError(
+            "refresh must be >= 2 because it counts both keyframe endpoints"
+        )
+    mode = infer_cfg.get("mode", "streaming_refresh")
+    if mode == "streaming":
+        mode = "streaming_refresh"
+    streaming_mode = infer_cfg.get("streaming_mode", "causal")
+    window_size = int(infer_cfg.get("window_size", 5))
+
+    selector = KeyframeSelector(
+        min_interval=keyframe_stride,
+        max_interval=keyframe_stride,
+        force_first=True,
+        mode="random" if keyframe_mode == "random" else "fixed",
+    )
+
+    out_root = output_cfg.get("root", "outputs")
+    _ensure_dir(out_root)
+    save_videos = bool(output_cfg.get("save_videos", True))
+    save_points = bool(output_cfg.get("save_points", True))
+    save_frame_points = bool(output_cfg.get("save_frame_points", True))
+    save_depth = bool(output_cfg.get("save_depth", True))
+    save_images = bool(output_cfg.get("save_images", True))
+    mask_sky = bool(output_cfg.get("mask_sky", True))
+    max_full_pointcloud_points = output_cfg.get("max_full_pointcloud_points", None)
+    if max_full_pointcloud_points is not None:
+        max_full_pointcloud_points = int(max_full_pointcloud_points)
+    max_frame_pointcloud_points = output_cfg.get("max_frame_pointcloud_points", None)
+    if max_frame_pointcloud_points is not None:
+        max_frame_pointcloud_points = int(max_frame_pointcloud_points)
+    skyseg_path = output_cfg.get(
+        "skyseg_path",
+        os.path.join(os.path.dirname(__file__), "..", "..", "skyseg.onnx"),
+    )
+
+    with torch.no_grad():
+        for seq in loader:
+            images = seq.images
+            B, S, C, H, W = images.shape
+            print(
+                f"[longstream] sequence {seq.name}: inference start ({S} frames)",
+                flush=True,
+            )
+
+            is_keyframe, keyframe_indices = selector.select_keyframes(
+                S, B, images.device
+            )
+
+            rel_pose_cfg = infer_cfg.get("rel_pose_head_cfg", {"num_iterations": 4})
+
+            if mode == "batch_refresh":
+                outputs = run_batch_refresh(
+                    model,
+                    images,
+                    is_keyframe,
+                    keyframe_indices,
+                    streaming_mode,
+                    keyframe_stride,
+                    refresh,
+                    rel_pose_cfg,
+                )
+            elif mode == "streaming_refresh":
+                outputs = run_streaming_refresh(
+                    model,
+                    images,
+                    is_keyframe,
+                    keyframe_indices,
+                    streaming_mode,
+                    window_size,
+                    refresh,
+                    rel_pose_cfg,
+                )
+            else:
+                raise ValueError(f"Unsupported inference mode: {mode}")
+            print(f"[longstream] sequence {seq.name}: inference done", flush=True)
+            if device_type == "cuda":
+                torch.cuda.empty_cache()
+
+            seq_dir = os.path.join(out_root, seq.name)
+            _ensure_dir(seq_dir)
+
+            frame_ids = list(range(S))
+            rgb = _to_uint8_rgb(images[0].permute(0, 2, 3, 1))
+
+            if "rel_pose_enc" in outputs:
+                rel_pose_enc = outputs["rel_pose_enc"][0]
+                abs_pose_enc = compose_abs_from_rel(rel_pose_enc, keyframe_indices[0])
+                extri, intri = pose_encoding_to_extri_intri(
+                    abs_pose_enc[None], image_size_hw=(H, W)
+                )
+                extri_np = extri[0].detach().cpu().numpy()
+                intri_np = intri[0].detach().cpu().numpy()
+
+                pose_dir = os.path.join(seq_dir, "poses")
+                _ensure_dir(pose_dir)
+                save_w2c_txt(
+                    os.path.join(pose_dir, "abs_pose.txt"), extri_np, frame_ids
+                )
+                save_intri_txt(os.path.join(pose_dir, "intri.txt"), intri_np, frame_ids)
+                save_rel_pose_txt(
+                    os.path.join(pose_dir, "rel_pose.txt"), rel_pose_enc, frame_ids
+                )
+            elif "pose_enc" in outputs:
+                pose_enc = outputs["pose_enc"][0]
+                extri, intri = pose_encoding_to_extri_intri(
+                    pose_enc[None], image_size_hw=(H, W)
+                )
+                extri_np = extri[0].detach().cpu().numpy()
+                intri_np = intri[0].detach().cpu().numpy()
+
+                pose_dir = os.path.join(seq_dir, "poses")
+                _ensure_dir(pose_dir)
+                save_w2c_txt(
+                    os.path.join(pose_dir, "abs_pose.txt"), extri_np, frame_ids
+                )
+                save_intri_txt(os.path.join(pose_dir, "intri.txt"), intri_np, frame_ids)
+
+            if save_images:
+                print(f"[longstream] sequence {seq.name}: saving rgb", flush=True)
+                rgb_dir = os.path.join(seq_dir, "images", "rgb")
+                save_image_sequence(rgb_dir, list(rgb))
+                if save_videos:
+                    save_video(
+                        os.path.join(seq_dir, "images", "rgb.mp4"),
+                        os.path.join(rgb_dir, "frame_*.png"),
+                    )
+
+            sky_masks = None
+            if mask_sky:
+                raw_sky_masks = compute_sky_mask(
+                    seq.image_paths, skyseg_path, os.path.join(seq_dir, "sky_masks")
+                )
+                if raw_sky_masks is not None:
+                    sky_masks = [
+                        _prepare_mask_for_model(
+                            mask,
+                            size=int(data_cfg.get("size", 518)),
+                            crop=bool(data_cfg.get("crop", False)),
+                            patch_size=int(data_cfg.get("patch_size", 14)),
+                            target_shape=(H, W),
+                        )
+                        for mask in raw_sky_masks
+                    ]
+
+            if save_depth and "depth" in outputs:
+                print(f"[longstream] sequence {seq.name}: saving depth", flush=True)
+                depth = outputs["depth"][0, :, :, :, 0].detach().cpu().numpy()
+                depth_dir = os.path.join(seq_dir, "depth", "dpt")
+                _ensure_dir(depth_dir)
+                color_dir = os.path.join(seq_dir, "depth", "dpt_plasma")
+                _ensure_dir(color_dir)
+
+                color_frames = []
+                for i in range(S):
+                    d = depth[i]
+                    if sky_masks is not None and sky_masks[i] is not None:
+                        d = _apply_sky_mask(d, sky_masks[i])
+                    np.save(os.path.join(depth_dir, f"frame_{i:06d}.npy"), d)
+                    colored = colorize_depth(d, cmap="plasma")
+                    Image.fromarray(colored).save(
+                        os.path.join(color_dir, f"frame_{i:06d}.png")
+                    )
+                    color_frames.append(colored)
+                if save_videos:
+                    save_video(
+                        os.path.join(seq_dir, "depth", "dpt_plasma.mp4"),
+                        os.path.join(color_dir, "frame_*.png"),
+                    )
+
+            if save_points:
+                print(
+                    f"[longstream] sequence {seq.name}: saving point clouds", flush=True
+                )
+                if "world_points" in outputs:
+                    if "rel_pose_enc" in outputs:
+                        abs_pose_enc = compose_abs_from_rel(
+                            outputs["rel_pose_enc"][0], keyframe_indices[0]
+                        )
+                        extri, intri = pose_encoding_to_extri_intri(
+                            abs_pose_enc[None], image_size_hw=(H, W)
+                        )
+                    else:
+                        extri, intri = pose_encoding_to_extri_intri(
+                            outputs["pose_enc"][0][None], image_size_hw=(H, W)
+                        )
+                    extri = extri[0]
+                    intri = intri[0]
+
+                    pts_dir = os.path.join(seq_dir, "points", "point_head")
+                    _ensure_dir(pts_dir)
+                    pts = outputs["world_points"][0].detach().cpu().numpy()
+                    full_pts = []
+                    full_cols = []
+                    for i in range(S):
+                        pts_world = _camera_points_to_world(
+                            pts[i], extri[i].detach().cpu().numpy()
+                        )
+                        pts_world = pts_world.reshape(pts[i].shape)
+                        pts_i, cols_i = _mask_points_and_colors(
+                            pts_world,
+                            rgb[i],
+                            None if sky_masks is None else sky_masks[i],
+                        )
+                        if save_frame_points:
+                            save_pointcloud(
+                                os.path.join(pts_dir, f"frame_{i:06d}.ply"),
+                                pts_i,
+                                colors=cols_i,
+                                max_points=max_frame_pointcloud_points,
+                                seed=i,
+                            )
+                        if len(pts_i):
+                            full_pts.append(pts_i)
+                            full_cols.append(cols_i)
+                    _save_full_pointcloud(
+                        os.path.join(seq_dir, "points", "point_head_full.ply"),
+                        full_pts,
+                        full_cols,
+                        max_points=max_full_pointcloud_points,
+                        seed=0,
+                    )
+
+                if "depth" in outputs and (
+                    "rel_pose_enc" in outputs or "pose_enc" in outputs
+                ):
+                    depth = outputs["depth"][0, :, :, :, 0]
+                    if "rel_pose_enc" in outputs:
+                        abs_pose_enc = compose_abs_from_rel(
+                            outputs["rel_pose_enc"][0], keyframe_indices[0]
+                        )
+                        extri, intri = pose_encoding_to_extri_intri(
+                            abs_pose_enc[None], image_size_hw=(H, W)
+                        )
+                    else:
+                        extri, intri = pose_encoding_to_extri_intri(
+                            outputs["pose_enc"][0][None], image_size_hw=(H, W)
+                        )
+
+                    extri = extri[0]
+                    intri = intri[0]
+                    dpt_pts_dir = os.path.join(seq_dir, "points", "dpt_unproj")
+                    _ensure_dir(dpt_pts_dir)
+                    full_pts = []
+                    full_cols = []
+
+                    for i in range(S):
+                        d = depth[i]
+                        pts_cam = unproject_depth_to_points(d[None], intri[i : i + 1])[
+                            0
+                        ]
+                        R = extri[i, :3, :3]
+                        t = extri[i, :3, 3]
+                        pts_world = (
+                            R.t() @ (pts_cam.reshape(-1, 3).t() - t[:, None])
+                        ).t()
+                        pts_world = pts_world.cpu().numpy().reshape(-1, 3)
+                        pts_i, cols_i = _mask_points_and_colors(
+                            pts_world,
+                            rgb[i],
+                            None if sky_masks is None else sky_masks[i],
+                        )
+                        if save_frame_points:
+                            save_pointcloud(
+                                os.path.join(dpt_pts_dir, f"frame_{i:06d}.ply"),
+                                pts_i,
+                                colors=cols_i,
+                                max_points=max_frame_pointcloud_points,
+                                seed=i,
+                            )
+                        if len(pts_i):
+                            full_pts.append(pts_i)
+                            full_cols.append(cols_i)
+                    _save_full_pointcloud(
+                        os.path.join(seq_dir, "points", "dpt_unproj_full.ply"),
+                        full_pts,
+                        full_cols,
+                        max_points=max_full_pointcloud_points,
+                        seed=1,
+                    )
+            del outputs
+            if device_type == "cuda":
+                torch.cuda.empty_cache()
+
+
+def run_inference(config_path: str):
+    with open(config_path, "r") as f:
+        cfg = yaml.safe_load(f)
+    run_inference_cfg(cfg)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", required=True)
+    args = parser.parse_args()
+    run_inference(args.config)
+
+
+if __name__ == "__main__":
+    main()

longstream/core/model.py

@@ -0,0 +1,69 @@
+import os
+import torch
+from typing import Dict, Any
+
+from longstream.models.longstream import LongStream
+from longstream.utils.hub import resolve_checkpoint_path
+
+
+class LongStreamModel(torch.nn.Module):
+    def __init__(self, cfg: Dict[str, Any] | None):
+        super().__init__()
+        cfg = cfg or {}
+
+        ckpt_path = resolve_checkpoint_path(
+            cfg.get("checkpoint", None), cfg.get("hf", None)
+        )
+
+        stream_cfg = dict(cfg.get("longstream_cfg", {}) or {})
+        rel_pose_cfg = stream_cfg.pop(
+            "rel_pose_head_cfg", cfg.get("rel_pose_head_cfg", None)
+        )
+        use_rel_pose_head = bool(stream_cfg.pop("use_rel_pose_head", False))
+        if use_rel_pose_head and rel_pose_cfg is not None:
+            stream_cfg["rel_pose_head_cfg"] = rel_pose_cfg
+        self.longstream = LongStream(**stream_cfg)
+
+        if ckpt_path:
+            self.load_checkpoint(ckpt_path, strict=bool(cfg.get("strict_load", True)))
+
+    def load_checkpoint(self, ckpt_path: str, strict: bool = True):
+        if not os.path.exists(ckpt_path):
+            raise FileNotFoundError(ckpt_path)
+        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        if isinstance(ckpt, dict):
+            if "model" in ckpt and isinstance(ckpt["model"], dict):
+                state = ckpt["model"]
+            elif "state_dict" in ckpt and isinstance(ckpt["state_dict"], dict):
+                state = ckpt["state_dict"]
+            else:
+                state = ckpt
+        else:
+            raise TypeError("Unsupported checkpoint format")
+
+        if state:
+            first_key = next(iter(state.keys()))
+            if first_key.startswith("sampler.longstream."):
+                state = {k.replace("sampler.", "", 1): v for k, v in state.items()}
+
+        missing, unexpected = self.load_state_dict(state, strict=False)
+        if missing or unexpected:
+            msg = f"checkpoint mismatch: missing={len(missing)} unexpected={len(unexpected)}"
+            if strict:
+                raise RuntimeError(msg)
+            print(msg)
+
+    def forward(self, *args, **kwargs):
+        return self.longstream(*args, **kwargs)
+
+    @property
+    def aggregator(self):
+        return self.longstream.aggregator
+
+    @property
+    def camera_head(self):
+        return getattr(self.longstream, "camera_head", None)
+
+    @property
+    def rel_pose_head(self):
+        return getattr(self.longstream, "rel_pose_head", None)

longstream/data/__init__.py

@@ -0,0 +1,3 @@
+from .dataloader import LongStreamDataLoader, LongStreamSequence, LongStreamSequenceInfo
+
+__all__ = ["LongStreamDataLoader", "LongStreamSequence", "LongStreamSequenceInfo"]

longstream/data/dataloader.py

@@ -0,0 +1,422 @@
+import os
+import glob
+from dataclasses import dataclass
+from typing import List, Dict, Any, Iterator, Optional, Tuple
+
+import torch
+
+from longstream.utils.vendor.dust3r.utils.image import load_images_for_eval
+
+dataset_metadata: Dict[str, Dict[str, Any]] = {
+    "davis": {
+        "img_path": "data/davis/DAVIS/JPEGImages/480p",
+        "mask_path": "data/davis/DAVIS/masked_images/480p",
+        "dir_path_func": lambda img_path, seq: os.path.join(img_path, seq),
+        "gt_traj_func": lambda img_path, anno_path, seq: None,
+        "traj_format": None,
+        "seq_list": None,
+        "full_seq": True,
+        "mask_path_seq_func": lambda mask_path, seq: os.path.join(mask_path, seq),
+        "skip_condition": None,
+        "process_func": None,
+    },
+    "kitti": {
+        "img_path": "data/kitti/sequences",
+        "anno_path": "data/kitti/poses",
+        "mask_path": None,
+        "dir_path_func": lambda img_path, seq: os.path.join(img_path, seq, "image_2"),
+        "gt_traj_func": lambda img_path, anno_path, seq: os.path.join(
+            anno_path, f"{seq}.txt"
+        )
+        if os.path.exists(os.path.join(anno_path, f"{seq}.txt"))
+        else None,
+        "traj_format": "kitti",
+        "seq_list": ["00", "01", "02", "03", "04", "05", "06", "07", "08", "09", "10"],
+        "full_seq": True,
+        "mask_path_seq_func": lambda mask_path, seq: None,
+        "skip_condition": None,
+        "process_func": None,
+    },
+    "bonn": {
+        "img_path": "data/bonn/rgbd_bonn_dataset",
+        "mask_path": None,
+        "dir_path_func": lambda img_path, seq: os.path.join(
+            img_path, f"rgbd_bonn_{seq}", "rgb_110"
+        ),
+        "gt_traj_func": lambda img_path, anno_path, seq: os.path.join(
+            img_path, f"rgbd_bonn_{seq}", "groundtruth_110.txt"
+        ),
+        "traj_format": "tum",
+        "seq_list": ["balloon2", "crowd2", "crowd3", "person_tracking2", "synchronous"],
+        "full_seq": False,
+        "mask_path_seq_func": lambda mask_path, seq: None,
+        "skip_condition": None,
+        "process_func": None,
+    },
+    "nyu": {
+        "img_path": "data/nyu-v2/val/nyu_images",
+        "mask_path": None,
+        "process_func": None,
+    },
+    "scannet": {
+        "img_path": "data/scannetv2",
+        "mask_path": None,
+        "dir_path_func": lambda img_path, seq: os.path.join(img_path, seq, "color_90"),
+        "gt_traj_func": lambda img_path, anno_path, seq: os.path.join(
+            img_path, seq, "pose_90.txt"
+        ),
+        "traj_format": "replica",
+        "seq_list": None,
+        "full_seq": True,
+        "mask_path_seq_func": lambda mask_path, seq: None,
+        "skip_condition": None,
+        "process_func": None,
+    },
+    "tum": {
+        "img_path": "data/tum",
+        "mask_path": None,
+        "dir_path_func": lambda img_path, seq: os.path.join(img_path, seq, "rgb_90"),
+        "gt_traj_func": lambda img_path, anno_path, seq: os.path.join(
+            img_path, seq, "groundtruth_90.txt"
+        ),
+        "traj_format": "tum",
+        "seq_list": None,
+        "full_seq": True,
+        "mask_path_seq_func": lambda mask_path, seq: None,
+        "skip_condition": None,
+        "process_func": None,
+    },
+    "sintel": {
+        "img_path": "data/sintel/training/final",
+        "anno_path": "data/sintel/training/camdata_left",
+        "mask_path": None,
+        "dir_path_func": lambda img_path, seq: os.path.join(img_path, seq),
+        "gt_traj_func": lambda img_path, anno_path, seq: os.path.join(anno_path, seq),
+        "traj_format": None,
+        "seq_list": [
+            "alley_2",
+            "ambush_4",
+            "ambush_5",
+            "ambush_6",
+            "cave_2",
+            "cave_4",
+            "market_2",
+            "market_5",
+            "market_6",
+            "shaman_3",
+            "sleeping_1",
+            "sleeping_2",
+            "temple_2",
+            "temple_3",
+        ],
+        "full_seq": False,
+        "mask_path_seq_func": lambda mask_path, seq: None,
+        "skip_condition": None,
+        "process_func": None,
+    },
+    "waymo": {
+        "img_path": "/horizon-bucket/saturn_v_4dlabel/004_vision/01_users/tao02.xie/datasets/scatt3r_evaluation/waymo_open_dataset_v1_4_3",
+        "anno_path": None,
+        "mask_path": None,
+        "dir_path_func": lambda img_path, seq: os.path.join(
+            img_path,
+            seq.split("_cam")[0] if "_cam" in seq else seq,
+            "images",
+            seq.split("_cam")[1] if "_cam" in seq else "00",
+        ),
+        "gt_traj_func": lambda img_path, anno_path, seq: os.path.join(
+            img_path,
+            seq.split("_cam")[0] if "_cam" in seq else seq,
+            "cameras",
+            seq.split("_cam")[1] if "_cam" in seq else "00",
+            "extri.yml",
+        ),
+        "traj_format": "waymo",
+        "seq_list": None,
+        "full_seq": True,
+        "mask_path_seq_func": lambda mask_path, seq: None,
+        "skip_condition": None,
+        "process_func": None,
+    },
+}
+
+
+@dataclass
+class LongStreamSequenceInfo:
+    name: str
+    scene_root: str
+    image_dir: str
+    image_paths: List[str]
+    camera: Optional[str]
+
+
+class LongStreamSequence:
+    def __init__(
+        self,
+        name: str,
+        images: torch.Tensor,
+        image_paths: List[str],
+        scene_root: Optional[str] = None,
+        image_dir: Optional[str] = None,
+        camera: Optional[str] = None,
+    ):
+        self.name = name
+        self.images = images
+        self.image_paths = image_paths
+        self.scene_root = scene_root
+        self.image_dir = image_dir
+        self.camera = camera
+
+
+def _read_list_file(path: str) -> List[str]:
+    with open(path, "r") as f:
+        lines = []
+        for line in f.readlines():
+            line = line.strip()
+            if not line:
+                continue
+            if line.startswith("#"):
+                continue
+            lines.append(line)
+    return lines
+
+
+def _is_generalizable_scene_root(path: str) -> bool:
+    return os.path.isdir(os.path.join(path, "images"))
+
+
+def _direct_image_files(dir_path: str) -> List[str]:
+    filelist = sorted(glob.glob(os.path.join(dir_path, "*.png")))
+    if not filelist:
+        filelist = sorted(glob.glob(os.path.join(dir_path, "*.jpg")))
+    if not filelist:
+        filelist = sorted(glob.glob(os.path.join(dir_path, "*.jpeg")))
+    return filelist
+
+
+class LongStreamDataLoader:
+    def __init__(self, cfg: Dict[str, Any]):
+        self.cfg = cfg
+        self.dataset = cfg.get("dataset", None)
+        meta = dataset_metadata.get(self.dataset, {})
+        self.img_path = cfg.get("img_path", meta.get("img_path"))
+        self.mask_path = cfg.get("mask_path", meta.get("mask_path"))
+        self.dir_path_func = meta.get("dir_path_func", lambda p, s: os.path.join(p, s))
+        self.mask_path_seq_func = meta.get("mask_path_seq_func", lambda p, s: None)
+        self.full_seq = bool(cfg.get("full_seq", meta.get("full_seq", True)))
+        self.seq_list = cfg.get("seq_list", None)
+        self.stride = int(cfg.get("stride", 1))
+        self.max_frames = cfg.get("max_frames", None)
+        self.size = int(cfg.get("size", 518))
+        self.crop = bool(cfg.get("crop", False))
+        self.patch_size = int(cfg.get("patch_size", 14))
+        self.format = cfg.get("format", "auto")
+        self.data_roots_file = cfg.get("data_roots_file", None)
+        self.split = cfg.get("split", None)
+        self.camera = cfg.get("camera", None)
+
+    def _infer_format(self) -> str:
+        if self.format in ["relpose", "generalizable"]:
+            return self.format
+        if self.img_path is None:
+            return "relpose"
+        if _is_generalizable_scene_root(self.img_path):
+            return "generalizable"
+        default_list = self.data_roots_file or "data_roots.txt"
+        if os.path.exists(os.path.join(self.img_path, default_list)):
+            return "generalizable"
+        return "relpose"
+
+    def _resolve_seq_list_generalizable(self) -> List[str]:
+        if self.seq_list is not None:
+            return list(self.seq_list)
+        if self.img_path is None or not os.path.isdir(self.img_path):
+            return []
+
+        if _is_generalizable_scene_root(self.img_path):
+            return [self.img_path]
+
+        candidates = []
+        if isinstance(self.data_roots_file, str) and self.data_roots_file:
+            candidates.append(self.data_roots_file)
+        if isinstance(self.split, str) and self.split:
+            split_name = self.split.lower()
+            if split_name in ["val", "valid", "validate"]:
+                split_name = "validate"
+            candidates.append(f"{split_name}_data_roots.txt")
+        candidates.append("data_roots.txt")
+        candidates.append("train_data_roots.txt")
+        candidates.append("validate_data_roots.txt")
+
+        for fname in candidates:
+            path = os.path.join(self.img_path, fname)
+            if os.path.exists(path):
+                return _read_list_file(path)
+
+        img_dirs = sorted(
+            glob.glob(os.path.join(self.img_path, "**", "images"), recursive=True)
+        )
+        scene_roots = [os.path.dirname(p) for p in img_dirs]
+
+        rels = []
+        for p in scene_roots:
+            try:
+                rels.append(os.path.relpath(p, self.img_path))
+            except ValueError:
+                rels.append(p)
+        return sorted(set(rels))
+
+    def _resolve_seq_list_relpose(self) -> List[str]:
+        if self.seq_list is not None:
+            return list(self.seq_list)
+        meta = dataset_metadata.get(self.dataset, {})
+        if self.full_seq:
+            if self.img_path is None or not os.path.isdir(self.img_path):
+                return []
+            seqs = [
+                s
+                for s in os.listdir(self.img_path)
+                if os.path.isdir(os.path.join(self.img_path, s))
+            ]
+            return sorted(seqs)
+        seqs = meta.get("seq_list", []) or []
+        return list(seqs)
+
+    def _resolve_seq_list(self) -> List[str]:
+        fmt = self._infer_format()
+        if fmt == "generalizable":
+            return self._resolve_seq_list_generalizable()
+        return self._resolve_seq_list_relpose()
+
+    def _resolve_scene_root(self, seq_entry: str) -> Tuple[str, str]:
+        if os.path.isabs(seq_entry) or os.path.sep in seq_entry:
+            scene_root = seq_entry
+            name = os.path.basename(os.path.normpath(seq_entry))
+        else:
+            scene_root = os.path.join(self.img_path, seq_entry)
+            name = seq_entry
+        return name, scene_root
+
+    def _resolve_image_dir_generalizable(self, scene_root: str) -> Optional[str]:
+        images_root = os.path.join(scene_root, "images")
+        if not os.path.isdir(images_root):
+            return None
+
+        if isinstance(self.camera, str) and self.camera:
+            cam_dir = os.path.join(images_root, self.camera)
+            if os.path.isdir(cam_dir):
+                return cam_dir
+
+        if _direct_image_files(images_root):
+            return images_root
+
+        cams = [
+            d
+            for d in os.listdir(images_root)
+            if os.path.isdir(os.path.join(images_root, d))
+        ]
+        if not cams:
+            return None
+        cams = sorted(cams)
+
+        frame_dirs = []
+        for name in cams:
+            child_dir = os.path.join(images_root, name)
+            child_images = _direct_image_files(child_dir)
+            if child_images:
+                frame_dirs.append((name, len(child_images)))
+
+        if (
+            len(cams) > 10
+            and len(frame_dirs) == len(cams)
+            and max(count for _, count in frame_dirs) == 1
+        ):
+            return images_root
+
+        return os.path.join(images_root, cams[0])
+
+    def _camera_from_image_dir(self, image_dir: str) -> Optional[str]:
+        parent = os.path.basename(os.path.dirname(image_dir))
+        if parent != "images":
+            return None
+        return os.path.basename(image_dir)
+
+    def _collect_filelist(self, dir_path: str) -> List[str]:
+        filelist = _direct_image_files(dir_path)
+        if not filelist:
+            nested = []
+            child_dirs = sorted(
+                d for d in glob.glob(os.path.join(dir_path, "*")) if os.path.isdir(d)
+            )
+            for child_dir in child_dirs:
+                child_images = _direct_image_files(child_dir)
+                if child_images:
+                    nested.append(child_images[0])
+            filelist = nested
+        if self.stride > 1:
+            filelist = filelist[:: self.stride]
+        if self.max_frames is not None:
+            filelist = filelist[: self.max_frames]
+        return filelist
+
+    def _load_images(self, filelist: List[str]) -> torch.Tensor:
+        views = load_images_for_eval(
+            filelist,
+            size=self.size,
+            verbose=False,
+            crop=self.crop,
+            patch_size=self.patch_size,
+        )
+        imgs = torch.cat([view["img"] for view in views], dim=0)
+        images = imgs.unsqueeze(0)
+        images = (images + 1.0) / 2.0
+        return images
+
+    def iter_sequence_infos(self) -> Iterator[LongStreamSequenceInfo]:
+        fmt = self._infer_format()
+        seqs = self._resolve_seq_list()
+        for seq_entry in seqs:
+            if fmt == "generalizable":
+                seq, scene_root = self._resolve_scene_root(seq_entry)
+                dir_path = self._resolve_image_dir_generalizable(scene_root)
+                if dir_path is None or not os.path.isdir(dir_path):
+                    continue
+                camera = self._camera_from_image_dir(dir_path)
+            else:
+                seq = seq_entry
+                scene_root = os.path.join(self.img_path, seq)
+                dir_path = self.dir_path_func(self.img_path, seq)
+                if not os.path.isdir(dir_path):
+                    continue
+                camera = None
+
+            filelist = self._collect_filelist(dir_path)
+            if not filelist:
+                continue
+            yield LongStreamSequenceInfo(
+                name=seq,
+                scene_root=scene_root,
+                image_dir=dir_path,
+                image_paths=filelist,
+                camera=camera,
+            )
+
+    def __iter__(self) -> Iterator[LongStreamSequence]:
+        for info in self.iter_sequence_infos():
+            print(
+                f"[longstream] loading sequence {info.name}: {len(info.image_paths)} frames",
+                flush=True,
+            )
+            images = self._load_images(info.image_paths)
+            print(
+                f"[longstream] loaded sequence {info.name}: {tuple(images.shape)}",
+                flush=True,
+            )
+            yield LongStreamSequence(
+                info.name,
+                images,
+                info.image_paths,
+                scene_root=info.scene_root,
+                image_dir=info.image_dir,
+                camera=info.camera,
+            )

longstream/demo/__init__.py

@@ -0,0 +1,11 @@
+from .backend import create_demo_session, load_frame_previews
+from .common import BRANCH_OPTIONS, DISPLAY_MODE_OPTIONS, branch_key, load_metadata
+
+__all__ = [
+    "BRANCH_OPTIONS",
+    "DISPLAY_MODE_OPTIONS",
+    "branch_key",
+    "create_demo_session",
+    "load_frame_previews",
+    "load_metadata",
+]

longstream/demo/backend.py

@@ -0,0 +1,495 @@
+import json
+import os
+import re
+import shutil
+import tempfile
+from datetime import datetime
+from typing import Iterable, List, Optional, Tuple
+
+import cv2
+import numpy as np
+import torch
+import yaml
+
+from longstream.core.cli import default_config_path
+from longstream.core.model import LongStreamModel
+from longstream.streaming.keyframe_selector import KeyframeSelector
+from longstream.streaming.refresh import run_batch_refresh, run_streaming_refresh
+from longstream.utils.camera import compose_abs_from_rel
+from longstream.utils.depth import colorize_depth
+from longstream.utils.hub import resolve_checkpoint_path
+from longstream.utils.sky_mask import compute_sky_mask
+from longstream.utils.vendor.dust3r.utils.image import load_images_for_eval
+from longstream.utils.vendor.models.components.utils.pose_enc import (
+    pose_encoding_to_extri_intri,
+)
+
+from .common import load_metadata, session_file
+
+_IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".bmp", ".webp")
+_MODEL_CACHE = {}
+
+
+def _resolve_file_path(item) -> str:
+    if item is None:
+        return ""
+    if isinstance(item, str):
+        return item
+    if isinstance(item, dict) and "name" in item:
+        return item["name"]
+    if hasattr(item, "name"):
+        return item.name
+    return str(item)
+
+
+def _natural_sort_key(path: str):
+    name = os.path.basename(path)
+    stem, _ = os.path.splitext(name)
+    parts = re.split(r"(\d+)", stem)
+    key = []
+    for part in parts:
+        if not part:
+            continue
+        if part.isdigit():
+            key.append((0, int(part)))
+        else:
+            key.append((1, part.lower()))
+    return key, name.lower()
+
+
+def _sorted_image_paths(image_dir: str) -> List[str]:
+    files = []
+    for name in os.listdir(image_dir):
+        if name.lower().endswith(_IMAGE_EXTS):
+            files.append(os.path.join(image_dir, name))
+    return sorted(files, key=_natural_sort_key)
+
+
+def _session_root() -> str:
+    root = os.path.join(tempfile.gettempdir(), "longstream_demo_sessions")
+    os.makedirs(root, exist_ok=True)
+    return root
+
+
+def _new_session_dir() -> str:
+    stamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S_%f")
+    return tempfile.mkdtemp(prefix=f"longstream_{stamp}_", dir=_session_root())
+
+
+def _copy_uploaded_images(uploaded_files: Iterable, session_dir: str) -> List[str]:
+    input_dir = os.path.join(session_dir, "input_images")
+    os.makedirs(input_dir, exist_ok=True)
+    copied = []
+    sources = sorted(
+        (_resolve_file_path(x) for x in uploaded_files if x),
+        key=_natural_sort_key,
+    )
+    for src in sources:
+        if not src or not os.path.isfile(src):
+            continue
+        dst = os.path.join(input_dir, os.path.basename(src))
+        shutil.copy2(src, dst)
+        copied.append(dst)
+    return copied
+
+
+def _extract_uploaded_video(uploaded_video, session_dir: str) -> List[str]:
+    src = _resolve_file_path(uploaded_video)
+    if not src:
+        return []
+    if not os.path.isfile(src):
+        raise FileNotFoundError(src)
+
+    input_dir = os.path.join(session_dir, "input_images")
+    os.makedirs(input_dir, exist_ok=True)
+    cap = cv2.VideoCapture(src)
+    if not cap.isOpened():
+        raise ValueError(f"unable to open video: {src}")
+
+    image_paths = []
+    frame_id = 0
+    while True:
+        ok, frame = cap.read()
+        if not ok:
+            break
+        dst = os.path.join(input_dir, f"{frame_id:06d}.png")
+        if not cv2.imwrite(dst, frame):
+            cap.release()
+            raise ValueError(f"failed to write extracted frame: {dst}")
+        image_paths.append(dst)
+        frame_id += 1
+    cap.release()
+
+    if not image_paths:
+        raise ValueError(f"no frames extracted from video: {src}")
+    return image_paths
+
+
+def _resize_long_edge(arr, long_edge_size, interpolation):
+    h, w = arr.shape[:2]
+    scale = float(long_edge_size) / float(max(h, w))
+    new_w = int(round(w * scale))
+    new_h = int(round(h * scale))
+    return cv2.resize(arr, (new_w, new_h), interpolation=interpolation)
+
+
+def _prepare_mask_for_model(
+    mask, size, crop, patch_size, target_shape, square_ok=False
+):
+    if mask is None:
+        return None
+    h0, w0 = mask.shape[:2]
+    long_edge = round(size * max(w0 / h0, h0 / w0)) if size == 224 else size
+    mask = _resize_long_edge(mask, long_edge, cv2.INTER_NEAREST)
+
+    h, w = mask.shape[:2]
+    cx, cy = w // 2, h // 2
+    if size == 224:
+        half = min(cx, cy)
+        if crop:
+            mask = mask[cy - half : cy + half, cx - half : cx + half]
+        else:
+            mask = cv2.resize(
+                mask, (2 * half, 2 * half), interpolation=cv2.INTER_NEAREST
+            )
+    else:
+        halfw = ((2 * cx) // patch_size) * (patch_size // 2)
+        halfh = ((2 * cy) // patch_size) * (patch_size // 2)
+        if not square_ok and w == h:
+            halfh = int(3 * halfw / 4)
+        if crop:
+            mask = mask[cy - halfh : cy + halfh, cx - halfw : cx + halfw]
+        else:
+            mask = cv2.resize(
+                mask, (2 * halfw, 2 * halfh), interpolation=cv2.INTER_NEAREST
+            )
+
+    if mask.shape[:2] != tuple(target_shape):
+        mask = cv2.resize(
+            mask, (target_shape[1], target_shape[0]), interpolation=cv2.INTER_NEAREST
+        )
+    return mask.astype(np.uint8, copy=False)
+
+
+def _load_base_config(config_path: Optional[str] = None) -> dict:
+    path = config_path or default_config_path()
+    with open(path, "r") as f:
+        return yaml.safe_load(f) or {}
+
+
+def _resolve_demo_checkpoint(checkpoint: str) -> str:
+    local_candidates = []
+    for candidate in [checkpoint, os.getenv("LONGSTREAM_CHECKPOINT", "")]:
+        if isinstance(candidate, str) and candidate:
+            local_candidates.append(candidate)
+
+    for candidate in local_candidates:
+        if os.path.exists(candidate):
+            return os.path.abspath(candidate)
+
+    hf_cfg = {
+        "repo_id": os.getenv("LONGSTREAM_HF_REPO"),
+        "filename": os.getenv("LONGSTREAM_HF_FILE"),
+        "revision": os.getenv("LONGSTREAM_HF_REVISION"),
+        "local_dir": os.getenv("LONGSTREAM_HF_LOCAL_DIR", "checkpoints"),
+    }
+    resolved = resolve_checkpoint_path(None, hf_cfg)
+    if resolved and os.path.exists(resolved):
+        return os.path.abspath(resolved)
+
+    if hf_cfg["repo_id"] and hf_cfg["filename"]:
+        raise FileNotFoundError(
+            "checkpoint not found locally and Hugging Face resolution failed: "
+            f"repo_id={hf_cfg['repo_id']} filename={hf_cfg['filename']}"
+        )
+
+    searched = ", ".join(local_candidates) if local_candidates else "<none>"
+    raise FileNotFoundError(
+        "checkpoint not found. "
+        f"searched local paths: {searched}. "
+        "You can also set LONGSTREAM_HF_REPO and LONGSTREAM_HF_FILE."
+    )
+
+
+def _model_device(device: str) -> str:
+    if device == "cuda" and not torch.cuda.is_available():
+        return "cpu"
+    return device
+
+
+def _cache_key(checkpoint: str, device: str, model_cfg: dict) -> Tuple[str, str, str]:
+    rel_cfg = json.dumps(model_cfg.get("longstream_cfg", {}), sort_keys=True)
+    return checkpoint, device, rel_cfg
+
+
+def get_or_load_model(checkpoint: str, device: str, model_cfg: dict) -> LongStreamModel:
+    device = _model_device(device)
+    cfg = json.loads(json.dumps(model_cfg))
+    cfg["checkpoint"] = checkpoint
+    key = _cache_key(checkpoint, device, cfg)
+    model = _MODEL_CACHE.get(key)
+    if model is None:
+        model = LongStreamModel(cfg).to(device)
+        model.eval()
+        _MODEL_CACHE.clear()
+        _MODEL_CACHE[key] = model
+    return model
+
+
+def _load_images(
+    image_paths: List[str], size: int, crop: bool, patch_size: int
+) -> torch.Tensor:
+    views = load_images_for_eval(
+        image_paths, size=size, verbose=False, crop=crop, patch_size=patch_size
+    )
+    imgs = torch.cat([view["img"] for view in views], dim=0)
+    images = (imgs.unsqueeze(0) + 1.0) / 2.0
+    return images
+
+
+def _select_keyframes(images: torch.Tensor, keyframe_stride: int, keyframe_mode: str):
+    selector = KeyframeSelector(
+        min_interval=keyframe_stride,
+        max_interval=keyframe_stride,
+        force_first=True,
+        mode="random" if keyframe_mode == "random" else "fixed",
+    )
+    return selector.select_keyframes(images.shape[1], images.shape[0], images.device)
+
+
+def _run_model(images: torch.Tensor, model: LongStreamModel, infer_cfg: dict):
+    keyframe_stride = int(infer_cfg.get("keyframe_stride", 8))
+    keyframe_mode = infer_cfg.get("keyframe_mode", "fixed")
+    refresh = int(infer_cfg.get("refresh", 4))
+    mode = infer_cfg.get("mode", "streaming_refresh")
+    streaming_mode = infer_cfg.get("streaming_mode", "causal")
+    window_size = int(infer_cfg.get("window_size", 48))
+    rel_pose_cfg = infer_cfg.get("rel_pose_head_cfg", {"num_iterations": 4})
+
+    is_keyframe, keyframe_indices = _select_keyframes(
+        images, keyframe_stride, keyframe_mode
+    )
+    if mode == "batch_refresh":
+        outputs = run_batch_refresh(
+            model,
+            images,
+            is_keyframe,
+            keyframe_indices,
+            streaming_mode,
+            keyframe_stride,
+            refresh,
+            rel_pose_cfg,
+        )
+    elif mode == "streaming_refresh":
+        outputs = run_streaming_refresh(
+            model,
+            images,
+            is_keyframe,
+            keyframe_indices,
+            streaming_mode,
+            window_size,
+            refresh,
+            rel_pose_cfg,
+        )
+    else:
+        raise ValueError(f"Unsupported demo inference mode: {mode}")
+    return outputs, keyframe_indices
+
+
+def _compute_pose_outputs(
+    outputs: dict, keyframe_indices: torch.Tensor, image_hw: Tuple[int, int]
+):
+    if "rel_pose_enc" in outputs:
+        rel_pose_enc = outputs["rel_pose_enc"][0]
+        abs_pose_enc = compose_abs_from_rel(rel_pose_enc, keyframe_indices[0])
+        extri, intri = pose_encoding_to_extri_intri(
+            abs_pose_enc[None], image_size_hw=image_hw
+        )
+        return (
+            rel_pose_enc.detach().cpu().numpy(),
+            extri[0].detach().cpu().numpy(),
+            intri[0].detach().cpu().numpy(),
+        )
+    if "pose_enc" in outputs:
+        pose_enc = outputs["pose_enc"][0]
+        extri, intri = pose_encoding_to_extri_intri(
+            pose_enc[None], image_size_hw=image_hw
+        )
+        return None, extri[0].detach().cpu().numpy(), intri[0].detach().cpu().numpy()
+    raise RuntimeError("Model outputs contain neither rel_pose_enc nor pose_enc")
+
+
+def _compute_sky_masks(
+    image_paths: List[str],
+    target_shape: Tuple[int, int],
+    data_cfg: dict,
+    skyseg_path: str,
+    session_dir: str,
+):
+    raw_masks = compute_sky_mask(
+        image_paths, skyseg_path, os.path.join(session_dir, "sky_masks_raw")
+    )
+    if raw_masks is None:
+        return None
+    masks = []
+    for mask in raw_masks:
+        masks.append(
+            _prepare_mask_for_model(
+                mask,
+                size=int(data_cfg.get("size", 518)),
+                crop=bool(data_cfg.get("crop", False)),
+                patch_size=int(data_cfg.get("patch_size", 14)),
+                target_shape=target_shape,
+            )
+        )
+    return np.stack(masks, axis=0)
+
+
+def create_demo_session(
+    image_dir: str,
+    uploaded_files,
+    uploaded_video,
+    checkpoint: str,
+    device: str,
+    mode: str,
+    streaming_mode: str,
+    keyframe_stride: int,
+    refresh: int,
+    window_size: int,
+    compute_sky: bool,
+    config_path: Optional[str] = None,
+) -> str:
+    checkpoint = _resolve_demo_checkpoint(checkpoint)
+
+    session_dir = _new_session_dir()
+    base_cfg = _load_base_config(config_path)
+    data_cfg = dict(base_cfg.get("data", {}))
+    model_cfg = dict(base_cfg.get("model", {}))
+    infer_cfg = dict(base_cfg.get("inference", {}))
+
+    if image_dir:
+        image_dir = os.path.abspath(image_dir)
+        if not os.path.isdir(image_dir):
+            raise FileNotFoundError(f"image_dir not found: {image_dir}")
+        image_paths = _sorted_image_paths(image_dir)
+        input_root = image_dir
+    elif uploaded_video:
+        image_paths = _extract_uploaded_video(uploaded_video, session_dir)
+        input_root = _resolve_file_path(uploaded_video)
+    else:
+        image_paths = _copy_uploaded_images(uploaded_files or [], session_dir)
+        input_root = os.path.dirname(image_paths[0]) if image_paths else ""
+
+    if not image_paths:
+        raise ValueError("No input images found")
+
+    data_cfg["size"] = int(data_cfg.get("size", 518))
+    data_cfg["crop"] = bool(data_cfg.get("crop", False))
+    data_cfg["patch_size"] = int(data_cfg.get("patch_size", 14))
+
+    device = _model_device(device)
+    model = get_or_load_model(checkpoint, device, model_cfg)
+
+    images = _load_images(
+        image_paths, data_cfg["size"], data_cfg["crop"], data_cfg["patch_size"]
+    )
+    infer_cfg.update(
+        {
+            "mode": mode,
+            "streaming_mode": streaming_mode,
+            "keyframe_stride": int(keyframe_stride),
+            "refresh": int(refresh),
+            "window_size": int(window_size),
+        }
+    )
+
+    with torch.no_grad():
+        outputs, keyframe_indices = _run_model(images, model, infer_cfg)
+        h, w = images.shape[-2:]
+        rel_pose_enc, extri, intri = _compute_pose_outputs(
+            outputs, keyframe_indices, (h, w)
+        )
+        point_head = (
+            outputs["world_points"][0]
+            .detach()
+            .cpu()
+            .numpy()
+            .astype(np.float32, copy=False)
+        )
+        depth = (
+            outputs["depth"][0, :, :, :, 0]
+            .detach()
+            .cpu()
+            .numpy()
+            .astype(np.float32, copy=False)
+        )
+
+    if device == "cuda":
+        torch.cuda.empty_cache()
+
+    images_uint8 = np.clip(
+        images[0].permute(0, 2, 3, 1).cpu().numpy() * 255.0, 0, 255
+    ).astype(np.uint8)
+    sky_masks = None
+    if compute_sky:
+        skyseg_path = os.path.join(
+            os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "skyseg.onnx"
+        )
+        sky_masks = _compute_sky_masks(
+            image_paths, (h, w), data_cfg, skyseg_path, session_dir
+        )
+
+    np.save(session_file(session_dir, "images.npy"), images_uint8)
+    np.save(session_file(session_dir, "depth.npy"), depth)
+    np.save(session_file(session_dir, "point_head.npy"), point_head)
+    np.save(session_file(session_dir, "w2c.npy"), extri)
+    np.save(session_file(session_dir, "intri.npy"), intri)
+    if rel_pose_enc is not None:
+        np.save(
+            session_file(session_dir, "rel_pose_enc.npy"),
+            rel_pose_enc.astype(np.float32, copy=False),
+        )
+    if sky_masks is not None:
+        np.save(
+            session_file(session_dir, "sky_masks.npy"),
+            sky_masks.astype(np.uint8, copy=False),
+        )
+
+    sky_removed_ratio = None
+    if sky_masks is not None:
+        sky_removed_ratio = float(1.0 - (sky_masks > 0).mean())
+
+    metadata = {
+        "session_dir": session_dir,
+        "created_at": datetime.utcnow().isoformat() + "Z",
+        "checkpoint": os.path.abspath(checkpoint),
+        "device": device,
+        "mode": mode,
+        "streaming_mode": streaming_mode,
+        "keyframe_stride": int(keyframe_stride),
+        "refresh": int(refresh),
+        "window_size": int(window_size),
+        "num_frames": int(images_uint8.shape[0]),
+        "height": int(images_uint8.shape[1]),
+        "width": int(images_uint8.shape[2]),
+        "input_root": input_root,
+        "image_paths": image_paths,
+        "has_sky_masks": bool(sky_masks is not None),
+        "sky_removed_ratio": sky_removed_ratio,
+    }
+    with open(session_file(session_dir, "metadata.json"), "w") as f:
+        json.dump(metadata, f, indent=2)
+
+    del outputs
+    return session_dir
+
+
+def load_frame_previews(session_dir: str, frame_index: int):
+    meta = load_metadata(session_dir)
+    frame_index = int(np.clip(frame_index, 0, meta["num_frames"] - 1))
+    images = np.load(session_file(session_dir, "images.npy"), mmap_mode="r")
+    depth = np.load(session_file(session_dir, "depth.npy"), mmap_mode="r")
+    rgb = np.array(images[frame_index])
+    depth_color = colorize_depth(np.array(depth[frame_index]), cmap="plasma")
+    label = f"Frame {frame_index + 1}/{meta['num_frames']}"
+    return rgb, depth_color, label

longstream/demo/common.py

@@ -0,0 +1,84 @@
+import json
+import os
+from typing import List
+
+import numpy as np
+
+BRANCH_OPTIONS = [
+    "Point Head + Pose",
+    "Depth Projection + Pose",
+]
+BRANCH_TO_KEY = {
+    "Point Head + Pose": "point_head",
+    "Depth Projection + Pose": "depth_projection",
+}
+DISPLAY_MODE_OPTIONS = [
+    "Current Frame",
+    "Accumulate to Frame",
+    "All Frames",
+]
+
+
+def branch_key(label: str) -> str:
+    return BRANCH_TO_KEY.get(label, "point_head")
+
+
+def session_file(session_dir: str, name: str) -> str:
+    return os.path.join(session_dir, name)
+
+
+def load_metadata(session_dir: str) -> dict:
+    with open(session_file(session_dir, "metadata.json"), "r") as f:
+        return json.load(f)
+
+
+def selected_frame_indices(
+    num_frames: int, frame_index: int, display_mode: str
+) -> List[int]:
+    if num_frames <= 0:
+        return []
+    frame_index = int(np.clip(frame_index, 0, num_frames - 1))
+    if display_mode == "Current Frame":
+        return [frame_index]
+    if display_mode == "Accumulate to Frame":
+        return list(range(frame_index + 1))
+    return list(range(num_frames))
+
+
+def as_4x4(w2c):
+    w2c = np.asarray(w2c, dtype=np.float64)
+    if w2c.shape == (4, 4):
+        return w2c
+    out = np.eye(4, dtype=np.float64)
+    out[:3, :4] = w2c
+    return out
+
+
+_VIEW_ROT = np.array(
+    [
+        [1.0, 0.0, 0.0],
+        [0.0, 0.0, 1.0],
+        [0.0, -1.0, 0.0],
+    ],
+    dtype=np.float64,
+)
+
+
+def world_to_view(points):
+    points = np.asarray(points, dtype=np.float64)
+    return points @ _VIEW_ROT.T
+
+
+def camera_center_from_w2c(w2c):
+    c2w = np.linalg.inv(as_4x4(w2c))
+    return c2w[:3, 3]
+
+
+def c2w_in_view_space(w2c, origin_shift=None):
+    c2w = np.linalg.inv(as_4x4(w2c))
+    out = np.eye(4, dtype=np.float64)
+    out[:3, :3] = _VIEW_ROT @ c2w[:3, :3]
+    out[:3, 3] = world_to_view(c2w[:3, 3][None])[0]
+    if origin_shift is not None:
+        out[:3, 3] -= np.asarray(origin_shift, dtype=np.float64)
+    return out

longstream/demo/export.py

@@ -0,0 +1,85 @@
+import os
+
+import numpy as np
+
+from .geometry import camera_geometry, collect_points
+
+_CAMERA_COLORS = np.array(
+    [
+        [239, 68, 68, 255],
+        [14, 165, 233, 255],
+        [34, 197, 94, 255],
+        [245, 158, 11, 255],
+    ],
+    dtype=np.uint8,
+)
+
+
+def _camera_mesh(center, corners, color):
+    import trimesh
+
+    vertices = np.vstack([center[None], corners]).astype(np.float32)
+    faces = np.array(
+        [
+            [0, 1, 2],
+            [0, 2, 3],
+            [0, 3, 4],
+            [0, 4, 1],
+            [1, 2, 3],
+            [1, 3, 4],
+        ],
+        dtype=np.int64,
+    )
+    mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False)
+    mesh.visual.face_colors = np.tile(color[None], (faces.shape[0], 1))
+    return mesh
+
+
+def export_glb(
+    session_dir: str,
+    branch: str,
+    display_mode: str,
+    frame_index: int,
+    mask_sky: bool,
+    show_cameras: bool,
+    camera_scale: float,
+    max_points: int,
+) -> str:
+    import trimesh
+
+    points, colors, _ = collect_points(
+        session_dir=session_dir,
+        branch=branch,
+        display_mode=display_mode,
+        frame_index=frame_index,
+        mask_sky=mask_sky,
+        max_points=max_points,
+        seed=13,
+    )
+    if len(points) == 0:
+        raise ValueError("No valid points to export")
+
+    scene = trimesh.Scene()
+    scene.add_geometry(trimesh.PointCloud(vertices=points, colors=colors))
+
+    if show_cameras:
+        _, frustums, _ = camera_geometry(
+            session_dir=session_dir,
+            display_mode=display_mode,
+            frame_index=frame_index,
+            camera_scale_ratio=camera_scale,
+            points_hint=points,
+        )
+        for idx, (center, corners) in enumerate(frustums):
+            scene.add_geometry(
+                _camera_mesh(center, corners, _CAMERA_COLORS[idx % len(_CAMERA_COLORS)])
+            )
+
+    export_dir = os.path.join(session_dir, "exports")
+    os.makedirs(export_dir, exist_ok=True)
+    branch_slug = branch.lower().replace(" + ", "_").replace(" ", "_")
+    mode_slug = display_mode.replace(" ", "_").lower()
+    filename = f"{branch_slug}_{mode_slug}_{frame_index:04d}_sky{int(mask_sky)}_cam{int(show_cameras)}.glb"
+    path = os.path.join(export_dir, filename)
+    scene.export(path)
+    return path

longstream/demo/geometry.py

@@ -0,0 +1,211 @@
+import os
+from typing import List, Optional, Tuple
+
+import numpy as np
+
+from .common import (
+    branch_key,
+    c2w_in_view_space,
+    load_metadata,
+    selected_frame_indices,
+    session_file,
+    world_to_view,
+)
+
+
+def _origin_shift(w2c_all) -> np.ndarray:
+    first = c2w_in_view_space(w2c_all[0])
+    return first[:3, 3].copy()
+
+
+def _sample_flat_indices(
+    valid_indices: np.ndarray, budget: Optional[int], rng: np.random.Generator
+) -> np.ndarray:
+    if budget is None or budget <= 0 or valid_indices.size <= budget:
+        return valid_indices
+    keep = rng.choice(valid_indices.size, size=int(budget), replace=False)
+    return valid_indices[keep]
+
+
+def _depth_points_from_flat(depth, intri, w2c, flat_indices):
+    h, w = depth.shape
+    ys = flat_indices // w
+    xs = flat_indices % w
+    z = depth.reshape(-1)[flat_indices].astype(np.float64)
+    fx = float(intri[0, 0])
+    fy = float(intri[1, 1])
+    cx = float(intri[0, 2])
+    cy = float(intri[1, 2])
+    x = (xs.astype(np.float64) - cx) * z / max(fx, 1e-12)
+    y = (ys.astype(np.float64) - cy) * z / max(fy, 1e-12)
+    pts_cam = np.stack([x, y, z], axis=1)
+    R = w2c[:3, :3].astype(np.float64)
+    t = w2c[:3, 3].astype(np.float64)
+    return (R.T @ (pts_cam.T - t[:, None])).T.astype(np.float32, copy=False)
+
+
+def _camera_points_to_world(points, w2c):
+    pts = np.asarray(points, dtype=np.float64).reshape(-1, 3)
+    R = w2c[:3, :3].astype(np.float64)
+    t = w2c[:3, 3].astype(np.float64)
+    return (R.T @ (pts.T - t[:, None])).T.astype(np.float32, copy=False)
+
+
+def collect_points(
+    session_dir: str,
+    branch: str,
+    display_mode: str,
+    frame_index: int,
+    mask_sky: bool,
+    max_points: Optional[int],
+    seed: int = 0,
+):
+    branch = branch_key(branch)
+    meta = load_metadata(session_dir)
+    frame_ids = selected_frame_indices(meta["num_frames"], frame_index, display_mode)
+    if not frame_ids:
+        return (
+            np.empty((0, 3), dtype=np.float32),
+            np.empty((0, 3), dtype=np.uint8),
+            np.zeros(3, dtype=np.float64),
+        )
+
+    images = np.load(session_file(session_dir, "images.npy"), mmap_mode="r")
+    w2c = np.load(session_file(session_dir, "w2c.npy"), mmap_mode="r")
+    origin_shift = _origin_shift(w2c)
+    sky = None
+    if mask_sky and os.path.exists(session_file(session_dir, "sky_masks.npy")):
+        sky = np.load(session_file(session_dir, "sky_masks.npy"), mmap_mode="r")
+
+    if branch == "point_head":
+        point_head = np.load(session_file(session_dir, "point_head.npy"), mmap_mode="r")
+        source = point_head
+        depth = None
+        intri = None
+    else:
+        source = None
+        depth = np.load(session_file(session_dir, "depth.npy"), mmap_mode="r")
+        intri = np.load(session_file(session_dir, "intri.npy"), mmap_mode="r")
+
+    per_frame_budget = None
+    if max_points is not None and max_points > 0:
+        per_frame_budget = max(int(max_points) // max(len(frame_ids), 1), 1)
+
+    rng = np.random.default_rng(seed)
+    points = []
+    colors = []
+    for idx in frame_ids:
+        rgb_flat = images[idx].reshape(-1, 3)
+        if branch == "point_head":
+            pts_map = source[idx]
+            valid = np.isfinite(pts_map).all(axis=-1).reshape(-1)
+            if sky is not None:
+                valid &= sky[idx].reshape(-1) > 0
+            flat = np.flatnonzero(valid)
+            if flat.size == 0:
+                continue
+            flat = _sample_flat_indices(flat, per_frame_budget, rng)
+            pts_cam = pts_map.reshape(-1, 3)[flat]
+            pts_world = _camera_points_to_world(pts_cam, w2c[idx])
+        else:
+            depth_i = depth[idx]
+            valid = (np.isfinite(depth_i) & (depth_i > 0)).reshape(-1)
+            if sky is not None:
+                valid &= sky[idx].reshape(-1) > 0
+            flat = np.flatnonzero(valid)
+            if flat.size == 0:
+                continue
+            flat = _sample_flat_indices(flat, per_frame_budget, rng)
+            pts_world = _depth_points_from_flat(depth_i, intri[idx], w2c[idx], flat)
+
+        pts_view = world_to_view(pts_world) - origin_shift[None]
+        points.append(pts_view.astype(np.float32, copy=False))
+        colors.append(rgb_flat[flat].astype(np.uint8, copy=False))
+
+    if not points:
+        return (
+            np.empty((0, 3), dtype=np.float32),
+            np.empty((0, 3), dtype=np.uint8),
+            origin_shift,
+        )
+    return np.concatenate(points, axis=0), np.concatenate(colors, axis=0), origin_shift
+
+
+def _frustum_corners_camera(intri, image_hw, depth_scale):
+    h, w = image_hw
+    fx = float(intri[0, 0])
+    fy = float(intri[1, 1])
+    cx = float(intri[0, 2])
+    cy = float(intri[1, 2])
+    corners = np.array(
+        [
+            [
+                (0.0 - cx) * depth_scale / max(fx, 1e-12),
+                (0.0 - cy) * depth_scale / max(fy, 1e-12),
+                depth_scale,
+            ],
+            [
+                ((w - 1.0) - cx) * depth_scale / max(fx, 1e-12),
+                (0.0 - cy) * depth_scale / max(fy, 1e-12),
+                depth_scale,
+            ],
+            [
+                ((w - 1.0) - cx) * depth_scale / max(fx, 1e-12),
+                ((h - 1.0) - cy) * depth_scale / max(fy, 1e-12),
+                depth_scale,
+            ],
+            [
+                (0.0 - cx) * depth_scale / max(fx, 1e-12),
+                ((h - 1.0) - cy) * depth_scale / max(fy, 1e-12),
+                depth_scale,
+            ],
+        ],
+        dtype=np.float64,
+    )
+    return corners
+
+
+def camera_geometry(
+    session_dir: str,
+    display_mode: str,
+    frame_index: int,
+    camera_scale_ratio: float,
+    points_hint=None,
+):
+    meta = load_metadata(session_dir)
+    frame_ids = selected_frame_indices(meta["num_frames"], frame_index, display_mode)
+    w2c = np.load(session_file(session_dir, "w2c.npy"), mmap_mode="r")
+    intri = np.load(session_file(session_dir, "intri.npy"), mmap_mode="r")
+    origin_shift = _origin_shift(w2c)
+
+    center_points = np.array(
+        [c2w_in_view_space(w2c[idx], origin_shift)[:3, 3] for idx in frame_ids],
+        dtype=np.float64,
+    )
+    center_extent = 1.0
+    if len(center_points) > 1:
+        center_extent = float(
+            np.linalg.norm(center_points.max(axis=0) - center_points.min(axis=0))
+        )
+
+    point_extent = 0.0
+    if points_hint is not None and len(points_hint) > 0:
+        lo = np.percentile(points_hint, 5, axis=0)
+        hi = np.percentile(points_hint, 95, axis=0)
+        point_extent = float(np.linalg.norm(hi - lo))
+
+    extent = max(center_extent, point_extent, 1.0)
+    depth_scale = extent * float(camera_scale_ratio)
+
+    centers = []
+    frustums = []
+    for idx in frame_ids:
+        c2w_view = c2w_in_view_space(w2c[idx], origin_shift)
+        center = c2w_view[:3, 3]
+        corners_cam = _frustum_corners_camera(
+            intri[idx], (meta["height"], meta["width"]), depth_scale
+        )
+        corners_world = (c2w_view[:3, :3] @ corners_cam.T).T + center[None]
+        centers.append(center)
+        frustums.append((center, corners_world))
+    return np.asarray(centers, dtype=np.float64), frustums, origin_shift

longstream/demo/viewer.py

@@ -0,0 +1,134 @@
+import numpy as np
+import plotly.graph_objects as go
+
+from longstream.demo.backend import load_frame_previews
+
+from .common import load_metadata
+from .geometry import camera_geometry, collect_points
+
+
+def _empty_figure(message: str):
+    fig = go.Figure()
+    fig.add_annotation(
+        text=message, x=0.5, y=0.5, xref="paper", yref="paper", showarrow=False
+    )
+    fig.update_layout(
+        template="plotly_white",
+        margin=dict(l=0, r=0, t=40, b=0),
+        scene=dict(aspectmode="data"),
+    )
+    return fig
+
+
+def _camera_lines(frustums):
+    xs, ys, zs = [], [], []
+    for center, corners in frustums:
+        order = [(0, 1), (1, 2), (2, 3), (3, 0)]
+        for a, b in order:
+            xs.extend([corners[a, 0], corners[b, 0], None])
+            ys.extend([corners[a, 1], corners[b, 1], None])
+            zs.extend([corners[a, 2], corners[b, 2], None])
+        for corner in corners:
+            xs.extend([center[0], corner[0], None])
+            ys.extend([center[1], corner[1], None])
+            zs.extend([center[2], corner[2], None])
+    return xs, ys, zs
+
+
+def build_interactive_figure(
+    session_dir: str,
+    branch: str,
+    display_mode: str,
+    frame_index: int,
+    point_size: float,
+    opacity: float,
+    preview_max_points: int,
+    show_cameras: bool,
+    camera_scale: float,
+    mask_sky: bool,
+):
+    meta = load_metadata(session_dir)
+    points, colors, _ = collect_points(
+        session_dir=session_dir,
+        branch=branch,
+        display_mode=display_mode,
+        frame_index=frame_index,
+        mask_sky=mask_sky,
+        max_points=preview_max_points,
+        seed=frame_index,
+    )
+    if len(points) == 0:
+        return _empty_figure("No valid points for the current selection")
+
+    fig = go.Figure()
+    fig.add_trace(
+        go.Scatter3d(
+            x=points[:, 0],
+            y=points[:, 1],
+            z=points[:, 2],
+            mode="markers",
+            marker=dict(
+                size=float(point_size),
+                color=[f"rgb({r},{g},{b})" for r, g, b in colors],
+                opacity=float(opacity),
+            ),
+            hoverinfo="skip",
+            name="points",
+        )
+    )
+
+    if show_cameras:
+        centers, frustums, _ = camera_geometry(
+            session_dir=session_dir,
+            display_mode=display_mode,
+            frame_index=frame_index,
+            camera_scale_ratio=camera_scale,
+            points_hint=points,
+        )
+        if len(centers) > 0:
+            fig.add_trace(
+                go.Scatter3d(
+                    x=centers[:, 0],
+                    y=centers[:, 1],
+                    z=centers[:, 2],
+                    mode="lines",
+                    line=dict(color="#16a34a", width=2),
+                    name="trajectory",
+                    hoverinfo="skip",
+                )
+            )
+            xs, ys, zs = _camera_lines(frustums)
+            fig.add_trace(
+                go.Scatter3d(
+                    x=xs,
+                    y=ys,
+                    z=zs,
+                    mode="lines",
+                    line=dict(color="#22c55e", width=1.5),
+                    name="cameras",
+                    hoverinfo="skip",
+                )
+            )
+
+    fig.update_layout(
+        template="plotly_white",
+        margin=dict(l=0, r=0, t=40, b=0),
+        scene=dict(
+            aspectmode="data",
+            xaxis_title="x_right",
+            yaxis_title="z_forward",
+            zaxis_title="y_up",
+            bgcolor="#f8fafc",
+            camera=dict(
+                up=dict(x=0.0, y=0.0, z=1.0),
+                eye=dict(x=-1.0, y=-1.8, z=0.9),
+            ),
+        ),
+        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="left", x=0.0),
+    )
+    return fig
+
+
+def build_frame_outputs(session_dir: str, frame_index: int):
+    rgb, depth, label = load_frame_previews(session_dir, frame_index)
+    return rgb, depth, label

longstream/eval/__init__.py

@@ -0,0 +1,3 @@
+from .evaluate import evaluate_predictions_cfg
+
+__all__ = ["evaluate_predictions_cfg"]

longstream/eval/evaluate.py

@@ -0,0 +1,551 @@
+import json
+import os
+
+import cv2
+import numpy as np
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+from longstream.data import LongStreamDataLoader
+from longstream.eval.io import (
+    frame_stems,
+    read_depth,
+    read_opencv_camera_yml,
+    read_pointcloud_xyz,
+    read_pred_w2c_txt,
+)
+from longstream.eval.metrics import ate_rmse, chamfer_and_f1, transform_points
+from longstream.utils.sky_mask import sky_mask_filename
+
+
+def _ensure_dir(path):
+    os.makedirs(path, exist_ok=True)
+
+
+def _sequence_output_dir(output_root, seq_name):
+    return os.path.join(output_root, seq_name)
+
+
+def _sequence_metrics_path(output_root, seq_name):
+    return os.path.join(output_root, "metrics", f"{seq_name}.json")
+
+
+def _sequence_plot_path(output_root, seq_name):
+    return os.path.join(output_root, "plots", f"{seq_name}_traj_3d.png")
+
+
+def _world_xyz_to_plot_xyz(xyz):
+    xyz = np.asarray(xyz, dtype=np.float64)
+    return np.stack([xyz[:, 0], xyz[:, 2], -xyz[:, 1]], axis=-1)
+
+
+def _set_equal_3d_axes(ax, xyz):
+    mins = xyz.min(axis=0)
+    maxs = xyz.max(axis=0)
+    center = 0.5 * (mins + maxs)
+    radius = 0.5 * np.max(np.maximum(maxs - mins, 1e-6))
+    ax.set_xlim(center[0] - radius, center[0] + radius)
+    ax.set_ylim(center[1] - radius, center[1] + radius)
+    ax.set_zlim(center[2] - radius, center[2] + radius)
+
+
+def _load_gt_pose_data(seq_info):
+    if seq_info.camera is not None:
+        cam_dir = os.path.join(seq_info.scene_root, "cameras", seq_info.camera)
+        extri_path = os.path.join(cam_dir, "extri.yml")
+        intri_path = os.path.join(cam_dir, "intri.yml")
+        if os.path.exists(extri_path):
+            extri, intri, image_sizes = read_opencv_camera_yml(extri_path, intri_path)
+            return extri, intri, image_sizes
+
+    extri_path = os.path.join(seq_info.scene_root, "extri.yml")
+    intri_path = os.path.join(seq_info.scene_root, "intri.yml")
+    if not os.path.exists(extri_path):
+        return None, None, None
+    extri, intri, image_sizes = read_opencv_camera_yml(extri_path, intri_path)
+    return extri, intri, image_sizes
+
+
+def _resolve_gt_depth_root(seq_info):
+    if seq_info.camera is not None:
+        camera_depth_root = os.path.join(seq_info.scene_root, "depths", seq_info.camera)
+        if os.path.isdir(camera_depth_root):
+            return camera_depth_root
+    depth_root = os.path.join(seq_info.scene_root, "depths")
+    if os.path.isdir(depth_root):
+        return depth_root
+    return None
+
+
+def _resolve_gt_depth_path(seq_info, depth_root, image_path, stem):
+    rel_path = os.path.relpath(image_path, seq_info.image_dir)
+    rel_stem = os.path.splitext(rel_path)[0]
+    file_stem = os.path.splitext(os.path.basename(image_path))[0]
+    candidates = [
+        os.path.join(depth_root, f"{stem}.exr"),
+        os.path.join(depth_root, rel_stem + ".exr"),
+        os.path.join(depth_root, stem, f"{file_stem}.exr"),
+    ]
+    for candidate in candidates:
+        if os.path.exists(candidate):
+            return candidate
+    return None
+
+
+def _resize_long_edge(arr, long_edge_size, interpolation):
+    h, w = arr.shape[:2]
+    scale = float(long_edge_size) / float(max(h, w))
+    new_w = int(round(w * scale))
+    new_h = int(round(h * scale))
+    return cv2.resize(arr, (new_w, new_h), interpolation=interpolation)
+
+
+def _prepare_map_for_eval(
+    arr, size, crop, patch_size, target_shape, interpolation, square_ok=False
+):
+    h0, w0 = arr.shape[:2]
+    long_edge = round(size * max(w0 / h0, h0 / w0)) if size == 224 else size
+    arr = _resize_long_edge(arr, long_edge, interpolation)
+
+    h, w = arr.shape[:2]
+    cx, cy = w // 2, h // 2
+
+    if size == 224:
+        half = min(cx, cy)
+        target_w = 2 * half
+        target_h = 2 * half
+        if crop:
+            arr = arr[cy - half : cy + half, cx - half : cx + half]
+        else:
+            arr = cv2.resize(arr, (target_w, target_h), interpolation=interpolation)
+    else:
+        halfw = ((2 * cx) // patch_size) * (patch_size // 2)
+        halfh = ((2 * cy) // patch_size) * (patch_size // 2)
+        if not square_ok and w == h:
+            halfh = int(3 * halfw / 4)
+        target_w = 2 * halfw
+        target_h = 2 * halfh
+        if crop:
+            arr = arr[cy - halfh : cy + halfh, cx - halfw : cx + halfw]
+        else:
+            arr = cv2.resize(arr, (target_w, target_h), interpolation=interpolation)
+
+    if arr.shape[:2] != tuple(target_shape):
+        arr = cv2.resize(
+            arr, (target_shape[1], target_shape[0]), interpolation=interpolation
+        )
+    return arr
+
+
+def _sky_mask_path(seq_dir, image_path):
+    return os.path.join(seq_dir, "sky_masks", sky_mask_filename(image_path))
+
+
+def _sample_frame_points(points, max_points, rng):
+    if max_points is None or len(points) <= max_points:
+        return points
+    keep = rng.choice(len(points), size=max_points, replace=False)
+    return points[keep]
+
+
+def _depth_to_world_points(depth, intri, extri, valid_mask):
+    ys, xs = np.nonzero(valid_mask)
+    if ys.size == 0:
+        return np.empty((0, 3), dtype=np.float32)
+
+    z = depth[ys, xs].astype(np.float64)
+    fx = float(intri[0, 0])
+    fy = float(intri[1, 1])
+    cx = float(intri[0, 2])
+    cy = float(intri[1, 2])
+
+    x = (xs.astype(np.float64) - cx) * z / max(fx, 1e-12)
+    y = (ys.astype(np.float64) - cy) * z / max(fy, 1e-12)
+    pts_cam = np.stack([x, y, z], axis=1)
+
+    R = extri[:3, :3]
+    t = extri[:3, 3]
+    pts_world = (R.T @ (pts_cam.T - t[:, None])).T
+    return pts_world.astype(np.float32, copy=False)
+
+
+def _load_gt_pointcloud(seq_info, seq_dir, gt_extri, gt_intri, eval_cfg):
+    if not gt_extri or not gt_intri:
+        return None
+
+    gt_dir = _resolve_gt_depth_root(seq_info)
+    if gt_dir is None:
+        return None
+
+    eval_max_points = int(eval_cfg.get("point_eval_max_points", 100000))
+    oversample_factor = int(eval_cfg.get("point_eval_oversample_factor", 4))
+    per_frame_budget = max(
+        (eval_max_points * oversample_factor) // max(len(seq_info.image_paths), 1), 1
+    )
+    rng = np.random.default_rng(0)
+    chunks = []
+
+    for image_path, stem in zip(
+        seq_info.image_paths, frame_stems(seq_info.image_paths)
+    ):
+        depth_path = _resolve_gt_depth_path(seq_info, gt_dir, image_path, stem)
+        if depth_path is None or stem not in gt_extri or stem not in gt_intri:
+            continue
+
+        depth = read_depth(depth_path)
+        valid = np.isfinite(depth) & (depth > 0)
+        if not np.any(valid):
+            continue
+
+        sky_path = _sky_mask_path(seq_dir, image_path)
+        if os.path.exists(sky_path):
+            sky_mask = cv2.imread(sky_path, cv2.IMREAD_GRAYSCALE)
+            if sky_mask is not None:
+                if sky_mask.shape[:2] != depth.shape[:2]:
+                    sky_mask = cv2.resize(
+                        sky_mask,
+                        (depth.shape[1], depth.shape[0]),
+                        interpolation=cv2.INTER_NEAREST,
+                    )
+                valid &= sky_mask > 0
+        if not np.any(valid):
+            continue
+
+        pts_world = _depth_to_world_points(depth, gt_intri[stem], gt_extri[stem], valid)
+        if len(pts_world) == 0:
+            continue
+        chunks.append(_sample_frame_points(pts_world, per_frame_budget, rng))
+
+    if not chunks:
+        return None
+    return np.concatenate(chunks, axis=0)
+
+
+def _evaluate_pointclouds(seq_info, seq_dir, eval_cfg, pose_align, gt_cloud):
+    if pose_align is None or gt_cloud is None:
+        return None
+
+    scale, R, t = pose_align
+    point_paths = {
+        "point_head": [
+            os.path.join(seq_dir, "points", "point_head_full.npy"),
+            os.path.join(seq_dir, "points", "point_head_full.npz"),
+            os.path.join(seq_dir, "points", "point_head_full.ply"),
+        ],
+        "dpt_unproj": [
+            os.path.join(seq_dir, "points", "dpt_unproj_full.npy"),
+            os.path.join(seq_dir, "points", "dpt_unproj_full.npz"),
+            os.path.join(seq_dir, "points", "dpt_unproj_full.ply"),
+        ],
+    }
+    threshold = float(eval_cfg.get("point_f1_threshold", 0.25))
+    max_points = int(eval_cfg.get("point_eval_max_points", 100000))
+    voxel_size = eval_cfg.get("point_eval_voxel_size", None)
+    voxel_size = None if voxel_size in (None, "", "null") else float(voxel_size)
+
+    metrics_by_branch = {}
+    for branch, candidates in point_paths.items():
+        path = next(
+            (candidate for candidate in candidates if os.path.exists(candidate)), None
+        )
+        if path is None:
+            continue
+        pred_cloud = read_pointcloud_xyz(path)
+        pred_cloud = transform_points(pred_cloud, scale, R, t)
+        metrics = chamfer_and_f1(
+            pred_cloud,
+            gt_cloud,
+            threshold=threshold,
+            max_points=max_points,
+            voxel_size=voxel_size,
+            seed=0 if branch == "point_head" else 1,
+        )
+        if metrics is not None:
+            metrics_by_branch[branch] = metrics
+    return metrics_by_branch or None
+
+
+def _evaluate_video_dpt(seq_info, seq_dir, eval_cfg, data_cfg):
+    pred_dir = os.path.join(seq_dir, "depth", "dpt")
+    gt_dir = _resolve_gt_depth_root(seq_info)
+    if not os.path.isdir(pred_dir) or gt_dir is None:
+        return None
+
+    size = int(data_cfg.get("size", 518))
+    crop = bool(data_cfg.get("crop", False))
+    patch_size = int(data_cfg.get("patch_size", 14))
+    rel_delta_threshold = float(eval_cfg.get("depth_rel_delta_threshold", 1.25))
+
+    abs_rel_sum = 0.0
+    rel_delta_hits = 0
+    valid_pixels = 0
+    evaluated_frames = 0
+
+    stems = frame_stems(seq_info.image_paths)
+    for frame_id, stem in enumerate(stems):
+        pred_path = os.path.join(pred_dir, f"frame_{frame_id:06d}.npy")
+        gt_path = _resolve_gt_depth_path(
+            seq_info, gt_dir, seq_info.image_paths[frame_id], stem
+        )
+        if not os.path.exists(pred_path) or gt_path is None:
+            continue
+
+        pred = np.load(pred_path).astype(np.float32)
+        gt = read_depth(gt_path)
+        gt = _prepare_map_for_eval(
+            gt,
+            size=size,
+            crop=crop,
+            patch_size=patch_size,
+            target_shape=pred.shape,
+            interpolation=cv2.INTER_NEAREST,
+        )
+
+        valid = np.isfinite(gt) & (gt > 0)
+        if not np.any(valid):
+            continue
+
+        sky_mask_path = _sky_mask_path(seq_dir, seq_info.image_paths[frame_id])
+        if os.path.exists(sky_mask_path):
+            sky_mask = cv2.imread(sky_mask_path, cv2.IMREAD_GRAYSCALE)
+            if sky_mask is not None:
+                sky_mask = _prepare_map_for_eval(
+                    sky_mask,
+                    size=size,
+                    crop=crop,
+                    patch_size=patch_size,
+                    target_shape=pred.shape,
+                    interpolation=cv2.INTER_NEAREST,
+                )
+                valid &= sky_mask > 0
+
+        valid &= np.isfinite(pred)
+        if not np.any(valid):
+            continue
+
+        pred_valid = pred[valid].astype(np.float64)
+        gt_valid = gt[valid].astype(np.float64)
+        pred_safe = np.clip(pred_valid, 1e-6, None)
+        gt_safe = np.clip(gt_valid, 1e-6, None)
+
+        abs_rel_sum += np.sum(np.abs(pred_valid - gt_valid) / gt_safe)
+        rel_ratio = np.maximum(gt_safe / pred_safe, pred_safe / gt_safe)
+        rel_delta_hits += int(np.sum(rel_ratio < rel_delta_threshold))
+        valid_pixels += int(gt_valid.size)
+        evaluated_frames += 1
+
+    if valid_pixels == 0:
+        return None
+
+    return {
+        "abs_rel": float(abs_rel_sum / valid_pixels),
+        "rel_delta": float(rel_delta_hits / valid_pixels),
+        "rel_delta_threshold": rel_delta_threshold,
+        "num_valid_pixels": int(valid_pixels),
+        "num_frames": int(evaluated_frames),
+    }
+
+
+def _extract_pose_pairs(seq_info, pred_pose_path, gt_extri):
+    frame_ids, pred_w2c = read_pred_w2c_txt(pred_pose_path)
+    if not pred_w2c:
+        return None
+
+    stems = frame_stems(seq_info.image_paths)
+    pred_xyz = []
+    gt_xyz = []
+
+    for frame_id, pred_mat in zip(frame_ids, pred_w2c):
+        if frame_id < 0 or frame_id >= len(stems):
+            continue
+        stem = stems[frame_id]
+        if stem not in gt_extri:
+            continue
+        pred_c2w = np.linalg.inv(pred_mat)
+        gt_c2w = np.linalg.inv(gt_extri[stem])
+        pred_xyz.append(pred_c2w[:3, 3])
+        gt_xyz.append(gt_c2w[:3, 3])
+
+    if len(pred_xyz) < 3:
+        return None
+    return np.asarray(pred_xyz, dtype=np.float64), np.asarray(gt_xyz, dtype=np.float64)
+
+
+def _save_traj_plot_3d(path, pred_xyz, gt_xyz):
+    _ensure_dir(os.path.dirname(path))
+    pred_plot = _world_xyz_to_plot_xyz(pred_xyz)
+    gt_plot = _world_xyz_to_plot_xyz(gt_xyz)
+    origin = gt_plot[:1]
+    pred_plot = pred_plot - origin
+    gt_plot = gt_plot - origin
+    all_plot = np.concatenate([pred_plot, gt_plot], axis=0)
+
+    fig = plt.figure(figsize=(7, 6))
+    ax = fig.add_subplot(111, projection="3d")
+    ax.plot(
+        gt_plot[:, 0],
+        gt_plot[:, 1],
+        gt_plot[:, 2],
+        label="gt",
+        linewidth=2.0,
+        color="#1f77b4",
+    )
+    ax.plot(
+        pred_plot[:, 0],
+        pred_plot[:, 1],
+        pred_plot[:, 2],
+        label="pred",
+        linewidth=2.0,
+        color="#d62728",
+    )
+    _set_equal_3d_axes(ax, all_plot)
+    ax.view_init(elev=24, azim=-118)
+    ax.set_xlabel("x_right")
+    ax.set_ylabel("z_forward")
+    ax.set_zlabel("y_up")
+    ax.legend(loc="best")
+    ax.set_title("Trajectory 3D (Sim3-aligned view)")
+    fig.tight_layout()
+    fig.savefig(path, dpi=180)
+    plt.close(fig)
+
+
+def evaluate_sequence(seq_info, output_root, eval_cfg, data_cfg):
+    seq_dir = _sequence_output_dir(output_root, seq_info.name)
+    result = {
+        "sequence": seq_info.name,
+        "output_dir": seq_dir,
+        "has_gt": False,
+        "has_gt_pose": False,
+        "has_gt_depth": False,
+    }
+
+    gt_extri, gt_intri, _ = _load_gt_pose_data(seq_info)
+    pose_align = None
+    if gt_extri:
+        result["has_gt"] = True
+        result["has_gt_pose"] = True
+
+        pred_pose_path = os.path.join(seq_dir, "poses", "abs_pose.txt")
+        pairs = _extract_pose_pairs(seq_info, pred_pose_path, gt_extri)
+        if pairs is not None:
+            pred_xyz, gt_xyz = pairs
+            pose_metrics = ate_rmse(
+                pred_xyz, gt_xyz, align_scale=bool(eval_cfg.get("align_scale", True))
+            )
+            sim3_scale = float(pose_metrics.get("sim3_scale", 1.0))
+            pred_xyz_aligned = transform_points(
+                pred_xyz,
+                sim3_scale,
+                np.asarray(pose_metrics["sim3_rotation"], dtype=np.float64),
+                np.asarray(pose_metrics["sim3_translation"], dtype=np.float64),
+            )
+            pose_align = (
+                sim3_scale,
+                np.asarray(pose_metrics["sim3_rotation"], dtype=np.float64),
+                np.asarray(pose_metrics["sim3_translation"], dtype=np.float64),
+            )
+            plot_path = _sequence_plot_path(output_root, seq_info.name)
+            _save_traj_plot_3d(plot_path, pred_xyz_aligned, gt_xyz)
+            pose_metrics.pop("sim3_scale", None)
+            pose_metrics["traj_3d_plot"] = plot_path
+            result["pose"] = pose_metrics
+
+    video_dpt_metrics = _evaluate_video_dpt(seq_info, seq_dir, eval_cfg, data_cfg)
+    if video_dpt_metrics is not None:
+        result["has_gt"] = True
+        result["has_gt_depth"] = True
+        result["video_dpt"] = video_dpt_metrics
+
+    gt_cloud = _load_gt_pointcloud(seq_info, seq_dir, gt_extri, gt_intri, eval_cfg)
+    pointcloud_metrics = _evaluate_pointclouds(
+        seq_info, seq_dir, eval_cfg, pose_align, gt_cloud
+    )
+    if pointcloud_metrics is not None:
+        result["has_gt"] = True
+        result["has_gt_depth"] = True
+        result["pointcloud"] = pointcloud_metrics
+
+    if not result["has_gt"]:
+        result["skipped"] = "missing_gt"
+
+    return result
+
+
+def _mean_metric(sequence_results, group_name, metric_name):
+    values = []
+    for item in sequence_results:
+        group = item
+        for key in group_name.split("."):
+            if not isinstance(group, dict):
+                group = None
+                break
+            group = group.get(key)
+        if not isinstance(group, dict):
+            continue
+        if metric_name in group:
+            values.append(float(group[metric_name]))
+    if not values:
+        return None
+    return float(np.mean(values))
+
+
+def evaluate_predictions_cfg(cfg):
+    data_cfg = dict(cfg.get("data", {}))
+    data_cfg["format"] = "generalizable"
+    output_cfg = cfg.get("output", {})
+    eval_cfg = cfg.get("evaluation", {})
+    output_root = output_cfg.get("root", "outputs")
+    _ensure_dir(output_root)
+
+    loader = LongStreamDataLoader(data_cfg)
+    sequence_results = []
+    for seq_info in loader.iter_sequence_infos():
+        print(f"[longstream] eval {seq_info.name}: start", flush=True)
+        metrics = evaluate_sequence(seq_info, output_root, eval_cfg, data_cfg)
+        sequence_results.append(metrics)
+        metrics_path = _sequence_metrics_path(output_root, seq_info.name)
+        _ensure_dir(os.path.dirname(metrics_path))
+        with open(metrics_path, "w") as f:
+            json.dump(metrics, f, indent=2)
+        print(f"[longstream] eval {seq_info.name}: wrote {metrics_path}", flush=True)
+
+    summary = {
+        "num_sequences": len(sequence_results),
+        "num_sequences_with_gt": sum(1 for x in sequence_results if x.get("has_gt")),
+        "num_sequences_with_pose_gt": sum(
+            1 for x in sequence_results if x.get("has_gt_pose")
+        ),
+        "num_sequences_with_depth_gt": sum(
+            1 for x in sequence_results if x.get("has_gt_depth")
+        ),
+        "ate_mean": _mean_metric(sequence_results, "pose", "ate_mean"),
+        "ate_rmse_mean": _mean_metric(sequence_results, "pose", "ate_rmse"),
+        "video_dpt_abs_rel_mean": _mean_metric(
+            sequence_results, "video_dpt", "abs_rel"
+        ),
+        "video_dpt_rel_delta_mean": _mean_metric(
+            sequence_results, "video_dpt", "rel_delta"
+        ),
+        "point_head_cd_mean": _mean_metric(
+            sequence_results, "pointcloud.point_head", "cd"
+        ),
+        "point_head_f1_mean": _mean_metric(
+            sequence_results, "pointcloud.point_head", "f1"
+        ),
+        "dpt_unproj_cd_mean": _mean_metric(
+            sequence_results, "pointcloud.dpt_unproj", "cd"
+        ),
+        "dpt_unproj_f1_mean": _mean_metric(
+            sequence_results, "pointcloud.dpt_unproj", "f1"
+        ),
+        "sequences": sequence_results,
+    }
+
+    summary_path = os.path.join(output_root, "summary.json")
+    with open(summary_path, "w") as f:
+        json.dump(summary, f, indent=2)
+    print(f"[longstream] eval: wrote {summary_path}", flush=True)
+    return summary

longstream/eval/io.py

@@ -0,0 +1,156 @@
+import os
+
+os.environ.setdefault("OPENCV_IO_ENABLE_OPENEXR", "1")
+
+import cv2
+import numpy as np
+
+
+def frame_stems(image_paths):
+    stems = [os.path.splitext(os.path.basename(p))[0] for p in image_paths]
+    if len(set(stems)) == len(stems):
+        return stems
+    parents = [os.path.basename(os.path.dirname(p)) for p in image_paths]
+    if len(set(parents)) == len(parents):
+        return parents
+    return stems
+
+
+def read_pred_w2c_txt(path):
+    frames = []
+    poses = []
+    if not os.path.exists(path):
+        return frames, poses
+    with open(path, "r") as f:
+        for line in f:
+            line = line.strip()
+            if not line or line.startswith("#"):
+                continue
+            vals = [float(x) for x in line.split()]
+            if len(vals) != 13:
+                continue
+            frame = int(vals[0])
+            mat = np.eye(4, dtype=np.float64)
+            mat[:3, :3] = np.asarray(vals[1:10], dtype=np.float64).reshape(3, 3)
+            mat[:3, 3] = np.asarray(vals[10:13], dtype=np.float64)
+            frames.append(frame)
+            poses.append(mat)
+    return frames, poses
+
+
+def read_opencv_camera_yml(extri_path, intri_path=None):
+    if not os.path.exists(extri_path):
+        return {}, {}, {}
+
+    fs_extri = cv2.FileStorage(extri_path, cv2.FILE_STORAGE_READ)
+    names_node = fs_extri.getNode("names")
+    names = []
+    for i in range(names_node.size()):
+        names.append(names_node.at(i).string())
+
+    extri = {}
+    for name in names:
+        rot = fs_extri.getNode(f"Rot_{name}").mat()
+        t = fs_extri.getNode(f"T_{name}").mat()
+        if rot is None or t is None:
+            continue
+        mat = np.eye(4, dtype=np.float64)
+        mat[:3, :3] = np.asarray(rot, dtype=np.float64)
+        mat[:3, 3] = np.asarray(t, dtype=np.float64).reshape(3)
+        extri[name] = mat
+    fs_extri.release()
+
+    intri = {}
+    image_sizes = {}
+    if intri_path is not None and os.path.exists(intri_path):
+        fs_intri = cv2.FileStorage(intri_path, cv2.FILE_STORAGE_READ)
+        for name in names:
+            K = fs_intri.getNode(f"K_{name}").mat()
+            if K is None:
+                continue
+            intri[name] = np.asarray(K, dtype=np.float64)
+            h_node = fs_intri.getNode(f"H_{name}")
+            w_node = fs_intri.getNode(f"W_{name}")
+            if not h_node.empty() and not w_node.empty():
+                image_sizes[name] = (int(h_node.real()), int(w_node.real()))
+        fs_intri.release()
+
+    return extri, intri, image_sizes
+
+
+def read_depth(path):
+    depth = cv2.imread(path, cv2.IMREAD_ANYDEPTH)
+    if depth is None:
+        raise FileNotFoundError(path)
+    return depth.astype(np.float32)
+
+
+def read_ply_xyz(path):
+    if not os.path.exists(path):
+        raise FileNotFoundError(path)
+
+    header = []
+    with open(path, "rb") as f:
+        while True:
+            line = f.readline()
+            if not line:
+                raise ValueError(f"Invalid PLY header: {path}")
+            text = line.decode("ascii").strip()
+            header.append(text)
+            if text == "end_header":
+                break
+
+        if "format binary_little_endian 1.0" not in header:
+            raise ValueError(f"Unsupported PLY format: {path}")
+
+        vertex_count = None
+        property_specs = []
+        in_vertex_block = False
+        for line in header:
+            if line.startswith("element vertex "):
+                vertex_count = int(line.split()[-1])
+                in_vertex_block = True
+                continue
+            if line.startswith("element ") and not line.startswith("element vertex "):
+                in_vertex_block = False
+            if in_vertex_block and line.startswith("property "):
+                _, dtype_name, prop_name = line.split()
+                property_specs.append((dtype_name, prop_name))
+
+        if vertex_count is None:
+            raise ValueError(f"Missing vertex count in PLY: {path}")
+
+        dtype_map = {
+            "float": "<f4",
+            "float32": "<f4",
+            "uchar": "u1",
+            "uint8": "u1",
+        }
+        vertex_dtype = []
+        for dtype_name, prop_name in property_specs:
+            if dtype_name not in dtype_map:
+                raise ValueError(
+                    f"Unsupported PLY property type {dtype_name} in {path}"
+                )
+            vertex_dtype.append((prop_name, dtype_map[dtype_name]))
+
+        data = np.fromfile(f, dtype=np.dtype(vertex_dtype), count=vertex_count)
+    return np.stack([data["x"], data["y"], data["z"]], axis=1).astype(
+        np.float32, copy=False
+    )
+
+
+def read_pointcloud_xyz(path):
+    ext = os.path.splitext(path)[1].lower()
+    if ext == ".npy":
+        data = np.load(path)
+        return np.asarray(data, dtype=np.float32).reshape(-1, 3)
+    if ext == ".npz":
+        data = np.load(path)
+        if "points" in data:
+            points = data["points"]
+        else:
+            first_key = next(iter(data.files))
+            points = data[first_key]
+        return np.asarray(points, dtype=np.float32).reshape(-1, 3)
+    return read_ply_xyz(path)

longstream/eval/metrics.py

@@ -0,0 +1,116 @@
+import numpy as np
+from scipy.spatial import cKDTree
+
+
+def similarity_align(src, dst, with_scale=True):
+    src = np.asarray(src, dtype=np.float64)
+    dst = np.asarray(dst, dtype=np.float64)
+    if src.shape != dst.shape or src.ndim != 2 or src.shape[1] != 3:
+        raise ValueError("Expected Nx3 source and target point sets")
+    if len(src) < 3:
+        return 1.0, np.eye(3), np.zeros(3)
+
+    src_mean = src.mean(axis=0)
+    dst_mean = dst.mean(axis=0)
+    src_centered = src - src_mean
+    dst_centered = dst - dst_mean
+
+    cov = (dst_centered.T @ src_centered) / len(src)
+    U, D, Vt = np.linalg.svd(cov)
+    S = np.eye(3)
+    if np.linalg.det(U @ Vt) < 0:
+        S[-1, -1] = -1.0
+    R = U @ S @ Vt
+
+    if with_scale:
+        var = np.mean(np.sum(src_centered ** 2, axis=1))
+        scale = float(np.trace(np.diag(D) @ S) / max(var, 1e-12))
+    else:
+        scale = 1.0
+    t = dst_mean - scale * (R @ src_mean)
+    return scale, R, t
+
+
+def transform_points(points, scale, R, t):
+    return (scale * (R @ points.T)).T + t[None]
+
+
+def ate_rmse(pred_xyz, gt_xyz, align_scale=True):
+    scale, R, t = similarity_align(pred_xyz, gt_xyz, with_scale=align_scale)
+    pred_aligned = transform_points(pred_xyz, scale, R, t)
+    err = np.linalg.norm(pred_aligned - gt_xyz, axis=1)
+    return {
+        "ate_rmse": float(np.sqrt(np.mean(err ** 2))),
+        "ate_mean": float(np.mean(err)),
+        "ate_median": float(np.median(err)),
+        "num_pose_pairs": int(len(err)),
+        "align_scale": bool(align_scale),
+        "sim3_scale": float(scale),
+        "sim3_rotation": R.tolist(),
+        "sim3_translation": t.tolist(),
+    }
+
+
+def _voxel_downsample(points, voxel_size):
+    if voxel_size is None:
+        return points
+    voxel_size = float(voxel_size)
+    if voxel_size <= 0 or len(points) == 0:
+        return points
+    coords = np.floor(points / voxel_size).astype(np.int64)
+    _, keep = np.unique(coords, axis=0, return_index=True)
+    keep.sort()
+    return points[keep]
+
+
+def _sample_points(points, max_points, seed):
+    if max_points is None or len(points) <= int(max_points):
+        return points
+    rng = np.random.default_rng(seed)
+    keep = rng.choice(len(points), size=int(max_points), replace=False)
+    return points[keep]
+
+
+def prepare_pointcloud(points, max_points=None, voxel_size=None, seed=0):
+    points = np.asarray(points, dtype=np.float64).reshape(-1, 3)
+    if len(points) == 0:
+        return points
+    valid = np.isfinite(points).all(axis=1)
+    points = points[valid]
+    points = _voxel_downsample(points, voxel_size)
+    points = _sample_points(points, max_points, seed)
+    return points
+
+
+def chamfer_and_f1(
+    pred_points, gt_points, threshold=0.25, max_points=None, voxel_size=None, seed=0
+):
+    pred = prepare_pointcloud(
+        pred_points, max_points=max_points, voxel_size=voxel_size, seed=seed
+    )
+    gt = prepare_pointcloud(
+        gt_points, max_points=max_points, voxel_size=voxel_size, seed=seed + 1
+    )
+    if len(pred) == 0 or len(gt) == 0:
+        return None
+
+    pred_tree = cKDTree(pred)
+    gt_tree = cKDTree(gt)
+    dist_pred_to_gt, _ = gt_tree.query(pred, k=1)
+    dist_gt_to_pred, _ = pred_tree.query(gt, k=1)
+
+    acc = float(np.mean(dist_pred_to_gt))
+    comp = float(np.mean(dist_gt_to_pred))
+    precision = float(np.mean(dist_pred_to_gt < threshold))
+    recall = float(np.mean(dist_gt_to_pred < threshold))
+    denom = precision + recall
+    f1 = 0.0 if denom <= 0 else float(2.0 * precision * recall / denom)
+    return {
+        "cd": float(acc + comp),
+        "acc": acc,
+        "comp": comp,
+        "f1": f1,
+        "f1_threshold": float(threshold),
+        "num_pred_points": int(len(pred)),
+        "num_gt_points": int(len(gt)),
+    }

longstream/io/save_images.py

@@ -0,0 +1,38 @@
+import os
+import subprocess
+from typing import List
+
+import numpy as np
+from PIL import Image
+
+
+def save_image_sequence(
+    path, images: List[np.ndarray], prefix: str = "frame", ext: str = "png"
+):
+    os.makedirs(path, exist_ok=True)
+    for i, img in enumerate(images):
+        out_path = os.path.join(path, f"{prefix}_{i:06d}.{ext}")
+        Image.fromarray(img).save(out_path)
+
+
+def save_video(output_path, pattern, fps=30):
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    cmd = [
+        "ffmpeg",
+        "-hide_banner",
+        "-loglevel",
+        "error",
+        "-y",
+        "-framerate",
+        str(fps),
+        "-pattern_type",
+        "glob",
+        "-i",
+        pattern,
+        "-c:v",
+        "libx264",
+        "-pix_fmt",
+        "yuv420p",
+        output_path,
+    ]
+    subprocess.run(cmd, check=True)

longstream/io/save_points.py

@@ -0,0 +1,71 @@
+import os
+import numpy as np
+
+
+def _maybe_downsample(points, colors=None, max_points=None, seed=0):
+    pts = np.asarray(points).reshape(-1, 3)
+    cols = None if colors is None else np.asarray(colors).reshape(-1, 3)
+    if max_points is None or pts.shape[0] <= int(max_points):
+        return pts, cols
+    rng = np.random.default_rng(seed)
+    keep = rng.choice(pts.shape[0], size=int(max_points), replace=False)
+    pts = pts[keep]
+    if cols is not None:
+        cols = cols[keep]
+    return pts, cols
+
+
+def save_pointcloud(path, points, colors=None, max_points=None, seed=0):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    pts, cols = _maybe_downsample(
+        points, colors=colors, max_points=max_points, seed=seed
+    )
+    pts = pts.astype(np.float32, copy=False)
+    if colors is not None:
+        if cols.max() <= 1.0:
+            cols = (cols * 255.0).astype(np.uint8)
+        else:
+            cols = cols.astype(np.uint8)
+        has_color = True
+    else:
+        cols = None
+        has_color = False
+
+    with open(path, "wb") as f:
+        f.write(b"ply\n")
+        f.write(b"format binary_little_endian 1.0\n")
+        f.write(f"element vertex {pts.shape[0]}\n".encode("ascii"))
+        f.write(b"property float x\n")
+        f.write(b"property float y\n")
+        f.write(b"property float z\n")
+        if has_color:
+            f.write(b"property uchar red\n")
+            f.write(b"property uchar green\n")
+            f.write(b"property uchar blue\n")
+        f.write(b"end_header\n")
+        if has_color:
+            vertex_dtype = np.dtype(
+                [
+                    ("x", "<f4"),
+                    ("y", "<f4"),
+                    ("z", "<f4"),
+                    ("red", "u1"),
+                    ("green", "u1"),
+                    ("blue", "u1"),
+                ]
+            )
+            vertex_data = np.empty(pts.shape[0], dtype=vertex_dtype)
+            vertex_data["x"] = pts[:, 0]
+            vertex_data["y"] = pts[:, 1]
+            vertex_data["z"] = pts[:, 2]
+            vertex_data["red"] = cols[:, 0]
+            vertex_data["green"] = cols[:, 1]
+            vertex_data["blue"] = cols[:, 2]
+            vertex_data.tofile(f)
+        else:
+            vertex_dtype = np.dtype([("x", "<f4"), ("y", "<f4"), ("z", "<f4")])
+            vertex_data = np.empty(pts.shape[0], dtype=vertex_dtype)
+            vertex_data["x"] = pts[:, 0]
+            vertex_data["y"] = pts[:, 1]
+            vertex_data["z"] = pts[:, 2]
+            vertex_data.tofile(f)

longstream/io/save_poses_txt.py

@@ -0,0 +1,43 @@
+import os
+import numpy as np
+
+
+def _ensure_dir(path):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+
+
+def save_w2c_txt(path, extri, frames):
+    _ensure_dir(path)
+    with open(path, "w") as f:
+        f.write("# w2c\n")
+        for i, frame in enumerate(frames):
+            mat = extri[i]
+            r = mat[:3, :3].reshape(-1)
+            t = mat[:3, 3].reshape(-1)
+            vals = [frame] + r.tolist() + t.tolist()
+            f.write(" ".join([str(v) for v in vals]) + "\n")
+
+
+def save_intri_txt(path, intri, frames):
+    _ensure_dir(path)
+    with open(path, "w") as f:
+        f.write("# fx fy cx cy\n")
+        for i, frame in enumerate(frames):
+            k = intri[i]
+            fx = float(k[0, 0])
+            fy = float(k[1, 1])
+            cx = float(k[0, 2])
+            cy = float(k[1, 2])
+            f.write(f"{frame} {fx} {fy} {cx} {cy}\n")
+
+
+def save_rel_pose_txt(path, rel_pose_enc, frames):
+    _ensure_dir(path)
+    arr = rel_pose_enc
+    if hasattr(arr, "detach"):
+        arr = arr.detach().cpu().numpy()
+    with open(path, "w") as f:
+        f.write("# tx ty tz qx qy qz qw fov_h fov_w\n")
+        for i, frame in enumerate(frames):
+            vals = [frame] + arr[i].tolist()
+            f.write(" ".join([str(v) for v in vals]) + "\n")

longstream/models/__init__.py

@@ -0,0 +1,3 @@
+from longstream.models.longstream import LongStream
+
+__all__ = ["LongStream"]

longstream/models/longstream.py

@@ -0,0 +1,370 @@
+from typing import Tuple, List, Optional, Dict
+import torch
+import torch.nn as nn
+from huggingface_hub import PyTorchModelHubMixin
+
+from longstream.utils.vendor.dust3r.utils.misc import freeze_all_params
+from longstream.utils.vendor.models.components.aggregator.streamaggregator import (
+    STreamAggregator,
+)
+from longstream.utils.vendor.models.components.heads.camera_head import (
+    CameraHead,
+    RelPoseHead,
+)
+from longstream.utils.vendor.models.components.heads.dpt_head import DPTHead
+
+
+class LongStream(nn.Module, PyTorchModelHubMixin):
+    def __init__(
+        self,
+        img_size=518,
+        patch_size=14,
+        embed_dim=1024,
+        freeze="none",
+        rel_pose_head_cfg=None,
+        use_role_embedding=True,
+        enable_scale_token=False,
+        scale_token_config=None,
+        disable_keyframe_distinction=False,
+        enable_camera_head=True,
+        use_segment_mask=False,
+        use_3d_rope=False,
+        rope_freq=100,
+        window_size=5000,
+    ):
+        super().__init__()
+
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.embed_dim = embed_dim
+        self.enable_scale_token = enable_scale_token
+        self.enable_camera_head = enable_camera_head
+        self.window_size = window_size
+
+        self.aggregator = STreamAggregator(
+            img_size=img_size,
+            patch_size=patch_size,
+            embed_dim=embed_dim,
+            use_role_embedding=use_role_embedding,
+            disable_keyframe_distinction=disable_keyframe_distinction,
+            use_segment_mask=use_segment_mask,
+            use_3d_rope=use_3d_rope,
+            rope_freq=rope_freq,
+            window_size=window_size,
+        )
+
+        if self.enable_camera_head:
+            self.camera_head = CameraHead(dim_in=2 * embed_dim, window_size=window_size)
+        else:
+            self.camera_head = None
+        self.point_head = DPTHead(
+            dim_in=2 * embed_dim,
+            output_dim=4,
+            activation="inv_log",
+            conf_activation="expp1",
+        )
+        self.depth_head = DPTHead(
+            dim_in=2 * embed_dim,
+            output_dim=2,
+            activation="exp",
+            conf_activation="expp1",
+        )
+
+        self.rel_pose_head = None
+        self.reinit_camera_head_when_rel_enabled = False
+
+        if rel_pose_head_cfg is not None:
+            enable = rel_pose_head_cfg.get("enabled", True)
+            if enable:
+
+                head_cfg = {
+                    "dim_in": 2 * embed_dim,
+                    "trunk_depth": rel_pose_head_cfg.get("trunk_depth", 4),
+                    "pose_mode": rel_pose_head_cfg.get("pose_mode", "SE3"),
+                    "num_heads": rel_pose_head_cfg.get("num_heads", 16),
+                    "mlp_ratio": rel_pose_head_cfg.get("mlp_ratio", 4),
+                    "init_values": rel_pose_head_cfg.get("init_values", 0.01),
+                    "trans_act": rel_pose_head_cfg.get("trans_act", "linear"),
+                    "quat_act": rel_pose_head_cfg.get("quat_act", "linear"),
+                    "fl_act": rel_pose_head_cfg.get("fl_act", "relu"),
+                    "use_global_scale": rel_pose_head_cfg.get(
+                        "use_global_scale", False
+                    ),
+                    "use_pair_cross_attn": rel_pose_head_cfg.get(
+                        "use_pair_cross_attn", False
+                    ),
+                    "detach_reference": rel_pose_head_cfg.get(
+                        "detach_reference", False
+                    ),
+                    "xattn_temperature": rel_pose_head_cfg.get(
+                        "xattn_temperature", 1.0
+                    ),
+                    "use_precat": rel_pose_head_cfg.get("use_precat", False),
+                    "use_kf_role_embed": rel_pose_head_cfg.get(
+                        "use_kf_role_embed", True
+                    ),
+                    "kf_role_embed_init_std": rel_pose_head_cfg.get(
+                        "kf_role_embed_init_std", 0.02
+                    ),
+                    "window_size": window_size,
+                }
+                self.rel_pose_head = RelPoseHead(**head_cfg)
+
+                self.reinit_camera_head_when_rel_enabled = rel_pose_head_cfg.get(
+                    "reinit_camera_head", False
+                )
+
+                if self.reinit_camera_head_when_rel_enabled:
+                    pass
+
+        if self.enable_scale_token:
+            self._init_scale_components(scale_token_config or {})
+
+        self.set_freeze(freeze)
+
+    def reinitialize_camera_head(self):
+        """
+        Reinitialize camera_head with fresh weights.
+
+        This is useful when:
+        1. Loading a pretrained checkpoint that has camera_head weights
+        2. But we want to train camera_head from scratch with new settings (e.g., quaternion normalization)
+
+        This method should be called AFTER checkpoint loading.
+        """
+
+        old_camera_head = self.camera_head
+        dim_in = old_camera_head.token_norm.normalized_shape[0]
+
+        self.camera_head = CameraHead(dim_in=dim_in)
+
+        device = next(old_camera_head.parameters()).device
+        self.camera_head = self.camera_head.to(device)
+
+    def _init_scale_components(self, config):
+        self.scale_token = nn.Parameter(torch.zeros(self.embed_dim))
+        torch.nn.init.trunc_normal_(self.scale_token, std=0.02)
+
+        self.scale_head = nn.Sequential(
+            nn.Linear(2 * self.embed_dim, 256),
+            nn.ReLU(),
+            nn.Linear(256, 128),
+            nn.ReLU(),
+            nn.Linear(128, 1),
+        )
+
+        for m in self.scale_head.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight, gain=1.0)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0.0)
+
+        import math
+
+        nn.init.constant_(self.scale_head[-1].bias, math.log(30.0))
+
+    def set_freeze(self, freeze):
+        self.freeze = freeze
+
+        to_be_frozen = {
+            "none": [],
+            "encoder": [self.aggregator.patch_embed],
+        }
+        freeze_all_params(to_be_frozen[freeze])
+
+    def forward(
+        self,
+        images: torch.Tensor,
+        mode: str = "causal",
+        aggregator_kv_cache_list: Optional[List[List[torch.Tensor]]] = None,
+        camera_head_kv_cache_list: Optional[List[List[List[torch.Tensor]]]] = None,
+        rel_pose_inputs: Optional[Dict] = None,
+        is_keyframe: Optional[torch.Tensor] = None,
+    ):
+
+        if len(images.shape) == 4:
+            images = images.unsqueeze(0)
+
+        batch_size = images.shape[0]
+
+        additional_tokens = None
+        if self.enable_scale_token:
+
+            scale_token_base = self.scale_token.unsqueeze(0).repeat(batch_size, 1)
+            additional_tokens = scale_token_base.unsqueeze(-1)
+
+        keyframe_indices = None
+        if rel_pose_inputs is not None and "keyframe_indices" in rel_pose_inputs:
+            keyframe_indices = rel_pose_inputs["keyframe_indices"]
+
+        if aggregator_kv_cache_list is not None:
+            (
+                aggregated_tokens_list,
+                patch_start_idx,
+                aggregator_kv_cache_list,
+                _,
+            ) = self.aggregator(
+                images,
+                mode=mode,
+                kv_cache_list=aggregator_kv_cache_list,
+                is_keyframe=is_keyframe,
+                keyframe_indices=keyframe_indices,
+                additional_tokens=additional_tokens,
+                reorder_keyframes_first=False,
+            )
+        else:
+            aggregated_tokens_list, patch_start_idx, _ = self.aggregator(
+                images,
+                mode=mode,
+                is_keyframe=is_keyframe,
+                keyframe_indices=keyframe_indices,
+                additional_tokens=additional_tokens,
+                reorder_keyframes_first=False,
+            )
+
+        predictions = {}
+
+        predicted_scale_factor = None
+        if self.enable_scale_token and additional_tokens is not None:
+
+            if len(aggregated_tokens_list) > 0:
+                last_layer_features = aggregated_tokens_list[-1]
+
+                scale_token_idx = patch_start_idx - 1
+                scale_token_output_features = last_layer_features[
+                    :, :, scale_token_idx, :
+                ]
+
+                scale_token_output_features = scale_token_output_features.mean(dim=1)
+
+                scale_logits = self.scale_head(scale_token_output_features).squeeze(-1)
+
+                predicted_scale_factor = torch.exp(scale_logits)
+
+                predictions["predicted_scale_factor"] = predicted_scale_factor
+                predictions["scale_token_features"] = scale_token_output_features
+
+        if self.enable_camera_head and self.camera_head is not None:
+            if camera_head_kv_cache_list is not None:
+                pose_enc_list, camera_head_kv_cache_list = self.camera_head(
+                    aggregated_tokens_list,
+                    mode=mode,
+                    kv_cache_list=camera_head_kv_cache_list,
+                )
+            else:
+                pose_enc_list = self.camera_head(aggregated_tokens_list, mode=mode)
+
+            final_pose_enc = pose_enc_list[-1]
+            if self.enable_scale_token and predicted_scale_factor is not None:
+                scale = predicted_scale_factor.view(-1, 1, 1)
+
+                scaled_t = final_pose_enc[..., :3] * scale
+                scaled_pose_enc = torch.cat([scaled_t, final_pose_enc[..., 3:]], dim=-1)
+                predictions["pose_enc"] = scaled_pose_enc
+            else:
+                predictions["pose_enc"] = final_pose_enc
+
+            if self.training:
+
+                if self.enable_scale_token and predicted_scale_factor is not None:
+                    scale = predicted_scale_factor.view(-1, 1, 1)
+                    scaled_pose_enc_list = []
+                    for pose_enc in pose_enc_list:
+
+                        scaled_t = pose_enc[..., :3] * scale
+                        scaled_pose_enc = torch.cat(
+                            [scaled_t, pose_enc[..., 3:]], dim=-1
+                        )
+                        scaled_pose_enc_list.append(scaled_pose_enc)
+                    predictions["pose_enc_list"] = scaled_pose_enc_list
+                else:
+                    predictions["pose_enc_list"] = pose_enc_list
+
+        if self.rel_pose_head is not None and rel_pose_inputs is not None:
+
+            rel_kwargs = dict(
+                aggregated_tokens_list=aggregated_tokens_list,
+                keyframe_indices=rel_pose_inputs.get("keyframe_indices"),
+                is_keyframe=rel_pose_inputs.get("is_keyframe", is_keyframe),
+                num_iterations=rel_pose_inputs.get("num_iterations", 4),
+                mode=mode,
+                kv_cache_list=rel_pose_inputs.get("kv_cache_list"),
+            )
+
+            rel_kwargs = {k: v for k, v in rel_kwargs.items() if v is not None}
+
+            rel_result = self.rel_pose_head(**rel_kwargs)
+
+            if isinstance(rel_result, dict):
+
+                pose_enc = rel_result["pose_enc"]
+                if pose_enc.dtype != torch.float32:
+                    pose_enc = pose_enc.float()
+
+                if self.enable_scale_token and predicted_scale_factor is not None:
+                    scale = predicted_scale_factor.view(-1, 1, 1)
+
+                    scaled_t = pose_enc[..., :3] * scale
+                    scaled_rel_pose_enc = torch.cat(
+                        [scaled_t, pose_enc[..., 3:]], dim=-1
+                    )
+                    predictions["rel_pose_enc"] = scaled_rel_pose_enc
+
+                    if "pose_enc_list" in rel_result:
+                        scaled_pose_enc_list = []
+                        for iter_pose in rel_result["pose_enc_list"]:
+                            scaled_t = iter_pose[..., :3] * scale
+                            scaled_iter_pose = torch.cat(
+                                [scaled_t, iter_pose[..., 3:]], dim=-1
+                            )
+                            scaled_pose_enc_list.append(scaled_iter_pose)
+                        predictions["rel_pose_enc_list"] = scaled_pose_enc_list
+                else:
+                    predictions["rel_pose_enc"] = pose_enc
+
+                    if "pose_enc_list" in rel_result:
+                        predictions["rel_pose_enc_list"] = rel_result["pose_enc_list"]
+
+                predictions["is_keyframe"] = rel_result.get("is_keyframe")
+                predictions["keyframe_indices"] = rel_result.get("keyframe_indices")
+
+                if "global_scale" in rel_result:
+                    predictions["global_scale"] = rel_result["global_scale"]
+
+            if "kv_cache_list" in rel_result:
+                predictions["rel_pose_kv_cache_list"] = rel_result["kv_cache_list"]
+
+        if self.point_head is not None:
+            pts3d, pts3d_conf = self.point_head(
+                aggregated_tokens_list, images=images, patch_start_idx=patch_start_idx
+            )
+
+            if self.enable_scale_token and predicted_scale_factor is not None:
+                scale = predicted_scale_factor.view(-1, 1, 1, 1, 1)
+                predictions["world_points"] = pts3d * scale
+            else:
+                predictions["world_points"] = pts3d
+            predictions["world_points_conf"] = pts3d_conf
+
+        if self.depth_head is not None:
+            depth, depth_conf = self.depth_head(
+                aggregated_tokens_list, images=images, patch_start_idx=patch_start_idx
+            )
+
+            if self.enable_scale_token and predicted_scale_factor is not None:
+                scale = predicted_scale_factor.view(-1, 1, 1, 1, 1)
+                predictions["depth"] = depth * scale
+            else:
+                predictions["depth"] = depth
+            predictions["depth_conf"] = depth_conf
+
+        if aggregator_kv_cache_list is not None:
+            predictions["aggregator_kv_cache_list"] = aggregator_kv_cache_list
+
+        if camera_head_kv_cache_list is not None:
+            predictions["camera_head_kv_cache_list"] = camera_head_kv_cache_list
+
+        if not self.training:
+            predictions["images"] = images
+
+        return predictions

longstream/streaming/keyframe_selector.py

@@ -0,0 +1,80 @@
+import random
+import torch
+from typing import Optional, Tuple
+
+
+class KeyframeSelector:
+    def __init__(
+        self,
+        min_interval: int = 8,
+        max_interval: int = 8,
+        force_first: bool = True,
+        motion_threshold: Optional[float] = None,
+        mode: str = "fixed",
+    ):
+        self.min_interval = int(min_interval)
+        self.max_interval = int(max_interval)
+        self.force_first = bool(force_first)
+        self.motion_threshold = motion_threshold
+        self.mode = mode
+
+    def select_keyframes(
+        self,
+        sequence_length: int,
+        batch_size: int = 1,
+        device: Optional[torch.device] = None,
+        poses: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        device = device or torch.device("cpu")
+        is_keyframe = torch.zeros(
+            batch_size, sequence_length, dtype=torch.bool, device=device
+        )
+        keyframe_indices = torch.zeros(
+            batch_size, sequence_length, dtype=torch.long, device=device
+        )
+
+        for b in range(batch_size):
+            last_keyframe_idx = 0
+            next_keyframe_target = None
+
+            if self.force_first or sequence_length == 1:
+                is_keyframe[b, 0] = True
+                keyframe_indices[b, 0] = 0
+                if self.mode == "random":
+                    interval = random.randint(self.min_interval, self.max_interval)
+                    next_keyframe_target = interval
+
+            for s in range(1, sequence_length):
+                keyframe_indices[b, s] = last_keyframe_idx
+                frames_since_last = s - last_keyframe_idx
+
+                if self.mode == "random" and next_keyframe_target is not None:
+                    if s >= next_keyframe_target:
+                        is_keyframe[b, s] = True
+                        last_keyframe_idx = s
+                        interval = random.randint(self.min_interval, self.max_interval)
+                        next_keyframe_target = s + interval
+                elif frames_since_last >= self.max_interval:
+                    is_keyframe[b, s] = True
+                    last_keyframe_idx = s
+                    if self.mode == "random":
+                        interval = random.randint(self.min_interval, self.max_interval)
+                        next_keyframe_target = s + interval
+                elif (
+                    frames_since_last >= self.min_interval
+                    and poses is not None
+                    and self.motion_threshold is not None
+                ):
+                    motion = torch.norm(
+                        poses[b, s, :3] - poses[b, last_keyframe_idx, :3]
+                    ).item()
+                    if motion > self.motion_threshold:
+                        is_keyframe[b, s] = True
+                        last_keyframe_idx = s
+                        if self.mode == "random":
+                            interval = random.randint(
+                                self.min_interval, self.max_interval
+                            )
+                            next_keyframe_target = s + interval
+
+        return is_keyframe, keyframe_indices

longstream/streaming/refresh.py

@@ -0,0 +1,217 @@
+import torch
+from typing import Dict, Any, List
+
+from longstream.streaming.stream_session import StreamSession
+
+_SEQUENCE_OUTPUT_KEYS = {
+    "pose_enc",
+    "rel_pose_enc",
+    "world_points",
+    "world_points_conf",
+    "depth",
+    "depth_conf",
+}
+_SCALAR_OUTPUT_KEYS = {
+    "predicted_scale_factor",
+    "global_scale",
+}
+
+
+def _refresh_intervals(refresh: int) -> int:
+    refresh = int(refresh)
+    if refresh < 2:
+        raise ValueError("refresh must be >= 2")
+    return refresh - 1
+
+
+def _model_device(model) -> torch.device:
+    return next(model.parameters()).device
+
+
+def _move_scalar_to_cpu(value: Any) -> Any:
+    if isinstance(value, torch.Tensor):
+        return value.detach().cpu()
+    return value
+
+
+def _append_batch_output(
+    stitched_tensors: Dict[str, List[torch.Tensor]],
+    stitched_scalars: Dict[str, Any],
+    output: Dict[str, Any],
+    actual_frames: int,
+    slice_start: int,
+) -> None:
+    for key in _SEQUENCE_OUTPUT_KEYS:
+        value = output.get(key)
+        if not isinstance(value, torch.Tensor):
+            continue
+        if value.ndim < 2 or value.shape[1] != actual_frames:
+            continue
+        stitched_tensors.setdefault(key, []).append(
+            value[:, slice_start:].detach().cpu()
+        )
+
+    for key in _SCALAR_OUTPUT_KEYS:
+        if key in output:
+            stitched_scalars[key] = _move_scalar_to_cpu(output[key])
+
+
+def _finalize_stitched_batches(
+    stitched_tensors: Dict[str, List[torch.Tensor]],
+    stitched_scalars: Dict[str, Any],
+) -> Dict[str, Any]:
+    stitched_output: Dict[str, Any] = {}
+    for key, chunks in stitched_tensors.items():
+        if not chunks:
+            continue
+        stitched_output[key] = (
+            chunks[0] if len(chunks) == 1 else torch.cat(chunks, dim=1)
+        )
+    stitched_output.update(stitched_scalars)
+    return stitched_output
+
+
+def run_batch_refresh(
+    model,
+    images,
+    is_keyframe,
+    keyframe_indices,
+    mode: str,
+    keyframe_stride: int,
+    refresh: int,
+    rel_pose_cfg,
+):
+    B, S = images.shape[:2]
+    device = _model_device(model)
+    refresh_intervals = _refresh_intervals(refresh)
+    frames_per_batch = refresh_intervals * keyframe_stride + 1
+    step_frames = refresh_intervals * keyframe_stride
+
+    stitched_tensors: Dict[str, List[torch.Tensor]] = {}
+    stitched_scalars: Dict[str, Any] = {}
+    num_batches = (S + step_frames - 1) // step_frames
+    for batch_idx in range(num_batches):
+        start_frame = batch_idx * step_frames
+        end_frame = min(start_frame + frames_per_batch, S)
+        batch_images = images[:, start_frame:end_frame].to(device, non_blocking=True)
+        batch_is_keyframe = (
+            is_keyframe[:, start_frame:end_frame].clone()
+            if is_keyframe is not None
+            else None
+        )
+        batch_keyframe_indices = (
+            keyframe_indices[:, start_frame:end_frame].clone()
+            if keyframe_indices is not None
+            else None
+        )
+
+        if batch_idx > 0 and batch_is_keyframe is not None:
+            batch_is_keyframe[:, 0] = True
+            if batch_keyframe_indices is not None:
+                batch_keyframe_indices[:, 0] = start_frame
+
+        if batch_keyframe_indices is not None:
+            batch_keyframe_indices = batch_keyframe_indices - start_frame
+            batch_keyframe_indices = torch.clamp(
+                batch_keyframe_indices, 0, end_frame - start_frame - 1
+            )
+
+        batch_rel_pose_inputs = None
+        if rel_pose_cfg is not None and batch_is_keyframe is not None:
+            batch_is_keyframe = batch_is_keyframe.to(device, non_blocking=True)
+            if batch_keyframe_indices is not None:
+                batch_keyframe_indices = batch_keyframe_indices.to(
+                    device, non_blocking=True
+                )
+            batch_rel_pose_inputs = {
+                "is_keyframe": batch_is_keyframe,
+                "keyframe_indices": batch_keyframe_indices,
+                "num_iterations": rel_pose_cfg.get("num_iterations", 4),
+            }
+        elif batch_is_keyframe is not None:
+            batch_is_keyframe = batch_is_keyframe.to(device, non_blocking=True)
+
+        batch_output = model(
+            images=batch_images,
+            mode=mode,
+            rel_pose_inputs=batch_rel_pose_inputs,
+            is_keyframe=batch_is_keyframe,
+        )
+
+        _append_batch_output(
+            stitched_tensors,
+            stitched_scalars,
+            batch_output,
+            actual_frames=end_frame - start_frame,
+            slice_start=0 if batch_idx == 0 else 1,
+        )
+        del batch_output
+        del batch_images
+        del batch_is_keyframe
+        del batch_keyframe_indices
+
+    return _finalize_stitched_batches(stitched_tensors, stitched_scalars)
+
+
+def run_streaming_refresh(
+    model,
+    images,
+    is_keyframe,
+    keyframe_indices,
+    mode: str,
+    window_size: int,
+    refresh: int,
+    rel_pose_cfg,
+):
+    B, S = images.shape[:2]
+    device = _model_device(model)
+    refresh_intervals = _refresh_intervals(refresh)
+    session = StreamSession(model, mode=mode, window_size=window_size)
+    keyframe_count = 0
+    segment_start = 0
+    for s in range(S):
+        frame_images = images[:, s : s + 1].to(device, non_blocking=True)
+        is_keyframe_s = (
+            is_keyframe[:, s : s + 1].to(device, non_blocking=True)
+            if is_keyframe is not None
+            else None
+        )
+        if keyframe_indices is not None:
+            keyframe_indices_s = keyframe_indices[:, s : s + 1].clone() - segment_start
+            keyframe_indices_s = torch.clamp(keyframe_indices_s, min=0)
+            keyframe_indices_s = keyframe_indices_s.to(device, non_blocking=True)
+        else:
+            keyframe_indices_s = None
+        session.forward_stream(
+            frame_images,
+            is_keyframe=is_keyframe_s,
+            keyframe_indices=keyframe_indices_s,
+            record=True,
+        )
+        if is_keyframe_s is None or not bool(is_keyframe_s.item()) or s <= 0:
+            del frame_images
+            if is_keyframe_s is not None:
+                del is_keyframe_s
+            if keyframe_indices_s is not None:
+                del keyframe_indices_s
+            continue
+        keyframe_count += 1
+        if keyframe_count % refresh_intervals == 0:
+            session.clear_cache_only()
+            segment_start = s
+            if keyframe_indices_s is not None:
+                keyframe_indices_self = torch.zeros_like(keyframe_indices_s)
+            else:
+                keyframe_indices_self = None
+            session.forward_stream(
+                frame_images,
+                is_keyframe=is_keyframe_s,
+                keyframe_indices=keyframe_indices_self,
+                record=False,
+            )
+        del frame_images
+        if is_keyframe_s is not None:
+            del is_keyframe_s
+        if keyframe_indices_s is not None:
+            del keyframe_indices_s
+    return session.get_all_predictions()

longstream/streaming/stream_session.py

@@ -0,0 +1,294 @@
+import torch
+
+
+class StreamSession:
+    def __init__(
+        self,
+        model,
+        mode: str,
+        window_size: int = 5,
+        keep_first_frame_anchor: bool = True,
+    ):
+        self.model = model
+        self.core_model = getattr(model, "longstream", model)
+        self.mode = mode
+        self.window_size = window_size
+        self.keep_first_frame_anchor = keep_first_frame_anchor
+
+        if self.mode not in ["causal", "window"]:
+            raise ValueError(f"Unsupported attention mode: {self.mode}")
+
+        self.aggregator_kv_cache_depth = self.core_model.aggregator.depth
+        self.use_camera_head = self.core_model.camera_head is not None
+        if self.use_camera_head:
+            self.camera_head_kv_cache_depth = self.core_model.camera_head.trunk_depth
+            self.camera_head_iterations = 4
+        else:
+            self.camera_head_kv_cache_depth = 0
+            self.camera_head_iterations = 0
+
+        self.use_rel_pose_head = (
+            hasattr(self.core_model, "rel_pose_head")
+            and self.core_model.rel_pose_head is not None
+        )
+        if self.use_rel_pose_head:
+            self.rel_pose_head_trunk_depth = self.core_model.rel_pose_head.trunk_depth
+            self.rel_pose_head_iterations = 4
+
+        self.clear()
+
+    def _clear_predictions(self):
+        self.sequence_predictions = {}
+        self.scalar_predictions = {}
+
+    def _update_predictions(self, predictions):
+        sequence_keys = [
+            "pose_enc",
+            "rel_pose_enc",
+            "world_points",
+            "world_points_conf",
+            "depth",
+            "depth_conf",
+        ]
+        scalar_keys = ["predicted_scale_factor", "global_scale"]
+
+        for k in sequence_keys:
+            if k in predictions:
+                self.sequence_predictions.setdefault(k, []).append(
+                    predictions[k].detach().cpu()
+                )
+
+        for k in scalar_keys:
+            if k in predictions:
+                value = predictions[k]
+                self.scalar_predictions[k] = (
+                    value.detach().cpu() if isinstance(value, torch.Tensor) else value
+                )
+
+    def _clear_cache(self):
+        self.aggregator_kv_cache_list = [
+            [None, None] for _ in range(self.aggregator_kv_cache_depth)
+        ]
+        if self.use_camera_head:
+            self.camera_head_kv_cache_list = [
+                [[None, None] for _ in range(self.camera_head_kv_cache_depth)]
+                for _ in range(self.camera_head_iterations)
+            ]
+        else:
+            self.camera_head_kv_cache_list = None
+        if self.use_rel_pose_head:
+            self.rel_pose_kv_cache_list = [
+                [[None, None] for _ in range(self.rel_pose_head_trunk_depth)]
+                for _ in range(self.rel_pose_head_iterations)
+            ]
+        else:
+            self.rel_pose_kv_cache_list = None
+
+    def _update_cache(
+        self, aggregator_kv_cache_list, camera_head_kv_cache_list, frame_hw
+    ):
+        if self.mode == "causal":
+            self.aggregator_kv_cache_list = aggregator_kv_cache_list
+            if self.use_camera_head:
+                self.camera_head_kv_cache_list = camera_head_kv_cache_list
+            return
+
+        if self.mode == "window":
+            h, w = frame_hw
+            P = (
+                h
+                * w
+                // self.core_model.aggregator.patch_size
+                // self.core_model.aggregator.patch_size
+                + self.core_model.aggregator.patch_start_idx
+            )
+
+            for k in range(2):
+                for i in range(self.aggregator_kv_cache_depth):
+                    cache_size = aggregator_kv_cache_list[i][k].size(2)
+                    if self.keep_first_frame_anchor:
+                        if cache_size <= P:
+                            self.aggregator_kv_cache_list[i][
+                                k
+                            ] = aggregator_kv_cache_list[i][k].contiguous()
+                        elif cache_size <= self.window_size * P:
+                            self.aggregator_kv_cache_list[i][
+                                k
+                            ] = aggregator_kv_cache_list[i][k].contiguous()
+                        else:
+                            anchor = aggregator_kv_cache_list[i][k][:, :, :P]
+                            recent_start = cache_size - (self.window_size - 1) * P
+                            recent = aggregator_kv_cache_list[i][k][:, :, recent_start:]
+                            self.aggregator_kv_cache_list[i][k] = torch.cat(
+                                [anchor, recent], dim=2
+                            ).contiguous()
+                    else:
+                        start_idx = max(0, cache_size - self.window_size * P)
+                        self.aggregator_kv_cache_list[i][k] = aggregator_kv_cache_list[
+                            i
+                        ][k][:, :, start_idx:].contiguous()
+
+            if camera_head_kv_cache_list is not None:
+                for k in range(2):
+                    for i in range(self.camera_head_iterations):
+                        for j in range(self.camera_head_kv_cache_depth):
+                            cache_size = camera_head_kv_cache_list[i][j][k].size(2)
+                            if self.keep_first_frame_anchor:
+                                if cache_size <= 1:
+                                    self.camera_head_kv_cache_list[i][j][
+                                        k
+                                    ] = camera_head_kv_cache_list[i][j][k].contiguous()
+                                elif cache_size <= self.window_size:
+                                    self.camera_head_kv_cache_list[i][j][
+                                        k
+                                    ] = camera_head_kv_cache_list[i][j][k].contiguous()
+                                else:
+                                    anchor = camera_head_kv_cache_list[i][j][k][
+                                        :, :, :1
+                                    ]
+                                    recent_start = cache_size - (self.window_size - 1)
+                                    recent = camera_head_kv_cache_list[i][j][k][
+                                        :, :, recent_start:
+                                    ]
+                                    self.camera_head_kv_cache_list[i][j][k] = torch.cat(
+                                        [anchor, recent], dim=2
+                                    ).contiguous()
+                            else:
+                                start_idx = max(0, cache_size - self.window_size)
+                                self.camera_head_kv_cache_list[i][j][
+                                    k
+                                ] = camera_head_kv_cache_list[i][j][k][
+                                    :, :, start_idx:
+                                ].contiguous()
+            return
+
+        raise ValueError(f"Unsupported attention mode: {self.mode}")
+
+    def _get_cache(self):
+        return self.aggregator_kv_cache_list, self.camera_head_kv_cache_list
+
+    def get_all_predictions(self):
+        predictions = {}
+        for key, chunks in self.sequence_predictions.items():
+            if not chunks:
+                continue
+            predictions[key] = (
+                chunks[0] if len(chunks) == 1 else torch.cat(chunks, dim=1)
+            )
+        predictions.update(self.scalar_predictions)
+        return predictions
+
+    def get_last_prediction(self):
+        last_predictions = {}
+        keys_to_extract = [
+            "pose_enc",
+            "rel_pose_enc",
+            "world_points",
+            "world_points_conf",
+            "depth",
+            "depth_conf",
+            "predicted_scale_factor",
+        ]
+        for k in keys_to_extract:
+            if k in self.sequence_predictions and self.sequence_predictions[k]:
+                last_predictions[k] = self.sequence_predictions[k][-1][:, -1:]
+            elif k in self.scalar_predictions:
+                last_predictions[k] = self.scalar_predictions[k]
+        return last_predictions
+
+    def clear(self):
+        self._clear_predictions()
+        self._clear_cache()
+        if self.use_rel_pose_head:
+            if hasattr(self.core_model.rel_pose_head, "_keyframe_tokens_cache"):
+                self.core_model.rel_pose_head._keyframe_tokens_cache = {}
+            if hasattr(self.core_model.rel_pose_head, "_current_frame_id"):
+                self.core_model.rel_pose_head._current_frame_id = 0
+            if hasattr(self.core_model.rel_pose_head, "_frame_info"):
+                self.core_model.rel_pose_head._frame_info = []
+
+    def clear_cache_only(self):
+        self._clear_cache()
+        if self.use_rel_pose_head:
+            if hasattr(self.core_model.rel_pose_head, "_keyframe_tokens_cache"):
+                self.core_model.rel_pose_head._keyframe_tokens_cache = {}
+            if hasattr(self.core_model.rel_pose_head, "_current_frame_id"):
+                self.core_model.rel_pose_head._current_frame_id = 0
+            if hasattr(self.core_model.rel_pose_head, "_frame_info"):
+                self.core_model.rel_pose_head._frame_info = []
+
+    def forward_stream(
+        self, images, is_keyframe=None, keyframe_indices=None, record: bool = True
+    ):
+        aggregator_kv_cache_list, camera_head_kv_cache_list = self._get_cache()
+
+        rel_pose_inputs = None
+        if (
+            self.use_rel_pose_head
+            and is_keyframe is not None
+            and keyframe_indices is not None
+        ):
+            rel_pose_inputs = {
+                "is_keyframe": is_keyframe,
+                "keyframe_indices": keyframe_indices,
+                "kv_cache_list": self.rel_pose_kv_cache_list,
+            }
+
+        outputs = self.model(
+            images=images,
+            mode=self.mode,
+            aggregator_kv_cache_list=aggregator_kv_cache_list,
+            camera_head_kv_cache_list=camera_head_kv_cache_list,
+            rel_pose_inputs=rel_pose_inputs,
+            is_keyframe=is_keyframe,
+        )
+
+        if record:
+            self._update_predictions(outputs)
+
+        camera_head_kv_cache_list = outputs.get("camera_head_kv_cache_list", None)
+        depth_hw = (
+            outputs["depth"].shape[2:4] if "depth" in outputs else images.shape[-2:]
+        )
+        self._update_cache(
+            outputs["aggregator_kv_cache_list"], camera_head_kv_cache_list, depth_hw
+        )
+
+        if self.use_rel_pose_head and "rel_pose_kv_cache_list" in outputs:
+            rel_pose_kv_cache = outputs["rel_pose_kv_cache_list"]
+            if self.mode == "causal":
+                self.rel_pose_kv_cache_list = rel_pose_kv_cache
+            elif self.mode == "window":
+                for k in range(2):
+                    for i in range(self.rel_pose_head_iterations):
+                        for j in range(self.rel_pose_head_trunk_depth):
+                            if rel_pose_kv_cache[i][j][k] is None:
+                                continue
+                            cache_len = rel_pose_kv_cache[i][j][k].size(2)
+                            if self.keep_first_frame_anchor:
+                                if cache_len <= 1:
+                                    self.rel_pose_kv_cache_list[i][j][
+                                        k
+                                    ] = rel_pose_kv_cache[i][j][k].contiguous()
+                                elif cache_len <= self.window_size:
+                                    self.rel_pose_kv_cache_list[i][j][
+                                        k
+                                    ] = rel_pose_kv_cache[i][j][k].contiguous()
+                                else:
+                                    anchor = rel_pose_kv_cache[i][j][k][:, :, :1]
+                                    recent_start = cache_len - (self.window_size - 1)
+                                    recent = rel_pose_kv_cache[i][j][k][
+                                        :, :, recent_start:
+                                    ]
+                                    self.rel_pose_kv_cache_list[i][j][k] = torch.cat(
+                                        [anchor, recent], dim=2
+                                    ).contiguous()
+                            else:
+                                start_idx = max(0, cache_len - self.window_size)
+                                self.rel_pose_kv_cache_list[i][j][
+                                    k
+                                ] = rel_pose_kv_cache[i][j][k][
+                                    :, :, start_idx:
+                                ].contiguous()
+
+        return outputs

longstream/utils/camera.py

@@ -0,0 +1,50 @@
+import torch
+from longstream.utils.vendor.models.components.utils.rotation import (
+    quat_to_mat,
+    mat_to_quat,
+)
+
+
+def compose_abs_from_rel(
+    rel_pose_enc: torch.Tensor, keyframe_indices: torch.Tensor
+) -> torch.Tensor:
+    squeeze_batch = False
+    if rel_pose_enc.ndim == 2:
+        rel_pose_enc = rel_pose_enc.unsqueeze(0)
+        squeeze_batch = True
+    if keyframe_indices.ndim == 1:
+        keyframe_indices = keyframe_indices.unsqueeze(0)
+    if rel_pose_enc.ndim != 3 or keyframe_indices.ndim != 2:
+        raise ValueError(
+            f"Expected rel_pose_enc [B,S,D] or [S,D] and keyframe_indices [B,S] or [S], "
+            f"got {tuple(rel_pose_enc.shape)} and {tuple(keyframe_indices.shape)}"
+        )
+
+    B, S, _ = rel_pose_enc.shape
+    device = rel_pose_enc.device
+    dtype = rel_pose_enc.dtype
+
+    rel_t = rel_pose_enc[..., :3]
+    rel_q = rel_pose_enc[..., 3:7]
+    rel_f = rel_pose_enc[..., 7:9]
+    rel_R = quat_to_mat(rel_q.reshape(-1, 4)).reshape(B, S, 3, 3)
+
+    abs_R = torch.zeros(B, S, 3, 3, device=device, dtype=dtype)
+    abs_t = torch.zeros(B, S, 3, device=device, dtype=dtype)
+    abs_f = torch.zeros(B, S, 2, device=device, dtype=dtype)
+
+    for b in range(B):
+        abs_R[b, 0] = rel_R[b, 0]
+        abs_t[b, 0] = rel_t[b, 0]
+        abs_f[b, 0] = rel_f[b, 0]
+        for s in range(1, S):
+            ref_idx = int(keyframe_indices[b, s].item())
+            abs_R[b, s] = rel_R[b, s] @ abs_R[b, ref_idx]
+            abs_t[b, s] = rel_t[b, s] + rel_R[b, s] @ abs_t[b, ref_idx]
+            abs_f[b, s] = rel_f[b, s]
+
+    abs_q = mat_to_quat(abs_R.reshape(-1, 3, 3)).reshape(B, S, 4)
+    abs_pose_enc = torch.cat([abs_t, abs_q, abs_f], dim=-1)
+    if squeeze_batch:
+        return abs_pose_enc[0]
+    return abs_pose_enc

longstream/utils/depth.py

@@ -0,0 +1,36 @@
+import numpy as np
+import torch
+import matplotlib.cm as cm
+
+
+def colorize_depth(depth: torch.Tensor, cmap: str = "plasma") -> np.ndarray:
+    if torch.is_tensor(depth):
+        depth_np = depth.detach().cpu().numpy()
+    else:
+        depth_np = depth
+    d_min = np.nanmin(depth_np)
+    d_max = np.nanmax(depth_np)
+    if d_max - d_min < 1e-6:
+        d_max = d_min + 1e-6
+    norm = (depth_np - d_min) / (d_max - d_min)
+    norm = np.clip(norm, 0.0, 1.0)
+    mapper = cm.get_cmap(cmap)
+    colored = mapper(norm)[..., :3]
+    return (colored * 255.0).astype(np.uint8)
+
+
+def unproject_depth_to_points(depth: torch.Tensor, intri: torch.Tensor) -> torch.Tensor:
+    B, H, W = depth.shape
+    fx = intri[:, 0, 0].view(B, 1, 1)
+    fy = intri[:, 1, 1].view(B, 1, 1)
+    cx = intri[:, 0, 2].view(B, 1, 1)
+    cy = intri[:, 1, 2].view(B, 1, 1)
+
+    ys = torch.arange(H, device=depth.device).view(1, H, 1).float()
+    xs = torch.arange(W, device=depth.device).view(1, 1, W).float()
+
+    x = (xs - cx) * depth / fx
+    y = (ys - cy) * depth / fy
+    z = depth
+    pts = torch.stack([x, y, z], dim=-1)
+    return pts

longstream/utils/hub.py

@@ -0,0 +1,42 @@
+import os
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class HFSpec:
+    repo_id: str
+    filename: str
+    revision: Optional[str] = None
+    local_dir: str = "checkpoints"
+
+
+def _is_nonempty_str(x) -> bool:
+    return isinstance(x, str) and len(x) > 0
+
+
+def resolve_checkpoint_path(
+    checkpoint: Optional[str], hf: Optional[dict]
+) -> Optional[str]:
+    if _is_nonempty_str(checkpoint):
+        return checkpoint
+    if not isinstance(hf, dict):
+        return None
+
+    repo_id = hf.get("repo_id")
+    filename = hf.get("filename")
+    revision = hf.get("revision", None)
+    local_dir = hf.get("local_dir", "checkpoints")
+
+    if not _is_nonempty_str(repo_id) or not _is_nonempty_str(filename):
+        return None
+
+    try:
+        from huggingface_hub import hf_hub_download
+    except Exception as e:
+        raise RuntimeError("huggingface_hub is required for auto-download") from e
+
+    os.makedirs(local_dir, exist_ok=True)
+    return hf_hub_download(
+        repo_id=repo_id, filename=filename, revision=revision, local_dir=local_dir
+    )

longstream/utils/sky_mask.py

@@ -0,0 +1,100 @@
+import os
+import copy
+import cv2
+import numpy as np
+import shutil
+import urllib.request
+
+try:
+    import onnxruntime
+except Exception:
+    onnxruntime = None
+
+SKYSEG_URL = "https://huggingface.co/JianyuanWang/skyseg/resolve/main/skyseg.onnx"
+SKYSEG_THRESHOLD = 0.5
+
+
+def run_skyseg(session, input_size, image):
+    temp_image = copy.deepcopy(image)
+    resize_image = cv2.resize(temp_image, dsize=(input_size[0], input_size[1]))
+    x = cv2.cvtColor(resize_image, cv2.COLOR_BGR2RGB)
+    x = np.array(x, dtype=np.float32)
+    mean = [0.485, 0.456, 0.406]
+    std = [0.229, 0.224, 0.225]
+    x = (x / 255 - mean) / std
+    x = x.transpose(2, 0, 1)
+    x = x.reshape(-1, 3, input_size[0], input_size[1]).astype("float32")
+    input_name = session.get_inputs()[0].name
+    result_map = session.run(None, {input_name: x})[0]
+    return result_map[0, 0]
+
+
+def _normalize_skyseg_output(result_map):
+    result_map = np.asarray(result_map, dtype=np.float32)
+    if result_map.size == 0:
+        return result_map
+    finite = np.isfinite(result_map)
+    if not np.any(finite):
+        return np.zeros_like(result_map, dtype=np.float32)
+    result_map = np.nan_to_num(result_map, nan=0.0, posinf=1.0, neginf=0.0)
+    max_value = float(result_map.max())
+    min_value = float(result_map.min())
+    if min_value >= 0.0 and max_value > 1.5:
+        result_map = result_map / 255.0
+    return np.clip(result_map, 0.0, 1.0)
+
+
+def sky_mask_filename(image_path):
+    parent = os.path.basename(os.path.dirname(image_path))
+    name = os.path.basename(image_path)
+    if parent:
+        return f"{parent}__{name}"
+    return name
+
+
+def segment_sky(image_path, session, mask_filename=None):
+    image = cv2.imread(image_path)
+    if image is None:
+        return None
+    result_map = run_skyseg(session, [320, 320], image)
+    result_map_original = cv2.resize(result_map, (image.shape[1], image.shape[0]))
+    result_map_original = _normalize_skyseg_output(result_map_original)
+    output_mask = np.zeros(result_map_original.shape, dtype=np.uint8)
+    output_mask[result_map_original < SKYSEG_THRESHOLD] = 255
+    if mask_filename is not None:
+        os.makedirs(os.path.dirname(mask_filename), exist_ok=True)
+        cv2.imwrite(mask_filename, output_mask)
+    return output_mask
+
+
+def compute_sky_mask(image_paths, model_path: str, target_dir: str = None):
+    if onnxruntime is None:
+        return None
+    if not os.path.exists(model_path):
+        os.makedirs(os.path.dirname(os.path.abspath(model_path)), exist_ok=True)
+        try:
+            print(f"[longstream] downloading skyseg.onnx to {model_path}", flush=True)
+            with urllib.request.urlopen(SKYSEG_URL) as src, open(
+                model_path, "wb"
+            ) as dst:
+                shutil.copyfileobj(src, dst)
+        except Exception as exc:
+            print(f"[longstream] failed to download skyseg.onnx: {exc}", flush=True)
+            return None
+    if not os.path.exists(model_path):
+        return None
+    session = onnxruntime.InferenceSession(model_path)
+    masks = []
+    for image_path in image_paths:
+        mask_filepath = None
+        if target_dir is not None:
+            name = sky_mask_filename(image_path)
+            mask_filepath = os.path.join(target_dir, name)
+            if os.path.exists(mask_filepath):
+                sky_mask = cv2.imread(mask_filepath, cv2.IMREAD_GRAYSCALE)
+            else:
+                sky_mask = segment_sky(image_path, session, mask_filepath)
+        else:
+            sky_mask = segment_sky(image_path, session, None)
+        masks.append(sky_mask)
+    return masks

longstream/utils/vendor/__init__.py

@@ -0,0 +1,2 @@
+
+

longstream/utils/vendor/croco/LICENSE

@@ -0,0 +1,52 @@
+CroCo, Copyright (c) 2022-present Naver Corporation, is licensed under the Creative Commons Attribution-NonCommercial-ShareAlike 4.0 license.
+
+A summary of the CC BY-NC-SA 4.0 license is located here:
+	https://creativecommons.org/licenses/by-nc-sa/4.0/
+
+The CC BY-NC-SA 4.0 license is located here:
+	https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode
+
+
+SEE NOTICE BELOW WITH RESPECT TO THE FILE: models/pos_embed.py, models/blocks.py
+
+***************************
+
+NOTICE WITH RESPECT TO THE FILE: models/pos_embed.py
+
+This software is being redistributed in a modifiled form. The original form is available here:
+
+https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py
+
+This software in this file incorporates parts of the following software available here:
+
+Transformer: https://github.com/tensorflow/models/blob/master/official/legacy/transformer/model_utils.py
+available under the following license: https://github.com/tensorflow/models/blob/master/LICENSE
+
+MoCo v3: https://github.com/facebookresearch/moco-v3
+available under the following license: https://github.com/facebookresearch/moco-v3/blob/main/LICENSE
+
+DeiT: https://github.com/facebookresearch/deit
+available under the following license: https://github.com/facebookresearch/deit/blob/main/LICENSE
+
+
+ORIGINAL COPYRIGHT NOTICE AND PERMISSION NOTICE AVAILABLE HERE IS REPRODUCE BELOW:
+
+https://github.com/facebookresearch/mae/blob/main/LICENSE
+
+Attribution-NonCommercial 4.0 International
+
+***************************
+
+NOTICE WITH RESPECT TO THE FILE: models/blocks.py
+
+This software is being redistributed in a modifiled form. The original form is available here:
+
+https://github.com/rwightman/pytorch-image-models
+
+ORIGINAL COPYRIGHT NOTICE AND PERMISSION NOTICE AVAILABLE HERE IS REPRODUCE BELOW:
+
+https://github.com/rwightman/pytorch-image-models/blob/master/LICENSE
+
+Apache License
+Version 2.0, January 2004
+http://www.apache.org/licenses/

longstream/utils/vendor/croco/NOTICE

@@ -0,0 +1,21 @@
+CroCo
+Copyright 2022-present NAVER Corp.
+
+This project contains subcomponents with separate copyright notices and license terms.
+Your use of the source code for these subcomponents is subject to the terms and conditions of the following licenses.
+
+====
+
+facebookresearch/mae
+https://github.com/facebookresearch/mae
+
+Attribution-NonCommercial 4.0 International
+
+====
+
+rwightman/pytorch-image-models
+https://github.com/rwightman/pytorch-image-models
+
+Apache License
+Version 2.0, January 2004
+http://www.apache.org/licenses/

longstream/utils/vendor/croco/README.MD

@@ -0,0 +1,124 @@
+# CroCo + CroCo v2 / CroCo-Stereo / CroCo-Flow
+
+[[`CroCo arXiv`](https://arxiv.org/abs/2210.10716)] [[`CroCo v2 arXiv`](https://arxiv.org/abs/2211.10408)] [[`project page and demo`](https://croco.europe.naverlabs.com/)]
+
+This repository contains the code for our CroCo model presented in our NeurIPS'22 paper [CroCo: Self-Supervised Pre-training for 3D Vision Tasks by Cross-View Completion](https://openreview.net/pdf?id=wZEfHUM5ri) and its follow-up extension published at ICCV'23 [Improved Cross-view Completion Pre-training for Stereo Matching and Optical Flow](https://openaccess.thecvf.com/content/ICCV2023/html/Weinzaepfel_CroCo_v2_Improved_Cross-view_Completion_Pre-training_for_Stereo_Matching_and_ICCV_2023_paper.html), refered to as CroCo v2:
+
+![image](assets/arch.jpg)
+
+```bibtex
+@inproceedings{croco,
+  title={{CroCo: Self-Supervised Pre-training for 3D Vision Tasks by Cross-View Completion}},
+  author={{Weinzaepfel, Philippe and Leroy, Vincent and Lucas, Thomas and Br\'egier, Romain and Cabon, Yohann and Arora, Vaibhav and Antsfeld, Leonid and Chidlovskii, Boris and Csurka, Gabriela and Revaud J\'er\^ome}},
+  booktitle={{NeurIPS}},
+  year={2022}
+}
+
+@inproceedings{croco_v2,
+  title={{CroCo v2: Improved Cross-view Completion Pre-training for Stereo Matching and Optical Flow}},
+  author={Weinzaepfel, Philippe and Lucas, Thomas and Leroy, Vincent and Cabon, Yohann and Arora, Vaibhav and Br{\'e}gier, Romain and Csurka, Gabriela and Antsfeld, Leonid and Chidlovskii, Boris and Revaud, J{\'e}r{\^o}me},
+  booktitle={ICCV},
+  year={2023}
+}
+```
+
+## License
+
+The code is distributed under the CC BY-NC-SA 4.0 License. See [LICENSE](LICENSE) for more information.
+Some components are based on code from [MAE](https://github.com/facebookresearch/mae) released under the CC BY-NC-SA 4.0 License and [timm](https://github.com/rwightman/pytorch-image-models) released under the Apache 2.0 License.
+Some components for stereo matching and optical flow are based on code from [unimatch](https://github.com/autonomousvision/unimatch) released under the MIT license.
+
+## Preparation
+
+1. Install dependencies on a machine with a NVidia GPU using e.g. conda. Note that `habitat-sim` is required only for the interactive demo and the synthetic pre-training data generation. If you don't plan to use it, you can ignore the line installing it and use a more recent python version.
+
+```bash
+conda create -n croco python=3.7 cmake=3.14.0
+conda activate croco
+conda install habitat-sim headless -c conda-forge -c aihabitat
+conda install pytorch torchvision -c pytorch
+conda install notebook ipykernel matplotlib
+conda install ipywidgets widgetsnbextension
+conda install scikit-learn tqdm quaternion opencv # only for pretraining / habitat data generation
+
+```
+
+2. Compile cuda kernels for RoPE
+
+CroCo v2 relies on RoPE positional embeddings for which you need to compile some cuda kernels.
+```bash
+cd models/curope/
+python setup.py build_ext --inplace
+cd ../../
+```
+
+This can be a bit long as we compile for all cuda architectures, feel free to update L9 of `models/curope/setup.py` to compile for specific architectures only.
+You might also need to set the environment `CUDA_HOME` in case you use a custom cuda installation.
+
+In case you cannot provide, we also provide a slow pytorch version, which will be automatically loaded.
+
+3. Download pre-trained model
+
+We provide several pre-trained models:
+
+| modelname                                                                                                                          | pre-training data | pos. embed. | Encoder | Decoder |
+|------------------------------------------------------------------------------------------------------------------------------------|-------------------|-------------|---------|---------|
+| [`CroCo.pth`](https://download.europe.naverlabs.com/ComputerVision/CroCo/CroCo.pth)                                                 | Habitat           | cosine      | ViT-B   | Small   |
+| [`CroCo_V2_ViTBase_SmallDecoder.pth`](https://download.europe.naverlabs.com/ComputerVision/CroCo/CroCo_V2_ViTBase_SmallDecoder.pth) | Habitat + real    | RoPE        | ViT-B   | Small   |
+| [`CroCo_V2_ViTBase_BaseDecoder.pth`](https://download.europe.naverlabs.com/ComputerVision/CroCo/CroCo_V2_ViTBase_BaseDecoder.pth)   | Habitat + real    | RoPE        | ViT-B   | Base    |
+| [`CroCo_V2_ViTLarge_BaseDecoder.pth`](https://download.europe.naverlabs.com/ComputerVision/CroCo/CroCo_V2_ViTLarge_BaseDecoder.pth) | Habitat + real    | RoPE        | ViT-L   | Base    |
+
+To download a specific model, i.e., the first one (`CroCo.pth`)
+```bash
+mkdir -p pretrained_models/
+wget https://download.europe.naverlabs.com/ComputerVision/CroCo/CroCo.pth -P pretrained_models/
+```
+
+## Reconstruction example
+
+Simply run after downloading the `CroCo_V2_ViTLarge_BaseDecoder` pretrained model (or update the corresponding line in `demo.py`)
+```bash
+python demo.py
+```
+
+## Interactive demonstration of cross-view completion reconstruction on the Habitat simulator
+
+First download the test scene from Habitat:
+```bash
+python -m habitat_sim.utils.datasets_download --uids habitat_test_scenes --data-path habitat-sim-data/
+```
+
+Then, run the Notebook demo `interactive_demo.ipynb`.
+
+In this demo, you should be able to sample a random reference viewpoint from an [Habitat](https://github.com/facebookresearch/habitat-sim) test scene. Use the sliders to change viewpoint and select a masked target view to reconstruct using CroCo.
+![croco_interactive_demo](https://user-images.githubusercontent.com/1822210/200516576-7937bc6a-55f8-49ed-8618-3ddf89433ea4.jpg)
+
+## Pre-training
+
+### CroCo
+
+To pre-train CroCo, please first generate the pre-training data from the Habitat simulator, following the instructions in [datasets/habitat_sim/README.MD](datasets/habitat_sim/README.MD) and then run the following command:
+```
+torchrun --nproc_per_node=4 pretrain.py --output_dir ./output/pretraining/
+```
+
+Our CroCo pre-training was launched on a single server with 4 GPUs.
+It should take around 10 days with A100 or 15 days with V100 to do the 400 pre-training epochs, but decent performances are obtained earlier in training.
+Note that, while the code contains the same scaling rule of the learning rate as MAE when changing the effective batch size, we did not experimented if it is valid in our case.
+The first run can take a few minutes to start, to parse all available pre-training pairs.
+
+### CroCo v2
+
+For CroCo v2 pre-training, in addition to the generation of the pre-training data from the Habitat simulator above, please pre-extract the crops from the real datasets following the instructions in [datasets/crops/README.MD](datasets/crops/README.MD).
+Then, run the following command for the largest model (ViT-L encoder, Base decoder):
+```
+torchrun --nproc_per_node=8 pretrain.py --model "CroCoNet(enc_embed_dim=1024, enc_depth=24, enc_num_heads=16, dec_embed_dim=768, dec_num_heads=12, dec_depth=12, pos_embed='RoPE100')" --dataset "habitat_release+ARKitScenes+MegaDepth+3DStreetView+IndoorVL" --warmup_epochs 12 --max_epoch 125 --epochs 250 --amp 0 --keep_freq 5 --output_dir ./output/pretraining_crocov2/
+```
+
+Our CroCo v2 pre-training was launched on a single server with 8 GPUs for the largest model, and on a single server with 4 GPUs for the smaller ones, keeping a batch size of 64 per gpu in all cases.
+The largest model should take around 12 days on A100.
+Note that, while the code contains the same scaling rule of the learning rate as MAE when changing the effective batch size, we did not experimented if it is valid in our case.
+
+## Stereo matching and Optical flow downstream tasks
+
+For CroCo-Stereo and CroCo-Flow, please refer to [stereoflow/README.MD](stereoflow/README.MD).

longstream/utils/vendor/croco/croco-stereo-flow-demo.ipynb

@@ -0,0 +1,182 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "9bca0f41",
+   "metadata": {},
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "80653ef7",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4f033862",
+   "metadata": {},
+   "source": [
+    "First download the model(s) of your choice by running\n",
+    "```\n",
+    "bash stereoflow/download_model.sh crocostereo.pth\n",
+    "bash stereoflow/download_model.sh crocoflow.pth\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1fb2e392",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "use_gpu = torch.cuda.is_available() and torch.cuda.device_count()>0\n",
+    "device = torch.device('cuda:0' if use_gpu else 'cpu')\n",
+    "import matplotlib.pylab as plt"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e0e25d77",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from stereoflow.test import _load_model_and_criterion\n",
+    "from stereoflow.engine import tiled_pred\n",
+    "from stereoflow.datasets_stereo import img_to_tensor, vis_disparity\n",
+    "from stereoflow.datasets_flow import flowToColor\n",
+    "tile_overlap=0.7 # recommended value, higher value can be slightly better but slower"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "86a921f5",
+   "metadata": {},
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "64e483cb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image1 = np.asarray(Image.open('<path_to_left_image>'))\n",
+    "image2 = np.asarray(Image.open('<path_to_right_image>'))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f0d04303",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model, _, cropsize, with_conf, task, tile_conf_mode = _load_model_and_criterion('stereoflow_models/crocostereo.pth', None, device)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "47dc14b5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "im1 = img_to_tensor(image1).to(device).unsqueeze(0)\n",
+    "im2 = img_to_tensor(image2).to(device).unsqueeze(0)\n",
+    "with torch.inference_mode():\n",
+    "    pred, _, _ = tiled_pred(model, None, im1, im2, None, conf_mode=tile_conf_mode, overlap=tile_overlap, crop=cropsize, with_conf=with_conf, return_time=False)\n",
+    "pred = pred.squeeze(0).squeeze(0).cpu().numpy()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "583b9f16",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.imshow(vis_disparity(pred))\n",
+    "plt.axis('off')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d2df5d70",
+   "metadata": {},
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9ee257a7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image1 = np.asarray(Image.open('<path_to_first_image>'))\n",
+    "image2 = np.asarray(Image.open('<path_to_second_image>'))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d5edccf0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model, _, cropsize, with_conf, task, tile_conf_mode = _load_model_and_criterion('stereoflow_models/crocoflow.pth', None, device)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b19692c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "im1 = img_to_tensor(image1).to(device).unsqueeze(0)\n",
+    "im2 = img_to_tensor(image2).to(device).unsqueeze(0)\n",
+    "with torch.inference_mode():\n",
+    "    pred, _, _ = tiled_pred(model, None, im1, im2, None, conf_mode=tile_conf_mode, overlap=tile_overlap, crop=cropsize, with_conf=with_conf, return_time=False)\n",
+    "pred = pred.squeeze(0).permute(1,2,0).cpu().numpy()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "26f79db3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.imshow(flowToColor(pred))\n",
+    "plt.axis('off')"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

longstream/utils/vendor/croco/datasets/__init__.py

@@ -0,0 +1,2 @@
+
+

longstream/utils/vendor/croco/datasets/crops/README.MD

@@ -0,0 +1,104 @@
+## Generation of crops from the real datasets
+
+The instructions below allow to generate the crops used for pre-training CroCo v2 from the following real-world datasets: ARKitScenes, MegaDepth, 3DStreetView and IndoorVL.
+
+### Download the metadata of the crops to generate
+
+First, download the metadata and put them in `./data/`:
+```
+mkdir -p data
+cd data/
+wget https://download.europe.naverlabs.com/ComputerVision/CroCo/data/crop_metadata.zip
+unzip crop_metadata.zip
+rm crop_metadata.zip
+cd ..
+```
+
+### Prepare the original datasets
+
+Second, download the original datasets in `./data/original_datasets/`.
+```
+mkdir -p data/original_datasets
+```
+
+##### ARKitScenes
+
+Download the `raw` dataset from https://github.com/apple/ARKitScenes/blob/main/DATA.md and put it in `./data/original_datasets/ARKitScenes/`.
+The resulting file structure should be like:
+```
+./data/original_datasets/ARKitScenes/
+└───Training
+    └───40753679
+     │  │   ultrawide
+     │  │   ...
+     └───40753686
+     │
+      ...
+```
+
+##### MegaDepth
+
+Download `MegaDepth v1 Dataset` from https://www.cs.cornell.edu/projects/megadepth/ and put it in `./data/original_datasets/MegaDepth/`.
+The resulting file structure should be like:
+
+```
+./data/original_datasets/MegaDepth/
+└───0000
+│   └───images
+│    │      │   1000557903_87fa96b8a4_o.jpg
+│    │      └ ...
+│    └─── ...
+└───0001
+│   │
+│   └ ...
+└─── ...
+```
+
+##### 3DStreetView
+
+Download `3D_Street_View` dataset from https://github.com/amir32002/3D_Street_View and put it in `./data/original_datasets/3DStreetView/`.
+The resulting file structure should be like:
+
+```
+./data/original_datasets/3DStreetView/
+└───dataset_aligned
+│   └───0002
+│    │      │   0000002_0000001_0000002_0000001.jpg
+│    │      └ ...
+│    └─── ...
+└───dataset_unaligned
+│   └───0003
+│    │      │   0000003_0000001_0000002_0000001.jpg
+│    │      └ ...
+│    └─── ...
+```
+
+##### IndoorVL
+
+Download the `IndoorVL` datasets using [Kapture](https://github.com/naver/kapture).
+
+```
+pip install kapture
+mkdir -p ./data/original_datasets/IndoorVL
+cd ./data/original_datasets/IndoorVL
+kapture_download_dataset.py update
+kapture_download_dataset.py install  "HyundaiDepartmentStore_*"
+kapture_download_dataset.py install  "GangnamStation_*"
+cd -
+```
+
+### Extract the crops
+
+Now, extract the crops for each of the dataset:
+```
+for dataset in ARKitScenes MegaDepth 3DStreetView IndoorVL;
+do
+  python3 datasets/crops/extract_crops_from_images.py --crops ./data/crop_metadata/${dataset}/crops_release.txt --root-dir ./data/original_datasets/${dataset}/ --output-dir ./data/${dataset}_crops/ --imsize 256 --nthread 8 --max-subdir-levels 5 --ideal-number-pairs-in-dir 500;
+done
+```
+
+##### Note for IndoorVL
+
+Due to some legal issues, we can only release 144,228 pairs out of the 1,593,689 pairs used in the paper.
+To account for it in terms of number of pre-training iterations, the pre-training command in this repository uses 125 training epochs including 12 warm-up epochs and learning rate cosine schedule of 250, instead of 100, 10 and 200 respectively.
+The impact on the performance is negligible.

longstream/utils/vendor/croco/datasets/crops/extract_crops_from_images.py

@@ -0,0 +1,175 @@
+import argparse
+import functools
+import math
+import os
+from multiprocessing import Pool
+
+from PIL import Image
+from tqdm import tqdm
+
+
+def arg_parser():
+    parser = argparse.ArgumentParser(
+        "Generate cropped image pairs from image crop list"
+    )
+
+    parser.add_argument("--crops", type=str, required=True, help="crop file")
+    parser.add_argument("--root-dir", type=str, required=True, help="root directory")
+    parser.add_argument(
+        "--output-dir", type=str, required=True, help="output directory"
+    )
+    parser.add_argument("--imsize", type=int, default=256, help="size of the crops")
+    parser.add_argument(
+        "--nthread", type=int, required=True, help="number of simultaneous threads"
+    )
+    parser.add_argument(
+        "--max-subdir-levels",
+        type=int,
+        default=5,
+        help="maximum number of subdirectories",
+    )
+    parser.add_argument(
+        "--ideal-number-pairs-in-dir",
+        type=int,
+        default=500,
+        help="number of pairs stored in a dir",
+    )
+    return parser
+
+
+def main(args):
+    listing_path = os.path.join(args.output_dir, "listing.txt")
+
+    print(f"Loading list of crops ... ({args.nthread} threads)")
+    crops, num_crops_to_generate = load_crop_file(args.crops)
+
+    print(f"Preparing jobs ({len(crops)} candidate image pairs)...")
+    num_levels = min(
+        math.ceil(math.log(num_crops_to_generate, args.ideal_number_pairs_in_dir)),
+        args.max_subdir_levels,
+    )
+    num_pairs_in_dir = math.ceil(num_crops_to_generate ** (1 / num_levels))
+
+    jobs = prepare_jobs(crops, num_levels, num_pairs_in_dir)
+    del crops
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    mmap = Pool(args.nthread).imap_unordered if args.nthread > 1 else map
+    call = functools.partial(save_image_crops, args)
+
+    print(f"Generating cropped images to {args.output_dir} ...")
+    with open(listing_path, "w") as listing:
+        listing.write("# pair_path\n")
+        for results in tqdm(mmap(call, jobs), total=len(jobs)):
+            for path in results:
+                listing.write(f"{path}\n")
+    print("Finished writing listing to", listing_path)
+
+
+def load_crop_file(path):
+    data = open(path).read().splitlines()
+    pairs = []
+    num_crops_to_generate = 0
+    for line in tqdm(data):
+        if line.startswith("#"):
+            continue
+        line = line.split(", ")
+        if len(line) < 8:
+            img1, img2, rotation = line
+            pairs.append((img1, img2, int(rotation), []))
+        else:
+            l1, r1, t1, b1, l2, r2, t2, b2 = map(int, line)
+            rect1, rect2 = (l1, t1, r1, b1), (l2, t2, r2, b2)
+            pairs[-1][-1].append((rect1, rect2))
+            num_crops_to_generate += 1
+    return pairs, num_crops_to_generate
+
+
+def prepare_jobs(pairs, num_levels, num_pairs_in_dir):
+    jobs = []
+    powers = [num_pairs_in_dir ** level for level in reversed(range(num_levels))]
+
+    def get_path(idx):
+        idx_array = []
+        d = idx
+        for level in range(num_levels - 1):
+            idx_array.append(idx // powers[level])
+            idx = idx % powers[level]
+        idx_array.append(d)
+        return "/".join(map(lambda x: hex(x)[2:], idx_array))
+
+    idx = 0
+    for pair_data in tqdm(pairs):
+        img1, img2, rotation, crops = pair_data
+        if -60 <= rotation and rotation <= 60:
+            rotation = 0
+        paths = [get_path(idx + k) for k in range(len(crops))]
+        idx += len(crops)
+        jobs.append(((img1, img2), rotation, crops, paths))
+    return jobs
+
+
+def load_image(path):
+    try:
+        return Image.open(path).convert("RGB")
+    except Exception as e:
+        print("skipping", path, e)
+        raise OSError()
+
+
+def save_image_crops(args, data):
+
+    img_pair, rot, crops, paths = data
+    try:
+        img1, img2 = [
+            load_image(os.path.join(args.root_dir, impath)) for impath in img_pair
+        ]
+    except OSError as e:
+        return []
+
+    def area(sz):
+        return sz[0] * sz[1]
+
+    tgt_size = (args.imsize, args.imsize)
+
+    def prepare_crop(img, rect, rot=0):
+
+        img = img.crop(rect)
+
+        interp = (
+            Image.Resampling.LANCZOS
+            if area(img.size) > 4 * area(tgt_size)
+            else Image.Resampling.BICUBIC
+        )
+        img = img.resize(tgt_size, resample=interp)
+
+        rot90 = (round(rot / 90) % 4) * 90
+        if rot90 == 90:
+            img = img.transpose(Image.Transpose.ROTATE_90)
+        elif rot90 == 180:
+            img = img.transpose(Image.Transpose.ROTATE_180)
+        elif rot90 == 270:
+            img = img.transpose(Image.Transpose.ROTATE_270)
+        return img
+
+    results = []
+    for (rect1, rect2), path in zip(crops, paths):
+        crop1 = prepare_crop(img1, rect1)
+        crop2 = prepare_crop(img2, rect2, rot)
+
+        fullpath1 = os.path.join(args.output_dir, path + "_1.jpg")
+        fullpath2 = os.path.join(args.output_dir, path + "_2.jpg")
+        os.makedirs(os.path.dirname(fullpath1), exist_ok=True)
+
+        assert not os.path.isfile(fullpath1), fullpath1
+        assert not os.path.isfile(fullpath2), fullpath2
+        crop1.save(fullpath1)
+        crop2.save(fullpath2)
+        results.append(path)
+
+    return results
+
+
+if __name__ == "__main__":
+    args = arg_parser().parse_args()
+    main(args)

longstream/utils/vendor/croco/datasets/habitat_sim/README.MD

@@ -0,0 +1,76 @@
+## Generation of synthetic image pairs using Habitat-Sim
+
+These instructions allow to generate pre-training pairs from the Habitat simulator.
+As we did not save metadata of the pairs used in the original paper, they are not strictly the same, but these data use the same setting and are equivalent.
+
+### Download Habitat-Sim scenes
+Download Habitat-Sim scenes:
+- Download links can be found here: https://github.com/facebookresearch/habitat-sim/blob/main/DATASETS.md
+- We used scenes from the HM3D, habitat-test-scenes, Replica, ReplicaCad and ScanNet datasets.
+- Please put the scenes under `./data/habitat-sim-data/scene_datasets/` following the structure below, or update manually paths in `paths.py`.
+```
+./data/
+└──habitat-sim-data/
+   └──scene_datasets/
+      ├──hm3d/
+      ├──gibson/
+      ├──habitat-test-scenes/
+      ├──replica_cad_baked_lighting/
+      ├──replica_cad/
+      ├──ReplicaDataset/
+      └──scannet/
+```
+
+### Image pairs generation
+We provide metadata to generate reproducible images pairs for pretraining and validation.
+Experiments described in the paper used similar data, but whose generation was not reproducible at the time.
+
+Specifications:
+- 256x256 resolution images, with 60 degrees field of view .
+- Up to 1000 image pairs per scene.
+- Number of scenes considered/number of images pairs per dataset:
+  - Scannet: 1097 scenes / 985 209 pairs
+  - HM3D:
+    - hm3d/train: 800 / 800k pairs
+    - hm3d/val: 100 scenes / 100k pairs
+    - hm3d/minival: 10 scenes / 10k pairs
+  - habitat-test-scenes: 3 scenes / 3k pairs
+  - replica_cad_baked_lighting: 13 scenes / 13k pairs
+
+- Scenes from hm3d/val and hm3d/minival pairs were not used for the pre-training but kept for validation purposes.
+
+Download metadata and extract it:
+```bash
+mkdir -p data/habitat_release_metadata/
+cd data/habitat_release_metadata/
+wget https://download.europe.naverlabs.com/ComputerVision/CroCo/data/habitat_release_metadata/multiview_habitat_metadata.tar.gz
+tar -xvf multiview_habitat_metadata.tar.gz
+cd ../..
+# Location of the metadata
+METADATA_DIR="./data/habitat_release_metadata/multiview_habitat_metadata"
+```
+
+Generate image pairs from metadata:
+- The following command will print a list of commandlines to generate image pairs for each scene:
+```bash
+# Target output directory
+PAIRS_DATASET_DIR="./data/habitat_release/"
+python datasets/habitat_sim/generate_from_metadata_files.py --input_dir=$METADATA_DIR --output_dir=$PAIRS_DATASET_DIR
+```
+- One can launch multiple of such commands in parallel e.g. using GNU Parallel:
+```bash
+python datasets/habitat_sim/generate_from_metadata_files.py --input_dir=$METADATA_DIR --output_dir=$PAIRS_DATASET_DIR | parallel -j 16
+```
+
+## Metadata generation
+
+Image pairs were randomly sampled using the following commands, whose outputs contain randomness and are thus not exactly reproducible:
+```bash
+# Print commandlines to generate image pairs from the different scenes available.
+PAIRS_DATASET_DIR=MY_CUSTOM_PATH
+python datasets/habitat_sim/generate_multiview_images.py --list_commands --output_dir=$PAIRS_DATASET_DIR
+
+# Once a dataset is generated, pack metadata files for reproducibility.
+METADATA_DIR=MY_CUSTON_PATH
+python datasets/habitat_sim/pack_metadata_files.py $PAIRS_DATASET_DIR  $METADATA_DIR
+```

longstream/utils/vendor/croco/datasets/habitat_sim/__init__.py

@@ -0,0 +1,2 @@
+
+

longstream/utils/vendor/croco/datasets/habitat_sim/generate_from_metadata.py

@@ -0,0 +1,121 @@
+"""
+Script to generate image pairs for a given scene reproducing poses provided in a metadata file.
+"""
+import argparse
+import json
+import os
+
+import cv2
+import PIL.Image
+import quaternion
+from datasets.habitat_sim.multiview_habitat_sim_generator import (
+    MultiviewHabitatSimGenerator,
+)
+from datasets.habitat_sim.paths import SCENES_DATASET
+from tqdm import tqdm
+
+
+def generate_multiview_images_from_metadata(
+    metadata_filename,
+    output_dir,
+    overload_params=dict(),
+    scene_datasets_paths=None,
+    exist_ok=False,
+):
+    """
+    Generate images from a metadata file for reproducibility purposes.
+    """
+
+    if scene_datasets_paths is not None:
+        scene_datasets_paths = dict(
+            sorted(scene_datasets_paths.items(), key=lambda x: len(x[0]), reverse=True)
+        )
+
+    with open(metadata_filename, "r") as f:
+        input_metadata = json.load(f)
+    metadata = dict()
+    for key, value in input_metadata.items():
+
+        if key in ("scene_dataset_config_file", "scene", "navmesh") and value != "":
+            if scene_datasets_paths is not None:
+                for dataset_label, dataset_path in scene_datasets_paths.items():
+                    if value.startswith(dataset_label):
+                        value = os.path.normpath(
+                            os.path.join(
+                                dataset_path, os.path.relpath(value, dataset_label)
+                            )
+                        )
+                        break
+        metadata[key] = value
+
+    for key, value in overload_params.items():
+        metadata[key] = value
+
+    generation_entries = dict(
+        [
+            (key, value)
+            for key, value in metadata.items()
+            if not (key in ("multiviews", "output_dir", "generate_depth"))
+        ]
+    )
+    generate_depth = metadata["generate_depth"]
+
+    os.makedirs(output_dir, exist_ok=exist_ok)
+
+    generator = MultiviewHabitatSimGenerator(**generation_entries)
+
+    for idx_label, data in tqdm(metadata["multiviews"].items()):
+        positions = data["positions"]
+        orientations = data["orientations"]
+        n = len(positions)
+        for oidx in range(n):
+            observation = generator.render_viewpoint(
+                positions[oidx], quaternion.from_float_array(orientations[oidx])
+            )
+            observation_label = f"{oidx + 1}"
+
+            img = PIL.Image.fromarray(observation["color"][:, :, :3])
+            filename = os.path.join(output_dir, f"{idx_label}_{observation_label}.jpeg")
+            img.save(filename)
+            if generate_depth:
+
+                filename = os.path.join(
+                    output_dir, f"{idx_label}_{observation_label}_depth.exr"
+                )
+                cv2.imwrite(
+                    filename,
+                    observation["depth"],
+                    [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_HALF],
+                )
+
+                camera_params = dict(
+                    [
+                        (key, observation[key].tolist())
+                        for key in ("camera_intrinsics", "R_cam2world", "t_cam2world")
+                    ]
+                )
+                filename = os.path.join(
+                    output_dir, f"{idx_label}_{observation_label}_camera_params.json"
+                )
+                with open(filename, "w") as f:
+                    json.dump(camera_params, f)
+
+    with open(os.path.join(output_dir, "metadata.json"), "w") as f:
+        json.dump(metadata, f)
+
+    generator.close()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--metadata_filename", required=True)
+    parser.add_argument("--output_dir", required=True)
+    args = parser.parse_args()
+
+    generate_multiview_images_from_metadata(
+        metadata_filename=args.metadata_filename,
+        output_dir=args.output_dir,
+        scene_datasets_paths=SCENES_DATASET,
+        overload_params=dict(),
+        exist_ok=True,
+    )

longstream/utils/vendor/croco/datasets/habitat_sim/generate_from_metadata_files.py

@@ -0,0 +1,34 @@
+"""
+Script generating commandlines to generate image pairs from metadata files.
+"""
+import argparse
+import glob
+import os
+
+from tqdm import tqdm
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input_dir", required=True)
+    parser.add_argument("--output_dir", required=True)
+    parser.add_argument(
+        "--prefix",
+        default="",
+        help="Commanline prefix, useful e.g. to setup environment.",
+    )
+    args = parser.parse_args()
+
+    input_metadata_filenames = glob.iglob(
+        f"{args.input_dir}/**/metadata.json", recursive=True
+    )
+
+    for metadata_filename in tqdm(input_metadata_filenames):
+        output_dir = os.path.join(
+            args.output_dir,
+            os.path.relpath(os.path.dirname(metadata_filename), args.input_dir),
+        )
+
+        if os.path.exists(os.path.join(output_dir, "metadata.json")):
+            continue
+        commandline = f"{args.prefix}python datasets/habitat_sim/generate_from_metadata.py --metadata_filename={metadata_filename} --output_dir={output_dir}"
+        print(commandline)