Skip to content

Commit

Permalink
Add pointcloud loading for blender parser (nerfstudio-project#2902)
Browse files Browse the repository at this point in the history
* Add pointcloud loading for blender parser

* Fix argument

* Fix docstring

* Fix metadata in dataparsers

* Add changes from review

* Avoid changing default behavior

* Apply `scale_factor` to loaded points

---------

Co-authored-by: Brent Yi <yibrenth@gmail.com>
  • Loading branch information
gchhablani and brentyi committed Apr 11, 2024
1 parent 03bdeff commit c302cc5
Showing 1 changed file with 22 additions and 1 deletion.
23 changes: 22 additions & 1 deletion nerfstudio/data/dataparsers/blender_dataparser.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,11 @@

from dataclasses import dataclass, field
from pathlib import Path
from typing import Type
from typing import Optional, Type

import imageio
import numpy as np
import open3d as o3d
import torch

from nerfstudio.cameras.cameras import Cameras, CameraType
Expand All @@ -42,6 +43,9 @@ class BlenderDataParserConfig(DataParserConfig):
"""How much to scale the camera origins by."""
alpha_color: str = "white"
"""alpha color of background"""
ply_path: Optional[Path] = None
"""Path to PLY file to load 3D points from, defined relative to the dataset directory. This is helpful for
Gaussian splatting and generally unused otherwise. If `None`, points are initialized randomly."""


@dataclass
Expand Down Expand Up @@ -94,12 +98,29 @@ def _generate_dataparser_outputs(self, split="train"):
camera_type=CameraType.PERSPECTIVE,
)

metadata = {}
if self.config.ply_path is not None:
metadata.update(self._load_3D_points(self.config.data / self.config.ply_path))

dataparser_outputs = DataparserOutputs(
image_filenames=image_filenames,
cameras=cameras,
alpha_color=self.alpha_color_tensor,
scene_box=scene_box,
dataparser_scale=self.scale_factor,
metadata=metadata,
)

return dataparser_outputs

def _load_3D_points(self, ply_file_path: Path):
pcd = o3d.io.read_point_cloud(str(ply_file_path))

points3D = torch.from_numpy(np.asarray(pcd.points, dtype=np.float32) * self.config.scale_factor)
points3D_rgb = torch.from_numpy((np.asarray(pcd.colors) * 255).astype(np.uint8))

out = {
"points3D_xyz": points3D,
"points3D_rgb": points3D_rgb,
}
return out

0 comments on commit c302cc5

Please sign in to comment.