feat: Add gs_utils for gs export

app.py

@@ -17,7 +17,7 @@ from torchvision import transforms
 from PIL import Image
 import open3d as o3d
 from backend_utils import improved_multiway_registration, pts2normal, point2mesh, combine_and_clean_point_clouds
-
+from gs_utils import point2gs
 
 # Default values
 DEFAULT_CKPT_PATH = './checkpoints/spann3r.pth'
@@ -29,15 +29,8 @@ OPENGL = np.array([[1, 0, 0, 0],
                    [0, 0, -1, 0],
                    [0, 0, 0, 1]])
 
-def export_geometry(geometry, as_pointcloud=False):
-    if as_pointcloud:
-        if not isinstance(geometry, o3d.geometry.PointCloud):
-            raise ValueError("Expected an Open3D PointCloud object when as_pointcloud is True")
-        output_path = tempfile.mktemp(suffix='.ply')
-    else:
-        if not isinstance(geometry, o3d.geometry.TriangleMesh):
-            raise ValueError("Expected an Open3D TriangleMesh object when as_pointcloud is False")
-        output_path = tempfile.mktemp(suffix='.obj')
+def export_geometry(geometry):
+    output_path = tempfile.mktemp(suffix='.obj')
 
     # Apply rotation
     rot = np.eye(4)
@@ -45,11 +38,7 @@ def export_geometry(geometry, as_pointcloud=False):
     transform = np.linalg.inv(OPENGL @ rot)
     geometry.transform(transform)
 
-    # Export the geometry
-    if as_pointcloud:
-        o3d.io.write_point_cloud(output_path, geometry, write_ascii=False, compressed=True)
-    else:
-        o3d.io.write_triangle_mesh(output_path, geometry, write_ascii=False, compressed=True)
+    o3d.io.write_triangle_mesh(output_path, geometry, write_ascii=False, compressed=True)
     
     return output_path
 
@@ -176,7 +165,7 @@ def generate_mask(image: np.ndarray):
     return mask_np
 @torch.no_grad()
 def reconstruct(video_path, conf_thresh, kf_every, 
-                as_pointcloud=False, remove_background=False, refine=False):
+                remove_background=False):
     # Extract frames from video
     demo_path = extract_frames(video_path)
     
