From c302cc530185bec6c3cc41c8f6ba22df24f2e7ef Mon Sep 17 00:00:00 2001
From: Gunjan Chhablani <chhablani.gunjan@gmail.com>
Date: Thu, 11 Apr 2024 16:41:14 -0400
Subject: [PATCH] Add pointcloud loading for blender parser (#2902)

* 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>
---
 .../data/dataparsers/blender_dataparser.py    | 23 ++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/nerfstudio/data/dataparsers/blender_dataparser.py b/nerfstudio/data/dataparsers/blender_dataparser.py
index 64ccf13de..571799f1d 100644
--- a/nerfstudio/data/dataparsers/blender_dataparser.py
+++ b/nerfstudio/data/dataparsers/blender_dataparser.py
@@ -17,10 +17,11 @@ from __future__ import annotations
 
 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
@@ -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
@@ -94,12 +98,29 @@ class Blender(DataParser):
             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