From 3bce69a71ce226184082397223aee46af663f9b7 Mon Sep 17 00:00:00 2001 From: Adraub Date: Fri, 13 Jun 2025 17:32:51 +0200 Subject: [PATCH] Add Pcd writing feature --- docs/io.rst | 2 +- src/pyntcloud/io/__init__.py | 3 +- src/pyntcloud/io/pcd.py | 87 ++++++++++++++++++++++++-- tests/data/diamond.pcd | 17 +++++ tests/integration/io/test_from_file.py | 1 + tests/integration/io/test_to_file.py | 1 + 6 files changed, 104 insertions(+), 7 deletions(-) create mode 100644 tests/data/diamond.pcd diff --git a/docs/io.rst b/docs/io.rst index 68465bd6..bc7e4f7f 100644 --- a/docs/io.rst +++ b/docs/io.rst @@ -19,7 +19,7 @@ generic array formats (more formats will be added in the near future): - `.npy / .npz `__ - `.obj `__ - `.off `__ (with color support) -- `.pcd `__ +- `.pcd `__ - `.ply `__ Reading diff --git a/src/pyntcloud/io/__init__.py b/src/pyntcloud/io/__init__.py index b34aa6b7..61f7e390 100644 --- a/src/pyntcloud/io/__init__.py +++ b/src/pyntcloud/io/__init__.py @@ -7,7 +7,7 @@ from .obj import read_obj, write_obj from .ply import read_ply, write_ply from .off import read_off -from .pcd import read_pcd +from .pcd import read_pcd, write_pcd FROM_FILE = { "ASC": read_ascii, @@ -32,6 +32,7 @@ "CSV": write_ascii, "NPZ": write_npz, "OBJ": write_obj, + "PCD": write_pcd, "PLY": write_ply, "PTS": write_ascii, "TXT": write_ascii, diff --git a/src/pyntcloud/io/pcd.py b/src/pyntcloud/io/pcd.py index 2173ce5e..effb6a88 100644 --- a/src/pyntcloud/io/pcd.py +++ b/src/pyntcloud/io/pcd.py @@ -24,7 +24,9 @@ def parse_header(lines): for ln in lines: if ln.startswith("#") or len(ln) < 2: continue - match = re.match(r"(\w+)\s+([\w\s\.]+)", ln) + # Should accept alphanumeric, decimal point, + # or negative sign (Viewpoint contains floats) + match = re.match(r"(\w+)\s+([\w\s\.\-]+)", ln) if not match: warnings.warn("warning: can't understand line: %s" % ln) continue @@ -119,10 +121,8 @@ def read_pcd(filename): if col in df.columns: # get the 'rgb' column from dataframe packed_rgb = df.rgb.values - # 'rgb' values are stored as float - # treat them as int - packed_rgb = packed_rgb.astype(np.float32).tostring() - packed_rgb = np.frombuffer(packed_rgb, dtype=np.int32) + # 'rgb' values are stored as float, treat them as int + packed_rgb = packed_rgb.astype(np.int32) # unpack 'rgb' into 'red', 'green' and 'blue' channel df["red"] = np.asarray((packed_rgb >> 16) & 255, dtype=np.uint8) df["green"] = np.asarray((packed_rgb >> 8) & 255, dtype=np.uint8) @@ -132,3 +132,80 @@ def read_pcd(filename): data["points"] = df return data + + +def write_pcd(filename, points): + """ + Write to a PCD file (ASCII format) populated with the given fields. + + Parameters + ---------- + filename : str + Path to the PCD file. + points : pandas.DataFrame + DataFrame containing the point cloud. Column names will be used + as FIELD names. + + Returns + ------- + boolean + True if no problems + + """ + if not filename.endswith("pcd"): + filename += ".pcd" + + # Some operation are needed for rgb + df = points.copy() + + # Handle color columns: combine red, green, blue into rgb as in PCL + rgb_columns = ["red", "green", "blue"] + # Run only if columns exists are unint8 + if (set(rgb_columns).issubset(df.columns) and + (df[rgb_columns].dtypes == np.uint8).all()): + # Pack RGB into a single float32 + df[rgb_columns] = df[rgb_columns].astype(np.int32) + rgb_int = ( + (df["red"].apply(lambda x: x << 16)) | + (df["green"].apply(lambda x: x << 8)) | + df["blue"] + ) + df["rgb"] = rgb_int.astype(np.float32) + df.drop(["red", "green", "blue"], axis=1, inplace=True) + + field_names = df.columns.tolist() + n_fields = len(field_names) + n_points = len(df) + + # Infer PCD data types and sizes from numpy dtypes + pcd_types = [] + pcd_sizes = [] + for col in field_names: + np_type = df[col].dtype + if np_type not in numpy_type_to_pcd_type: + raise ValueError(f"Unsupported dtype {np_type} for column '{col}'") + t, s = numpy_type_to_pcd_type[np_type] + pcd_types.append(t) + pcd_sizes.append(str(s)) + + # Header + header = [ + "# .PCD v.7 - Point Cloud Data file format", + "VERSION .7", + "FIELDS " + " ".join(field_names), + "SIZE " + " ".join(pcd_sizes), + "TYPE " + " ".join(pcd_types), + "COUNT " + " ".join(["1"] * n_fields), + f"WIDTH {n_points}", + "HEIGHT 1", + "VIEWPOINT 0 0 0 1 0 0 0", + f"POINTS {n_points}", + "DATA ascii" + ] + + # Write file + with open(filename, "w") as f: + f.write("\n".join(header) + "\n") + df.to_csv(f, sep=" ", header=False, index=False) + + return True diff --git a/tests/data/diamond.pcd b/tests/data/diamond.pcd new file mode 100644 index 00000000..6835b787 --- /dev/null +++ b/tests/data/diamond.pcd @@ -0,0 +1,17 @@ +# .PCD v.7 - Point Cloud Data file format +VERSION .7 +FIELDS x y z nx ny nz rgb +SIZE 4 4 4 4 4 4 4 +TYPE F F F F F F F +COUNT 1 1 1 1 1 1 1 +WIDTH 6 +HEIGHT 1 +VIEWPOINT 0 0 0 1 0 0 0 +POINTS 6 +DATA ascii +0.5 0.0 0.5 0.0 -1.0 0.0 16711680.0 +0.0 0.5 0.5 -1.0 0.0 0.0 16711680.0 +0.5 0.5 0.0 0.0 0.0 -1.0 65280.0 +1.0 0.5 0.5 1.0 0.0 0.0 16711680.0 +0.5 1.0 0.5 0.0 1.0 0.0 16711680.0 +0.5 0.5 1.0 0.0 0.0 1.0 255.0 diff --git a/tests/integration/io/test_from_file.py b/tests/integration/io/test_from_file.py index ede194cd..50195637 100644 --- a/tests/integration/io/test_from_file.py +++ b/tests/integration/io/test_from_file.py @@ -42,6 +42,7 @@ def assert_mesh(data): (".ply", True, True, False), ("_ascii.ply", True, True, True), ("_ascii_vertex_index.ply", True, True, True), + (".pcd", True, False, False), (".npz", True, True, False), (".obj", False, True, False), (".off", False, False, False), diff --git a/tests/integration/io/test_to_file.py b/tests/integration/io/test_to_file.py index 3072866b..da25c86f 100644 --- a/tests/integration/io/test_to_file.py +++ b/tests/integration/io/test_to_file.py @@ -13,6 +13,7 @@ (".npz", True, True, False), (".obj", False, True, False), (".bin", False, False, False), + (".pcd", True, False, False), ], ) def test_to_file(tmpdir, diamond, extension, color, mesh, comments):