89 lines
2.8 KiB
Python
89 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import csv
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Iterable
|
|
|
|
import numpy as np
|
|
|
|
|
|
def read_point_cloud(path: str | Path) -> np.ndarray:
|
|
path = Path(path)
|
|
suffix = path.suffix.lower()
|
|
if suffix == ".ply":
|
|
return _read_ascii_ply(path)
|
|
if suffix in {".xyz", ".txt", ".csv"}:
|
|
return np.loadtxt(path, dtype=np.float64)
|
|
raise ValueError(f"Unsupported point cloud format: {path.suffix}")
|
|
|
|
|
|
def _read_ascii_ply(path: Path) -> np.ndarray:
|
|
header_lines = 0
|
|
vertex_count = None
|
|
with path.open("r", encoding="ascii", errors="strict") as fh:
|
|
for line in fh:
|
|
header_lines += 1
|
|
stripped = line.strip()
|
|
if stripped.startswith("element vertex"):
|
|
vertex_count = int(stripped.split()[-1])
|
|
if stripped == "end_header":
|
|
break
|
|
points = np.loadtxt(path, skiprows=header_lines, dtype=np.float64)
|
|
if points.ndim == 1:
|
|
points = points[None, :]
|
|
if points.shape[1] < 3:
|
|
raise ValueError(f"PLY file does not contain XYZ coordinates: {path}")
|
|
points = points[:, :3]
|
|
if vertex_count is not None and len(points) != vertex_count:
|
|
raise ValueError(
|
|
f"PLY vertex count mismatch: header={vertex_count}, data={len(points)}"
|
|
)
|
|
return points
|
|
|
|
|
|
def ensure_output_dir(path: str | Path) -> Path:
|
|
out_dir = Path(path)
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
return out_dir
|
|
|
|
|
|
def write_json(path: str | Path, payload: dict) -> None:
|
|
Path(path).write_text(
|
|
json.dumps(payload, ensure_ascii=False, indent=2),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def write_csv(path: str | Path, fieldnames: list[str], rows: Iterable[dict]) -> None:
|
|
with Path(path).open("w", encoding="utf-8", newline="") as fh:
|
|
writer = csv.DictWriter(fh, fieldnames=fieldnames)
|
|
writer.writeheader()
|
|
for row in rows:
|
|
writer.writerow(row)
|
|
|
|
|
|
def write_colored_ply(
|
|
path: str | Path,
|
|
points: np.ndarray,
|
|
colors: np.ndarray,
|
|
) -> None:
|
|
path = Path(path)
|
|
if len(points) != len(colors):
|
|
raise ValueError("Point/color length mismatch for PLY export.")
|
|
points = np.asarray(points, dtype=np.float64)
|
|
colors = np.asarray(colors, dtype=np.uint8)
|
|
with path.open("w", encoding="ascii", newline="\n") as fh:
|
|
fh.write("ply\n")
|
|
fh.write("format ascii 1.0\n")
|
|
fh.write(f"element vertex {len(points)}\n")
|
|
fh.write("property float x\n")
|
|
fh.write("property float y\n")
|
|
fh.write("property float z\n")
|
|
fh.write("property uchar red\n")
|
|
fh.write("property uchar green\n")
|
|
fh.write("property uchar blue\n")
|
|
fh.write("end_header\n")
|
|
for (x, y, z), (r, g, b) in zip(points, colors, strict=True):
|
|
fh.write(f"{x:.6f} {y:.6f} {z:.6f} {int(r)} {int(g)} {int(b)}\n")
|