update demo

README.md

@@ -83,6 +83,32 @@ pip install -e ".[vis]"
 
 # 🎬 Demo
 
+Run `demo.py` for interactive 3D visualization via a browser-based [viser](https://github.com/nerfstudio-project/viser) viewer (default `http://localhost:8080`).
+
+### Try the Example Scenes
+
+We provide three example scenes in `example/` that you can run out of the box:
+
+| Scene | Frames | Description |
+|:---|:---|:---|
+| `church` | 286 | Outdoor church with complex geometry |
+| `oxford` | 320 | Oxford street-level walkthrough |
+| `university4` | 324 | University campus outdoor scene |
+
+```bash
+# Church scene
+python demo.py --model_path /path/to/checkpoint.pt \
+    --image_folder example/church --mask_sky
+
+# Oxford scene with sky masking (outdoor)
+python demo.py --model_path /path/to/checkpoint.pt \
+    --image_folder example/oxford --mask_sky
+
+# University scene
+python demo.py --model_path /path/to/checkpoint.pt \
+    --image_folder example/university4 --mask_sky
+```
+
 ### Streaming Inference from Images
 
 ```bash
@@ -99,8 +125,7 @@ python demo.py --model_path /path/to/checkpoint.pt \
 
 ### Streaming with Keyframe Interval
 
-Use `--keyframe_interval` to reduce KV cache memory by only keeping every N-th frame as a keyframe. Non-keyframe frames still produce predictions but are not stored in the cache. This is useful for long sequences 
+Use `--keyframe_interval` to reduce KV cache memory by only keeping every N-th frame as a keyframe. Non-keyframe frames still produce predictions but are not stored in the cache. This is useful for long sequences which exceed 320 frames.
-which excesses 320 frames.
 
 ```bash
 python demo.py --model_path /path/to/checkpoint.pt \
@@ -108,6 +133,7 @@ python demo.py --model_path /path/to/checkpoint.pt \
 ```
 
 ### Windowed Inference (for long sequences, >3000 frames)
+
 ```bash
 python demo.py --model_path /path/to/checkpoint.pt \
     --video_path video.mp4 --fps 10 \
@@ -137,7 +163,23 @@ python demo.py --model_path /path/to/checkpoint.pt \
     --image_folder /path/to/images/ --mask_sky
 ```
 
-Sky masks are cached in `<image_folder>_sky_masks/` so subsequent runs skip regeneration.
+Sky masks are cached in `<image_folder>_sky_masks/` so subsequent runs skip regeneration. You can also specify a custom cache directory with `--sky_mask_dir`, or save side-by-side mask visualizations with `--sky_mask_visualization_dir`:
+
+```bash
+python demo.py --model_path /path/to/checkpoint.pt \
+    --image_folder /path/to/images/ --mask_sky \
+    --sky_mask_dir /path/to/cached_masks/ \
+    --sky_mask_visualization_dir /path/to/mask_viz/
+```
+
+### Visualization Options
+
+| Argument | Default | Description |
+|:---|:---|:---|
+| `--port` | `8080` | Viser viewer port |
+| `--conf_threshold` | `1.5` | Visibility threshold for filtering low-confidence points |
+| `--point_size` | `0.00001` | Point cloud point size |
+| `--downsample_factor` | `10` | Spatial downsampling for point cloud display |
 
 ### Without FlashInfer (SDPA fallback)
 

demo.py

@@ -39,7 +39,12 @@ from lingbot_map.utils.load_fn import load_and_preprocess_images
 
 def load_images(image_folder=None, video_path=None, fps=10, image_ext=".jpg,.png",
                 first_k=None, stride=1, image_size=518, patch_size=14, num_workers=8):
-    """Load images from folder or video and preprocess into a tensor."""
+    """Load images from folder or video and preprocess into a tensor.
+
+    Returns:
+        (images, paths, resolved_image_folder): preprocessed tensor, file paths,
+        and the folder containing the source images (for sky mask caching etc.).
+    """
     if video_path is not None:
         video_name = os.path.splitext(os.path.basename(video_path))[0]
         out_dir = os.path.join(os.path.dirname(video_path), f"{video_name}_frames")
@@ -63,6 +68,7 @@ def load_images(image_folder=None, video_path=None, fps=10, image_ext=".jpg,.png
         pbar.close()
         cap.release()
         paths = saved
+        resolved_folder = out_dir
         print(f"Extracted {len(paths)} frames from video ({total_frames} total, interval={interval})")
     else:
         exts = image_ext.split(",")
@@ -70,11 +76,12 @@ def load_images(image_folder=None, video_path=None, fps=10, image_ext=".jpg,.png
         for ext in exts:
             paths.extend(glob.glob(os.path.join(image_folder, f"*{ext}")))
         paths = sorted(paths)
+        resolved_folder = image_folder
 
-    if stride > 1:
-        paths = paths[::stride]
     if first_k is not None and first_k > 0:
         paths = paths[:first_k]
+    if stride > 1:
+        paths = paths[::stride]
 
     print(f"Loading {len(paths)} images...")
     images = load_and_preprocess_images(
@@ -85,7 +92,7 @@ def load_images(image_folder=None, video_path=None, fps=10, image_ext=".jpg,.png
     )
     h, w = images.shape[-2:]
     print(f"Preprocessed images to {w}x{h} using canonical crop mode")
-    return images, paths
+    return images, paths, resolved_folder
 
 
 # =============================================================================
@@ -261,8 +268,14 @@ def main():
     parser.add_argument("--port", type=int, default=8080)
     parser.add_argument("--conf_threshold", type=float, default=1.5)
     parser.add_argument("--downsample_factor", type=int, default=10)
-    parser.add_argument("--point_size", type=float, default=0.0007)
+    parser.add_argument("--point_size", type=float, default=0.00001)
     parser.add_argument("--mask_sky", action="store_true", help="Apply sky segmentation to filter out sky points")
+    parser.add_argument("--sky_mask_dir", type=str, default=None,
+                        help="Directory for cached sky masks (default: <image_folder>_sky_masks/)")
+    parser.add_argument("--sky_mask_visualization_dir", type=str, default=None,
+                        help="Save sky mask visualizations (original | mask | overlay) to this directory")
+    parser.add_argument("--export_preprocessed", type=str, default=None,
+                        help="Export stride-sampled, resized/cropped images to this folder")
 
     args = parser.parse_args()
     assert args.image_folder or args.video_path, \
@@ -272,11 +285,24 @@ def main():
 
     # ── Load images & model ──────────────────────────────────────────────────
     t0 = time.time()
-    images, paths = load_images(
+    images, paths, resolved_image_folder = load_images(
         image_folder=args.image_folder, video_path=args.video_path,
         fps=args.fps, first_k=args.first_k, stride=args.stride,
         image_size=args.image_size, patch_size=args.patch_size,
     )
+
+    # Export preprocessed images if requested
+    if args.export_preprocessed:
+        os.makedirs(args.export_preprocessed, exist_ok=True)
+        print(f"Exporting {images.shape[0]} preprocessed images to {args.export_preprocessed}...")
+        for i in range(images.shape[0]):
+            img = (images[i].permute(1, 2, 0).numpy() * 255).clip(0, 255).astype(np.uint8)
+            cv2.imwrite(
+                os.path.join(args.export_preprocessed, f"{i:06d}.png"),
+                cv2.cvtColor(img, cv2.COLOR_RGB2BGR),
+            )
+        print(f"Exported to {args.export_preprocessed}")
+
     model = load_model(args, device)
     print(f"Total load time: {time.time() - t0:.1f}s")
 
@@ -330,6 +356,9 @@ def main():
             downsample_factor=args.downsample_factor,
             point_size=args.point_size,
             mask_sky=args.mask_sky,
+            image_folder=resolved_image_folder,
+            sky_mask_dir=args.sky_mask_dir,
+            sky_mask_visualization_dir=args.sky_mask_visualization_dir,
         )
         print(f"3D viewer at http://localhost:{args.port}")
         viewer.run()