feature: write to text files (closes #42)

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)
 

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)