diff --git a/pointcloud_nodes.py b/pointcloud_nodes.py index 0866c0b..37dc0af 100644 --- a/pointcloud_nodes.py +++ b/pointcloud_nodes.py @@ -504,24 +504,40 @@ class LoadPointCloud: arr = np.load(file_path) tensor_pc = torch.from_numpy(arr) return (tensor_pc,) - coords = [] - colors = [] - with open(file_path, 'r') as f: - line = f.readline().strip() - while not line.startswith("end_header"): + + if o3d is None: + logging.warning("[camera-comfyUI] open3d is not installed. Falling back to manual PLY parser.") + coords = [] + colors = [] + with open(file_path, 'r') as f: line = f.readline().strip() - for line in f: - parts = line.strip().split() - if len(parts) < 7: - continue - x, y, z = map(float, parts[0:3]) - r, g, b, a = map(int, parts[3:7]) - coords.append((x, y, z)) - colors.append((r, g, b, a)) - np_coords = np.array(coords, dtype=np.float32) - np_colors = np.array(colors, dtype=np.float32)/255.0 - combined = np.concatenate([np_coords, np_colors], axis=1) - tensor_pc = torch.from_numpy(combined) + while not line.startswith("end_header"): + line = f.readline().strip() + for line in f: + parts = line.strip().split() + if len(parts) < 7: + continue + x, y, z = map(float, parts[0:3]) + r, g, b, a = map(float, parts[3:7]) + coords.append((x, y, z)) + colors.append((r, g, b, a)) + np_coords = np.array(coords, dtype=np.float32) + np_colors = np.array(colors, dtype=np.float32) + else: + pc = o3d.t.io.read_point_cloud(file_path) + np_coords = pc.point["positions"].numpy().astype(np.float32) + if "colors" in pc.point: + cols = pc.point["colors"].numpy().astype(np.float32) + else: + cols = np.ones((np_coords.shape[0], 3), dtype=np.float32) + if "alpha" in pc.point: + alpha = pc.point["alpha"].numpy().astype(np.float32) + else: + alpha = np.ones((np_coords.shape[0], 1), dtype=np.float32) + np_colors = np.concatenate([cols, alpha], axis=1) + + combined = np.concatenate([np_coords, np_colors], axis=1) + tensor_pc = torch.from_numpy(combined) return (tensor_pc,) @classmethod @@ -586,24 +602,36 @@ class SavePointCloud: os.makedirs(full_output_folder, exist_ok=True) base_name = filename.replace("%batch_num%", "0") if save_as == "ply": - ply_name = f"{base_name}_{counter:05}.ply" - ply_path = os.path.join(full_output_folder, ply_name) - coords = pointcloud[:, :3].cpu().numpy() - colors = pointcloud[:, 3:].cpu().numpy().clip(0,1) - with open(ply_path, 'w') as f: - f.write("ply\n") - f.write("format ascii 1.0\n") - f.write(f"element vertex {coords.shape[0]}\n") - f.write("property float x\n") - f.write("property float y\n") - f.write("property float z\n") - f.write("property uchar red\n") - f.write("property uchar green\n") - f.write("property uchar blue\n") - f.write("property uchar alpha\n") - f.write("end_header\n") - for (x,y,z), (r,g,b,a) in zip(coords, colors): - f.write(f"{x} {y} {z} {int(r*255)} {int(g*255)} {int(b*255)} {int(a*255)}\n") + ply_name = f"{base_name}_{counter:05}.ply" + ply_path = os.path.join(full_output_folder, ply_name) + coords = pointcloud[:, :3].cpu().numpy().astype(np.float32) + colors = pointcloud[:, 3:].cpu().numpy().clip(0, 1).astype(np.float32) + + if o3d is None: + logging.warning("[camera-comfyUI] open3d is not installed. Falling back to manual ASCII PLY writer.") + with open(ply_path, 'w') as f: + f.write("ply\n") + f.write("format ascii 1.0\n") + f.write(f"element vertex {coords.shape[0]}\n") + f.write("property float x\n") + f.write("property float y\n") + f.write("property float z\n") + f.write("property float red\n") + f.write("property float green\n") + f.write("property float blue\n") + f.write("property float alpha\n") + f.write("end_header\n") + for (x, y, z), (r, g, b, a) in zip(coords, colors): + f.write(f"{x} {y} {z} {r} {g} {b} {a}\n") + else: + pc = o3d.t.geometry.PointCloud() + pc.point["positions"] = o3d.core.Tensor(coords, o3d.core.float32) + pc.point["colors"] = o3d.core.Tensor(colors[:, :3], o3d.core.float32) + if colors.shape[1] > 3: + pc.point["alpha"] = o3d.core.Tensor(colors[:, 3:], o3d.core.float32) + else: + pc.point["alpha"] = o3d.core.Tensor(np.ones((coords.shape[0], 1), dtype=np.float32), o3d.core.float32) + o3d.t.io.write_point_cloud(ply_path, pc) file_name = ply_name else: npy_name = f"{base_name}_{counter:05}.npy"