fix: Refactor visualization

spann3r/tools/vis.py

@@ -14,41 +14,31 @@ import numpy as np
 import imageio
 import os.path as osp
 
-def render_frames(pts_all, image_all, camera_parameters, output_dir, mask=None, save_video=True, save_camera=True):
-    t, h, w, _ = pts_all.shape
-
+def render_frames(o3d_geometry, camera_all, output_dir, save_video=True, save_camera=True):
     vis = o3d.visualization.Visualizer()
-    vis.create_window(width=1920, height=1080)
+    w, h = camera_all[0].intrinsic.width, camera_all[0].intrinsic.height
+    vis.create_window(width=w, height=h)
 
     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)
-
-    if save_camera:
-        o3d.io.write_pinhole_camera_parameters(os.path.join(render_frame_path, 'camera.json'), camera_parameters)
+    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)
 
-    pcd = o3d.geometry.PointCloud()
-    vis.add_geometry(pcd)
-
-    for i in range(t):
-        new_pts = pts_all[i].reshape(-1, 3)
-        new_colors = image_all[i].reshape(-1, 3)
-
-        if mask is not None:
-            new_pts = new_pts[mask[i].reshape(-1)]
-            new_colors = new_colors[mask[i].reshape(-1)]
-
-        pcd.points.extend(o3d.utility.Vector3dVector(new_pts))
-        pcd.colors.extend(o3d.utility.Vector3dVector(new_colors))
-
-        vis.clear_geometries()
-        vis.add_geometry(pcd)
+    vis.add_geometry(o3d_geometry)
 
+    for i, camera_params in enumerate(camera_all):
+        if camera_params is None:
+            continue
         ctr = vis.get_view_control()
-        ctr.convert_from_pinhole_camera_parameters(camera_parameters)
+        ctr.convert_from_pinhole_camera_parameters(camera_params, True)
+        
+        if save_camera:
+            o3d.io.write_pinhole_camera_parameters(os.path.join(render_camera_path, f'camera_{i:03d}.json'), 
+                                                   camera_params)
 
         opt = vis.get_render_option()
         opt.point_size = 1
@@ -70,7 +60,7 @@ def render_frames(pts_all, image_all, camera_parameters, output_dir, mask=None,
         writer.close()
 
     vis.destroy_window()
-
+    return video_path
 
 
 def find_render_cam(pcd):