feat: Add featureless visualization utils

app.py

@@ -16,7 +16,7 @@ from transformers import AutoModelForImageSegmentation
 from torchvision import transforms
 from PIL import Image
 import open3d as o3d
-from spann3r.tools.vis import render_frames
+from vis_utils import render_frames
 from backend_utils import improved_multiway_registration, pts2normal, point2mesh, combine_and_clean_point_clouds
 from gs_utils import point2gs
 from pose_utils import solve_cemara

vis_utils.py

@@ -0,0 +1,171 @@
+import os
+import cv2
+import imageio
+import numpy as np
+import open3d as o3d
+import os.path as osp
+import matplotlib.pyplot as plt
+import matplotlib.colors as mcolors
+
+def render_frames(o3d_geometry, camera_all, output_dir, save_video=True, save_camera=True):
+    # Create off-screen renderer
+    render = o3d.visualization.rendering.OffscreenRenderer(
+        width=camera_all[0].intrinsic.width,
+        height=camera_all[0].intrinsic.height
+    )
+    
+    render_frame_path = os.path.join(output_dir, 'render_frames')
+    render_camera_path = os.path.join(output_dir, 'render_cameras')
+    os.makedirs(render_frame_path, exist_ok=True)
+    os.makedirs(render_camera_path, exist_ok=True)
+
+    video_path = os.path.join(output_dir, 'render_frame.mp4')
+    if save_video:
+        writer = imageio.get_writer(video_path, fps=10)
+
+    material = o3d.visualization.rendering.MaterialRecord()
+    material.shader = 'defaultUnlit'  # Use unlit shader for point clouds
+    material.point_size = 1.0  # Match original point size
+    material.base_color = [1.0, 1.0, 1.0, 1.0] 
+    
+    for i, camera_params in enumerate(camera_all):
+        if camera_params is None:
+            continue
+            
+        # Set camera view
+        render.setup_camera(
+            camera_params.intrinsic.intrinsic_matrix,
+            camera_params.extrinsic,
+            camera_params.intrinsic.width,
+            camera_params.intrinsic.height
+        )
+
+        if save_camera:
+            o3d.io.write_pinhole_camera_parameters(
+                os.path.join(render_camera_path, f'camera_{i:03d}.json'),
+                camera_params
+            )
+
+        # Render
+        render.scene.add_geometry("points", o3d_geometry, material)
+        img = render.render_to_image()
+        render.scene.remove_geometry("points")
+
+        # Save frame
+        image_uint8 = (np.asarray(img) * 255).astype(np.uint8)
+        frame_filename = f'frame_{i:03d}.png'
+        imageio.imwrite(osp.join(render_frame_path, frame_filename), image_uint8)
+
+        if save_video:
+            writer.append_data(image_uint8)
+
+    if save_video:
+        writer.close()
+
+    return video_path
+
+def find_render_cam(pcd, width=1920, height=1080):
+    # For headless servers, we'll need to pre-define camera parameters
+    # This creates a default viewing angle looking at the center of the point cloud
+    
+    # Calculate point cloud center and scale
+    center = pcd.get_center()
+    scale = np.max(pcd.get_max_bound() - pcd.get_min_bound())
+    
+    # Create default camera parameters
+    camera_params = o3d.camera.PinholeCameraParameters()
+    
+    # Set intrinsic parameters
+    intrinsic = o3d.camera.PinholeCameraIntrinsic()
+    intrinsic.set_intrinsics(
+        width=width,
+        height=height,
+        fx=width,
+        fy=width,
+        cx=width/2,
+        cy=height/2
+    )
+    camera_params.intrinsic = intrinsic
+    
+    # Set extrinsic parameters (looking at center from a 45-degree angle)
+    camera_params.extrinsic = np.array([
+        [1, 0, 0, 0],
+        [0, np.cos(np.pi/4), -np.sin(np.pi/4), 0],
+        [0, np.sin(np.pi/4), np.cos(np.pi/4), 2*scale],
+        [0, 0, 0, 1]
+    ])
+    
+    return camera_params
+
+def vis_pred_and_imgs(pts_all, save_path, images_all=None, conf_all=None, save_video=True):
+    # Set matplotlib backend to non-interactive
+    plt.switch_backend('Agg')
+    
+    # Normalization
+    min_val = pts_all.min(axis=(0, 1, 2), keepdims=True)
+    max_val = pts_all.max(axis=(0, 1, 2), keepdims=True)
+    pts_all = (pts_all - min_val) / (max_val - min_val)
+
+    pts_save_path = osp.join(save_path, 'pts')
+    os.makedirs(pts_save_path, exist_ok=True)
+
+    if images_all is not None:
+        images_save_path = osp.join(save_path, 'imgs')
+        os.makedirs(images_save_path, exist_ok=True)
+    
+    if conf_all is not None:
+        conf_save_path = osp.join(save_path, 'confs')
+        os.makedirs(conf_save_path, exist_ok=True)
+
+    if save_video:
+        pts_video_path = osp.join(save_path, 'pts.mp4')
+        pts_writer = imageio.get_writer(pts_video_path, fps=10)
+
+        if images_all is not None:
+            imgs_video_path = osp.join(save_path, 'imgs.mp4')
+            imgs_writer = imageio.get_writer(imgs_video_path, fps=10)
+
+        if conf_all is not None:
+            conf_video_path = osp.join(save_path, 'confs.mp4')
+            conf_writer = imageio.get_writer(conf_video_path, fps=10)
+    
+    for frame_id in range(pts_all.shape[0]):
+        # Points visualization
+        pt_vis = pts_all[frame_id].astype(np.float32)
+        pt_vis_rgb = mcolors.hsv_to_rgb(1-pt_vis)
+        pt_vis_rgb_uint8 = (pt_vis_rgb * 255).astype(np.uint8)
+
+        # Use matplotlib in non-interactive mode
+        fig, ax = plt.subplots()
+        ax.imshow(pt_vis_rgb_uint8)
+        plt.savefig(osp.join(pts_save_path, f'pts_{frame_id:04d}.png'))
+        plt.close(fig)
+
+        if save_video:
+            pts_writer.append_data(pt_vis_rgb_uint8)
+
+        if images_all is not None:
+            image = images_all[frame_id]
+            image_uint8 = (image * 255).astype(np.uint8)
+            imageio.imwrite(osp.join(images_save_path, f'img_{frame_id:04d}.png'), image_uint8)
+
+            if save_video:
+                imgs_writer.append_data(image_uint8)
+        
+        if conf_all is not None:
+            fig, ax = plt.subplots()
+            conf_image = plt.cm.jet(conf_all[frame_id])
+            ax.imshow(conf_image)
+            plt.savefig(osp.join(conf_save_path, f'conf_{frame_id:04d}.png'))
+            plt.close(fig)
+
+            conf_image_uint8 = (conf_image * 255).astype(np.uint8)
+            if save_video:
+                conf_writer.append_data(conf_image_uint8)
+    
+    if save_video:
+        pts_writer.close()
+        if images_all is not None:
+            imgs_writer.close()
+        if conf_all is not None:
+            conf_writer.close()