diff --git a/point_e/util/point_cloud.py b/point_e/util/point_cloud.py index 424f67d..f2b34b3 100644 --- a/point_e/util/point_cloud.py +++ b/point_e/util/point_cloud.py @@ -1,10 +1,11 @@ import random from dataclasses import dataclass -from typing import BinaryIO, Dict, List, Optional, Union +from typing import BinaryIO, Dict, List, Optional, Union, TextIO import numpy as np from .ply_util import write_ply +from .txt_util import write_txt COLORS = frozenset(["R", "G", "B", "A"]) @@ -65,6 +66,29 @@ def write_ply(self, raw_f: BinaryIO): ), ) + def write_txt(self, f: Union[str, TextIO]): + """ + Save the point cloud to a .txt file. + """ + rgb = ( + np.stack([self.channels[x] for x in "RGB"], axis=1 + ) if all(x in self.channels for x in "RGB") + else None + ) + if isinstance(f, str): + with open(f, "w") as writer: + write_txt( + writer, + coords=self.coords, + rgb=rgb, + ) + else: + write_txt( + f, + coords=self.coords, + rgb=rgb, + ) + def random_sample(self, num_points: int, **subsample_kwargs) -> "PointCloud": """ Sample a random subset of this PointCloud. @@ -80,7 +104,7 @@ def random_sample(self, num_points: int, **subsample_kwargs) -> "PointCloud": return self.subsample(indices, **subsample_kwargs) def farthest_point_sample( - self, num_points: int, init_idx: Optional[int] = None, **subsample_kwargs + self, num_points: int, init_idx: Optional[int] = None, **subsample_kwargs ) -> "PointCloud": """ Sample a subset of the point cloud that is evenly distributed in space. @@ -104,7 +128,7 @@ def farthest_point_sample( init_idx = random.randrange(len(self.coords)) if init_idx is None else init_idx indices = np.zeros([num_points], dtype=np.int64) indices[0] = init_idx - sq_norms = np.sum(self.coords**2, axis=-1) + sq_norms = np.sum(self.coords ** 2, axis=-1) def compute_dists(idx: int): # Utilize equality: ||A-B||^2 = ||A||^2 + ||B||^2 - 2*(A @ B). @@ -156,11 +180,11 @@ def nearest_points(self, points: np.ndarray, batch_size: int = 16384) -> np.ndar make the computation faster. :return: an [N] array of indices into self.coords. """ - norms = np.sum(self.coords**2, axis=-1) + norms = np.sum(self.coords ** 2, axis=-1) all_indices = [] for i in range(0, len(points), batch_size): - batch = points[i : i + batch_size] - dists = norms + np.sum(batch**2, axis=-1)[:, None] - 2 * (batch @ self.coords.T) + batch = points[i: i + batch_size] + dists = norms + np.sum(batch ** 2, axis=-1)[:, None] - 2 * (batch @ self.coords.T) all_indices.append(np.argmin(dists, axis=-1)) return np.concatenate(all_indices, axis=0) diff --git a/point_e/util/txt_util.py b/point_e/util/txt_util.py new file mode 100644 index 0000000..1998777 --- /dev/null +++ b/point_e/util/txt_util.py @@ -0,0 +1,23 @@ +from typing import TextIO, Optional + +import numpy as np + + +def write_txt( + f: TextIO, + coords: np.ndarray, + rgb: Optional[np.ndarray] = None, +): + """ + Write a text file for a point cloud. + + :param f: an i/o stream to write to, such as the stream returned by open() + :param coords: an [N x 3] array of floating point coordinates. + :param rgb: an [N x 3] array of vertex colors, in the range [0.0, 1.0]. + + """ + points = coords + if rgb is not None: + rgb = (rgb * 255.499).round().astype(int) + points = np.concatenate([points, rgb], axis=1) + np.savetxt(f, points)