Merge pull request #6 from Alexankharin/codex/rewrite-i/o-operations-for-pointclouds
Update PLY IO with Open3D
This commit is contained in:
+63
-35
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user