Add a pointcloud export from polycam to process-data (#3124)

* add a pointcloud export from polycam mesh to process-data

* bugfix

nerfstudio/process_data/polycam_utils.py

@@ -17,7 +17,9 @@
 import json
 import sys
 from pathlib import Path
-from typing import List, Tuple
+from typing import List, Optional, Tuple
+
+import numpy as np
 
 from nerfstudio.process_data import process_data_utils
 from nerfstudio.process_data.process_data_utils import CAMERA_MODELS
@@ -28,6 +30,7 @@ from nerfstudio.utils.rich_utils import CONSOLE
 def polycam_to_json(
     image_filenames: List[Path],
     depth_filenames: List[Path],
+    glb_filename: Optional[Path],
     cameras_dir: Path,
     output_dir: Path,
     min_blur_score: float = 0.0,
@@ -80,6 +83,31 @@ def polycam_to_json(
         frames.append(frame)
     data["frames"] = frames
 
+    if glb_filename is not None:
+        # If the .glb is populated, use it to save a pointcloud for splatfacto init
+        import open3d as o3d
+
+        mesh = o3d.io.read_triangle_mesh(str(glb_filename), enable_post_processing=True)
+        textures = np.asarray(mesh.textures[0])  # 2D images of color
+        vert_points = np.asarray(mesh.vertices)  # 3D positions of verts
+        tri_ids = np.asarray(mesh.triangles)  # indices of the vertices
+        points = vert_points[tri_ids.flatten()]  # get the 3D positions of the vertices
+        uvs = np.asarray(mesh.triangle_uvs)  # get the uv coords of the vertices
+        # convert uv coord to texture integer index
+        tex_ids = (uvs[:, 1] * textures.shape[0]).astype(int), (uvs[:, 0] * textures.shape[1]).astype(int)
+        colors = textures[tex_ids[0], tex_ids[1]]
+        pointcloud = o3d.geometry.PointCloud(o3d.utility.Vector3dVector(points))
+        pointcloud.colors = o3d.utility.Vector3dVector(colors.astype(np.float64) / 255.0)
+        pointcloud = pointcloud.remove_duplicated_points()
+        # align the pointcloud to the coord system of cameras, which is provided inside the mesh_info.json file
+        mesh_info_json = json.load(open(glb_filename.parent / "mesh_info.json"))
+        transform = np.array(mesh_info_json["alignmentTransform"]).reshape(4, 4).T
+        pointcloud = pointcloud.transform(np.linalg.inv(transform))
+        # shift the axes coordinates to match the nerfstudio ones (same as the cameras' coord system)
+        pointcloud.points = o3d.utility.Vector3dVector(np.array(pointcloud.points)[:, [2, 0, 1]])
+        o3d.io.write_point_cloud(str(output_dir / "point_cloud.ply"), pointcloud)
+        data["ply_file_path"] = "point_cloud.ply"
+
     with open(output_dir / "transforms.json", "w", encoding="utf-8") as f:
         json.dump(data, f, indent=4)
 

nerfstudio/scripts/process_data.py

@@ -200,6 +200,7 @@ class ProcessPolycam(BaseConverterToNerfstudioDataset):
             polycam_utils.polycam_to_json(
                 image_filenames=polycam_image_filenames,
                 depth_filenames=polycam_depth_filenames,
+                glb_filename=self.data / "raw.glb" if (self.data / "raw.glb").exists() else None,
                 cameras_dir=polycam_cameras_dir,
                 output_dir=self.output_dir,
                 min_blur_score=self.min_blur_score,