@@ -220,31 +209,21 @@ def reconstruct(video_path, conf_thresh, kf_every,
         pcds.append(pcd)
     
     pcd_combined = combine_and_clean_point_clouds(pcds, voxel_size=0.001)
-    
-    if as_pointcloud:
-        o3d_geometry = pcd_combined
-    else:
-        o3d_geometry = point2mesh(pcd_combined)
+    o3d_geometry = point2mesh(pcd_combined)
         
     # Create coarse result
-    coarse_output_path = export_geometry(o3d_geometry, as_pointcloud)
+    coarse_output_path = export_geometry(o3d_geometry)
     
     yield coarse_output_path, None
     
-    if refine:
-        # Perform global optimization 
-        print("Performing global registration...")
-        transformed_pcds, _, _ = improved_multiway_registration(pcds, voxel_size=0.001)
-        
-        if as_pointcloud:
-            o3d_geometry = transformed_pcds
-        else:
-            o3d_geometry = point2mesh(transformed_pcds)
-            
-        # Create coarse result
-        refined_output_path = export_geometry(o3d_geometry, as_pointcloud)
-        
-        yield coarse_output_path, refined_output_path
+    # Perform global optimization 
+    print("Performing global registration...")
+    transformed_pcds, _, _ = improved_multiway_registration(pcds, voxel_size=0.01)
+
+    # Create coarse result
+    refined_output_path = tempfile.mktemp(suffix='.ply')
+    point2gs(refined_output_path, transformed_pcds)
+    yield coarse_output_path, refined_output_path
 
     # Clean up temporary directory
     os.system(f"rm -rf {demo_path}")
@@ -320,19 +299,19 @@ with gr.Blocks(
                 kf_every = gr.Slider(1, 30, step=1, value=1, label="Keyframe Interval")
             with gr.Row():
                 remove_background = gr.Checkbox(label="Remove Background", value=False)
-                refine = gr.Checkbox(label="Enable Backend", value=False)
-                as_pointcloud = gr.Checkbox(label="As Pointcloud", value=False)
             reconstruct_btn = gr.Button("Reconstruct")
         
         with gr.Column(scale=2):
             with gr.Tab("Coarse Model"):
-                coarse_model = gr.Model3D(label="Coarse 3D Model", display_mode="solid", clear_color=[0.0, 0.0, 0.0, 0.0])
+                coarse_model = gr.Model3D(label="Coarse 3D Model", display_mode="solid", 
+                                          clear_color=[0.0, 0.0, 0.0, 0.0])
             with gr.Tab("Refined Model"):
-                refined_model = gr.Model3D(label="Refined 3D Model", display_mode="solid", clear_color=[0.0, 0.0, 0.0, 0.0])
+                refined_model = gr.Model3D(label="Refined Gaussian Splatting", display_mode="solid", 
+                                           clear_color=[0.0, 0.0, 0.0, 0.0])
     
     reconstruct_btn.click(
         fn=reconstruct,
-        inputs=[video_input, conf_thresh, kf_every, as_pointcloud, remove_background, refine],
+        inputs=[video_input, conf_thresh, kf_every, remove_background],
         outputs=[coarse_model, refined_model]
     )
 

gs_utils.py

@@ -0,0 +1,106 @@
+import numpy as np
+import os
+from plyfile import PlyElement, PlyData
+import open3d as o3d
+
+def get_f_dc(colors):
+    return RGB2SH(colors)[:, :, np.newaxis]
+
+def get_f_rest(points, max_sh_degree=3):
+    f_rest_shape = (points.shape[0], (max_sh_degree + 1) ** 2 - 1, 3)
+    return np.zeros(f_rest_shape)
+
+def get_opacity(points):
+    return inverse_sigmoid(0.5 * np.ones((points.shape[0], 1)))
+
+def get_scales(points):
+    scales = np.ones((points.shape[0], 3)) * 0.0015
+    scales[:, 2] = 1e-6
+    
+    return np.log(scales)
+
+def get_rotation(normals):
+    if normals is not None and np.any(normals):
+        return normal2rotation(normals)
+    else:
+        return np.zeros((normals.shape[0], 4))
+
+def RGB2SH(rgb):
+    return (rgb - 0.5) / 0.28209479177387814
+
+def inverse_sigmoid(x):
+    return np.log(x / (1 - x))
+
+def normal2rotation(n):
+    n = n / np.linalg.norm(n, axis=1, keepdims=True)
+    w0 = np.tile([[1, 0, 0]], (n.shape[0], 1))
+    R0 = w0 - np.sum(w0 * n, axis=1, keepdims=True) * n
+    R0 *= np.sign(R0[:, :1])
+    R0 /= np.linalg.norm(R0, axis=1, keepdims=True)
+    R1 = np.cross(n, R0)
+    R1 *= np.sign(R1[:, 1:2]) * np.sign(n[:, 2:])
+    R = np.stack([R0, R1, n], axis=-1)
+    q = rotmat2quaternion(R)
+    return q
+
+def rotmat2quaternion(R, normalize=False):
+    tr = R[:, 0, 0] + R[:, 1, 1] + R[:, 2, 2] + 1e-6
+    r = np.sqrt(1 + tr) / 2
+    q = np.stack([
+        r,
+        (R[:, 2, 1] - R[:, 1, 2]) / (4 * r),
+        (R[:, 0, 2] - R[:, 2, 0]) / (4 * r),
+        (R[:, 1, 0] - R[:, 0, 1]) / (4 * r)
+    ], axis=-1)
+    if normalize:
+        q /= np.linalg.norm(q, axis=-1, keepdims=True)
+    return q
+
+def point2gs(path, pcd, scale=None, max_sh_degree=1):
+    # Ensure the directory exists
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+
+    # Get point cloud data
+    xyz = np.asarray(pcd.points)
+    normals = np.asarray(pcd.normals) if pcd.has_normals() else np.zeros_like(xyz)
+    colors = np.asarray(pcd.colors) if pcd.has_colors() else np.ones_like(xyz)
+
+    # Generate additional attributes
+    f_dc = get_f_dc(colors).reshape(xyz.shape[0], -1)
+    f_rest = get_f_rest(xyz, max_sh_degree).reshape(xyz.shape[0], -1)
+    opacities = get_opacity(xyz)
+    if scale is not None:
+        scale = np.log(scale)
+    else:
+        scale = get_scales(xyz)
+    rotation = get_rotation(normals)
+    
+    # Construct list of attributes
+    attribute_names = ['x', 'y', 'z', 'nx', 'ny', 'nz']
+    attribute_names.extend([f'f_dc_{i}' for i in range(f_dc.shape[-1])])
+    attribute_names.extend([f'f_rest_{i}' for i in range(f_rest.shape[-1])])
+    attribute_names.append('opacity')
+    attribute_names.extend([f'scale_{i}' for i in range(scale.shape[1])])
+    attribute_names.extend([f'rot_{i}' for i in range(rotation.shape[1])])
+
+    # Create dtype for numpy structured array
+    dtype_full = [(attribute, 'f4') for attribute in attribute_names]
+
+    # Combine all attributes
+    attributes = np.concatenate((
+        xyz, normals, 
+        f_dc,
+        f_rest,
+        opacities, scale, rotation
+    ), axis=1)
+    
+    # Ensure attributes match the dtype
+    assert attributes.shape[1] == len(dtype_full), f"Mismatch in attribute count. Expected {len(dtype_full)}, got {attributes.shape[1]}"
+
+    # Create structured array
+    elements = np.empty(xyz.shape[0], dtype=dtype_full)
+    elements[:] = list(map(tuple, attributes))
+
+    # Create PlyElement and save
+    el = PlyElement.describe(elements, 'vertex')
+    PlyData([el]).write(path)

requirements.txt

@@ -18,4 +18,5 @@ transformers
 kornia
 timm
 numpy==1.26.4
-open3d
+open3d
+plyfile