From cddaf02bb0f251b0394d3c3f7896d9cf49f064be Mon Sep 17 00:00:00 2001
From: Graham Knapp <graham.knapp@acernis.fr>
Date: Tue, 27 Dec 2022 09:51:43 +0100
Subject: [PATCH 1/2] feature: write to text files

---
 point_e/util/point_cloud.py | 35 +++++++++++++++++++++++++++++------
 point_e/util/txt_util.py    | 23 +++++++++++++++++++++++
 2 files changed, 52 insertions(+), 6 deletions(-)
 create mode 100644 point_e/util/txt_util.py

diff --git a/point_e/util/point_cloud.py b/point_e/util/point_cloud.py
index 424f67d..453dc24 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,28 @@ 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, "wb") as writer:
+                write_txt(
+                    writer,
+                    coords=self.coords,
+                    rgb=rgb,
+                )
+        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 +103,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 +127,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 +179,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)

From a701727d42d0ee58d413fc6297c1d9693328e13f Mon Sep 17 00:00:00 2001
From: Graham Knapp <graham.knapp@acernis.fr>
Date: Tue, 27 Dec 2022 14:44:42 +0100
Subject: [PATCH 2/2] feature: write to text files

---
 point_e/util/point_cloud.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/point_e/util/point_cloud.py b/point_e/util/point_cloud.py
index 453dc24..f2b34b3 100644
--- a/point_e/util/point_cloud.py
+++ b/point_e/util/point_cloud.py
@@ -76,17 +76,18 @@ def write_txt(self, f: Union[str, TextIO]):
             else None
         )
         if isinstance(f, str):
-            with open(f, "wb") as writer:
+            with open(f, "w") as writer:
                 write_txt(
                     writer,
                     coords=self.coords,
                     rgb=rgb,
                 )
-        write_txt(
-            f,
-            coords=self.coords,
-            rgb=rgb,
-        )
+        else:
+            write_txt(
+                f,
+                coords=self.coords,
+                rgb=rgb,
+            )
 
     def random_sample(self, num_points: int, **subsample_kwargs) -> "PointCloud":
         """