fix: Replace render output with point cloud

app.py

@@ -16,7 +16,6 @@ from transformers import AutoModelForImageSegmentation
 from torchvision import transforms
 from PIL import Image
 import open3d as o3d
-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
@@ -306,23 +305,19 @@ def reconstruct(video_path, conf_thresh, kf_every,
     
     # Create coarse result
     coarse_output_path = export_geometry(o3d_geometry_centered)
-    yield coarse_output_path, None
 
-    gs_output_path = tempfile.mktemp(suffix='.ply')
     if enable_registration:
-        transformed_pcds, _, _ = improved_multiway_registration(pcds, voxel_size=0.01)
-        transformed_pcds = center_pcd(transformed_pcds)
-        point2gs(gs_output_path, transformed_pcds)
-    else:
-        point2gs(gs_output_path, pcd_combined)
+        pcd_combined, _, _ = improved_multiway_registration(pcds, voxel_size=0.01)
+        pcd_combined = center_pcd(pcd_combined)
 
     if output_3d_model:
+        gs_output_path = tempfile.mktemp(suffix='.ply')
+        point2gs(gs_output_path, pcd_combined)
         # Create 3D model result using gaussian splatting
-        yield coarse_output_path, gs_output_path
+        return coarse_output_path, gs_output_path
     else:
-        gs_output_path = tempfile.mktemp(suffix='.ply')
-        render_video_path = render_frames(o3d_geometry, cameras_all, demo_path)
-        yield coarse_output_path, render_video_path
+        pcd_output_path = export_geometry(pcd_combined)
+        return coarse_output_path, pcd_output_path
 
     # Clean up temporary directory
     os.system(f"rm -rf {demo_path}")
@@ -406,7 +401,7 @@ with gr.Blocks(
                 output_3d_model = gr.Checkbox(
                     label="Output Splat", 
                     value=True,
-                    info="Generate Splat (PLY) instead of video render"
+                    info="Generate Splat (PLY) instead of Point Cloud (PLY)"
                 )
             reconstruct_btn = gr.Button("Start Reconstruction")
         
@@ -414,33 +409,16 @@ with gr.Blocks(
             with gr.Tab("3D Models"):
                 with gr.Group():
                     initial_model = gr.Model3D(
-                        label="Initial 3D Model", 
+                        label="Reconstructed Mesh", 
                         display_mode="solid",
                         clear_color=[0.0, 0.0, 0.0, 0.0]
                     )
-                    gr.Markdown(
-                        """
-                        <div class="model-description">
-                        This is the initial 3D model generated from the video. Finish within 10 seconds.
-                        </div>
-                        """
-                    )
                 
                 with gr.Group():
-                    output_model = gr.File(
-                        label="Refined Result (Splat or Video)", 
-                        file_types=[".ply", ".mp4"],
-                        file_count="single"
-                    )
-                    gr.Markdown(
-                        """
-                        <div class="model-description">
-                        Downloads as either:
-                        - PLY file: Gaussin Splat Model (when "Output Splat" is enabled)
-                        - MP4 file: 360° rotating render video (when "Output Splat" is disabled)
-                        <br>Time: ~60 seconds with refinement, ~30 seconds without
-                        </div>
-                        """
+                    initial_model = gr.Model3D(
+                        label="Reconstructed PointCloud or Splat", 
+                        display_mode="solid",
+                        clear_color=[0.0, 0.0, 0.0, 0.0]
                     )
     
     Examples(

vis_utils.py

@@ -1,171 +0,0 @@
-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()