From 13530d0ad2c32608d0b46a5c16ad4ac0c25b59f1 Mon Sep 17 00:00:00 2001
From: wu_xinjun <bent20140204@outlook.com>
Date: 星期五, 13 五月 2022 14:08:50 +0800
Subject: [PATCH] support for 6 columns ply

---
 src/utils/tools.py |   58 ++++++++++++++++++++++++++++++++++++++++++++++------------
 1 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/src/utils/tools.py b/src/utils/tools.py
index ee2806a..8f9fa00 100644
--- a/src/utils/tools.py
+++ b/src/utils/tools.py
@@ -20,28 +20,62 @@
 
     return target_files
 
-def read_ply(path:str,start_line:int = 10):
-    ply_array = np.load(path,skiprows=start_line)
+def get_skip_rows(path):
+    with open(path,"r") as ply_file:
+        count = 0
+        # read the line
+        for line in ply_file.readlines():
+            # strip the str of the header
+            flag = line.strip()
+            if  flag == "end_header":
+                return count + 1
+            else:
+                if count >= 50:
+                    return None
+                count += 1
+    return None
+
+def read_ply(path:str):
+    start_line = get_skip_rows(path)
+    if start_line is not None:
+        ply_array = np.loadtxt(path,skiprows=start_line)
+    else:
+        ply_array = None
     return ply_array
 
 
 def write_ply(data:np.array,folder:str,name:str):
     file = os.path.join(folder,name+".ply")
     n_p = data.shape[0]
-    header = ["ply\nformat ascii 1.0\n",
-        "comment write_ply generated\n",
-        f"element vertex {n_p}\n",
-        "property float x\n",
-        "property float y\n",
-        "property float z\n",
-        "element face 0\n",
-        "property list uchar int vertex_indices\n",
-        "end_header"]
+    n_cols = data.shape[1]
+
+    if n_cols == 3:
+        header = ["ply\nformat ascii 1.0\n",
+            "comment write_ply generated\n",
+            f"element vertex {n_p}\n",
+            "property float x\n",
+            "property float y\n",
+            "property float z\n",
+            "end_header"]
+        fmt = '%.6f','%.6f','%.6f'
+
+    elif n_cols == 6:
+        header = ["ply\nformat ascii 1.0\n",
+            "comment write_ply generated\n",
+            f"element vertex {n_p}\n",
+            "property float x\n",
+            "property float y\n",
+            "property float z\n",
+            "property uchar red\n",
+            "property uchar green\n",
+            "property uchar blue\n",
+            "end_header"]
+        fmt = '%.6f','%.6f','%.6f','%d','%d','%d'
     np.savetxt(file,data,
                 delimiter=" ",
                 header="".join(header),
                 comments="",
-                fmt='%.6f')
+                fmt=fmt)
     print(f"{time.asctime()} {name}.ply have been saved with {n_p} points!")
 
 

--
Gitblit v1.9.1