diff --git a/nerfstudio/process_data/polycam_utils.py b/nerfstudio/process_data/polycam_utils.py index 0b1ba7a62..a84de5a31 100644 --- a/nerfstudio/process_data/polycam_utils.py +++ b/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) diff --git a/nerfstudio/scripts/process_data.py b/nerfstudio/scripts/process_data.py index 0ab2bc8f0..643587298 100644 --- a/nerfstudio/scripts/process_data.py +++ b/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,