From 3a9751dc724877933fc883320100796cef23489d Mon Sep 17 00:00:00 2001
From: "Xinliang.Lu" <1209425127@qq.com>
Date: Tue, 23 Mar 2021 16:01:55 +0800
Subject: [PATCH] Update preprocess.py

Improve txt reading efficiency and speed
---
 preprocess.py | 50 +++++++++++++-------------------------------------
 1 file changed, 13 insertions(+), 37 deletions(-)

diff --git a/preprocess.py b/preprocess.py
index 63e09f3..bae7e29 100644
--- a/preprocess.py
+++ b/preprocess.py
@@ -1,29 +1,14 @@
 import os
 import subprocess
-import shutil
 import open3d
-
+import pandas as pd
+import numpy as np
 from dataset.semantic_dataset import all_file_prefixes
 
 
-def wc(file_name):
-    out = subprocess.Popen(
-        ["wc", "-l", file_name], stdout=subprocess.PIPE, stderr=subprocess.STDOUT
-    ).communicate()[0]
-    return int(out.partition(b" ")[0])
-
-
-def prepend_line(file_name, line):
-    with open(file_name, "r+") as f:
-        content = f.read()
-        f.seek(0, 0)
-        f.write(line.rstrip("\r\n") + "\n" + content)
-
-
 def point_cloud_txt_to_pcd(raw_dir, file_prefix):
     # File names
     txt_file = os.path.join(raw_dir, file_prefix + ".txt")
-    pts_file = os.path.join(raw_dir, file_prefix + ".pts")
     pcd_file = os.path.join(raw_dir, file_prefix + ".pcd")
 
     # Skip if already done
@@ -31,28 +16,18 @@ def point_cloud_txt_to_pcd(raw_dir, file_prefix):
         print("pcd {} exists, skipped".format(pcd_file))
         return
 
-    # .txt to .pts
-    # We could just prepend the line count, however, there are some intensity value
-    # which are non-integers.
-    print("[txt->pts]")
+    # .txt -> .pcd
+    print("[txt->pcd]")
     print("txt: {}".format(txt_file))
-    print("pts: {}".format(pts_file))
-    with open(txt_file, "r") as txt_f, open(pts_file, "w") as pts_f:
-        for line in txt_f:
-            # x, y, z, i, r, g, b
-            tokens = line.split()
-            tokens[3] = str(int(float(tokens[3])))
-            line = " ".join(tokens)
-            pts_f.write(line + "\n")
-    prepend_line(pts_file, str(wc(txt_file)))
-
-    # .pts -> .pcd
-    print("[pts->pcd]")
-    print("pts: {}".format(pts_file))
     print("pcd: {}".format(pcd_file))
-    point_cloud = open3d.read_point_cloud(pts_file)
-    open3d.write_point_cloud(pcd_file, point_cloud)
-    os.remove(pts_file)
+    pcd = open3d.geometry.PointCloud()
+    point_cloud = ((pd.read_csv(txt_file, index_col=False, header=None, sep=" ")).dropna(axis=1)).values
+    pcd.points = open3d.utility.Vector3dVector(point_cloud[:, :3])
+    pcd.colors = open3d.utility.Vector3dVector(point_cloud[:, 4:].astype(np.uint8))
+    open3d.io.write_point_cloud(pcd_file, pcd)
+    print(pcd_file + " DONE!")
+
+    return
 
 
 if __name__ == "__main__":
@@ -64,3 +39,4 @@ def point_cloud_txt_to_pcd(raw_dir, file_prefix):
 
     for file_prefix in all_file_prefixes:
         point_cloud_txt_to_pcd(raw_dir, file_prefix)
+    print("ALL DONE!")