Merge pull request #6 from Alexankharin/codex/rewrite-i/o-operations-for-pointclouds

Update PLY IO with Open3D
This commit is contained in:
Alexander Kharin
2025-06-04 20:06:54 +03:00
committed by GitHub
+63 -35
View File
@@ -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"