From da6ff0858f4997f79f5d8cdc88b7367a5cb710b1 Mon Sep 17 00:00:00 2001
From: wu_xinjun <bent20140204@outlook.com>
Date: 星期二, 07 六月 2022 10:55:26 +0800
Subject: [PATCH] Merge commit 'mqdev' into mqaddFeature

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

diff --git a/src/utils/tools.py b/src/utils/tools.py
index 80b55d1..8f9fa00 100644
--- a/src/utils/tools.py
+++ b/src/utils/tools.py
@@ -1,3 +1,4 @@
+from email import header
 import os, sys
 import numpy as np
 import time
@@ -19,51 +20,63 @@
 
     return target_files
 
-def read_ply(path:str,start_line:int = 10):
-    count = 0
-    ply_list = []
+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
-            if count < start_line:
-                pass
+            flag = line.strip()
+            if  flag == "end_header":
+                return count + 1
             else:
-                line_str_list = line.strip().split(' ')
-                # convert the str type to float type
-                line_float_list = []
-                for value in line_str_list:
-                    line_float_list.append(float(value))
+                if count >= 50:
+                    return None
+                count += 1
+    return None
 
-                ply_list.append(line_float_list)
-            count += 1
-    ply_array = np.array(ply_list,dtype=np.float16)
-
+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")
-    with open (file,'w') as f:
-        f.write("ply\n")
-        f.write("format ascii 1.0\n")
-        f.write("comment write_ply generated\n")
-        f.write("element vertex %d\n" % data.shape[0])
-        f.write("property float x\n")
-        f.write("property float y\n")
-        f.write("property float z\n")
-        f.write("element face 0\n")
-        f.write("property list uchar int vertex_indices\n")
-        f.write("end_header\n")
-        n_p = 0
+    n_p = data.shape[0]
+    n_cols = data.shape[1]
 
-        for p in tqdm(data):
-            x = p[0]
-            y = p[1]
-            z = p[2]
-            f.write("{:.5f} {:.5f} {:.5f}\n".format(x, y, z))
-            n_p += 1
-        print(time.asctime(),f"Frame {name}.ply have been saved with {n_p} points!")
+    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=fmt)
+    print(f"{time.asctime()} {name}.ply have been saved with {n_p} points!")
 
 
 

--
Gitblit v1.9.1