#!/usr/bin/env python3
"""Merge survey point clouds into one voxel-downsampled PLY."""
import argparse
import sys
import xml.etree.ElementTree as ET
from pathlib import Path
import numpy as np
import rclpy.serialization
import rosbag2_py
import tf2_ros
from rclpy.duration import Duration
from rclpy.time import Time
from sensor_msgs.msg import JointState, PointCloud2
from sensor_msgs_py import point_cloud2
from tf2_msgs.msg import TFMessage
from tf_transformations import euler_matrix, quaternion_matrix, rotation_matrix
CLOUD_TOPIC = "/points/xyzrgba"
def xyz_rpy_matrix(xyz, rpy):
matrix = euler_matrix(*rpy, axes="sxyz")
matrix[:3, 3] = xyz
return matrix
class UrdfFkSolver:
"""Compute transforms from a URDF and joint positions for legacy bags."""
def __init__(self, path):
self.parent_joint = {}
for element in ET.parse(path).getroot().findall("joint"):
origin = element.find("origin")
axis = element.find("axis")
xyz = (
[float(value) for value in origin.get("xyz", "0 0 0").split()]
if origin is not None
else [0, 0, 0]
)
rpy = (
[float(value) for value in origin.get("rpy", "0 0 0").split()]
if origin is not None
else [0, 0, 0]
)
child = element.find("child").get("link")
self.parent_joint[child] = {
"name": element.get("name"),
"type": element.get("type"),
"parent": element.find("parent").get("link"),
"origin": xyz_rpy_matrix(xyz, rpy),
"axis": np.asarray(
[float(value) for value in axis.get("xyz", "1 0 0").split()]
if axis is not None
else [1, 0, 0]
),
}
def chain_to_root(self, link):
chain = []
seen = set()
while link in self.parent_joint:
if link in seen:
raise ValueError(f"cycle in URDF at {link}")
seen.add(link)
joint = self.parent_joint[link]
chain.append(joint)
link = joint["parent"]
return reversed(chain)
def joint_matrix(self, joint, angles):
if joint["type"] not in {"revolute", "continuous", "prismatic"}:
return joint["origin"]
if joint["name"] not in angles:
raise ValueError(f"joint {joint['name']!r} missing from /joint_states")
value = angles[joint["name"]]
if joint["type"] == "prismatic":
motion = np.eye(4)
motion[:3, 3] = joint["axis"] * value
else:
motion = rotation_matrix(value, joint["axis"])
return joint["origin"] @ motion
def world_matrix(self, link, angles):
matrix = np.eye(4)
for joint in self.chain_to_root(link):
matrix = matrix @ self.joint_matrix(joint, angles)
return matrix
def lookup(self, target, source, angles):
roots = {joint["parent"] for joint in self.parent_joint.values()}
for frame in (target, source):
if frame not in self.parent_joint and frame not in roots:
raise ValueError(f"frame {frame!r} not found in URDF")
return np.linalg.inv(self.world_matrix(target, angles)) @ self.world_matrix(
source, angles
)
def open_bag(path):
reader = rosbag2_py.SequentialReader()
reader.open(
rosbag2_py.StorageOptions(uri=str(path), storage_id=""),
rosbag2_py.ConverterOptions(
input_serialization_format="cdr", output_serialization_format="cdr"
),
)
topic_types = {item.name: item.type for item in reader.get_all_topics_and_types()}
return reader, topic_types
def transform_matrix(transform):
translation = transform.transform.translation
rotation = transform.transform.rotation
matrix = quaternion_matrix([rotation.x, rotation.y, rotation.z, rotation.w])
matrix[:3, 3] = [translation.x, translation.y, translation.z]
return matrix
def extract_points(cloud):
fields = {field.name for field in cloud.fields}
color_field = "rgba" if "rgba" in fields else "rgb" if "rgb" in fields else None
names = ["x", "y", "z"] + ([color_field] if color_field else [])
points = point_cloud2.read_points(
cloud, field_names=names, skip_nans=True, reshape_organized_cloud=False
)
xyz = np.column_stack(
[
points["x"].astype(np.float64),
points["y"].astype(np.float64),
points["z"].astype(np.float64),
]
)
if color_field is None or not len(xyz):
return xyz, None
packed = np.ascontiguousarray(points[color_field]).view(np.uint32).reshape(-1)
rgb = np.column_stack(
[
((packed >> 16) & 0xFF).astype(np.uint8),
((packed >> 8) & 0xFF).astype(np.uint8),
(packed & 0xFF).astype(np.uint8),
]
)
return xyz, rgb
def voxel_downsample(xyz, rgb, leaf):
if leaf <= 0 or not len(xyz):
return xyz, rgb
cells = np.floor(xyz / leaf).astype(np.int64)
_, groups, counts = np.unique(
cells, axis=0, return_inverse=True, return_counts=True
)
xyz_sum = np.zeros((len(counts), 3), dtype=np.float64)
np.add.at(xyz_sum, groups, xyz)
xyz = xyz_sum / counts[:, None]
if rgb is not None:
rgb_sum = np.zeros((len(counts), 3), dtype=np.float64)
np.add.at(rgb_sum, groups, rgb)
rgb = np.clip(rgb_sum / counts[:, None], 0, 255).astype(np.uint8)
return xyz, rgb
def write_ply(path, xyz, rgb):
fields = [("x", "<f4"), ("y", "<f4"), ("z", "<f4")]
header = [
"ply",
"format binary_little_endian 1.0",
f"element vertex {len(xyz)}",
"property float x",
"property float y",
"property float z",
]
if rgb is not None:
fields += [("r", "u1"), ("g", "u1"), ("b", "u1")]
header += ["property uchar red", "property uchar green", "property uchar blue"]
header += ["end_header", ""]
points = np.empty(len(xyz), dtype=np.dtype(fields))
points["x"], points["y"], points["z"] = xyz.T
if rgb is not None:
points["r"], points["g"], points["b"] = rgb.T
with open(path, "wb") as output:
output.write("\n".join(header).encode("ascii"))
output.write(points.tobytes())
def self_check():
xyz = np.array([[0.0, 0, 0], [0.4, 0, 0], [1.1, 0, 0]])
rgb = np.array([[0, 0, 0], [100, 100, 100], [200, 200, 200]], dtype=np.uint8)
sampled_xyz, sampled_rgb = voxel_downsample(xyz, rgb, 1.0)
assert np.allclose(sampled_xyz, [[0.2, 0, 0], [1.1, 0, 0]])
assert np.array_equal(sampled_rgb, [[50, 50, 50], [200, 200, 200]])
assert np.allclose(xyz_rpy_matrix([1, 2, 3], [0, 0, 0])[:3, 3], [1, 2, 3])
print("self-check passed")
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--bag")
parser.add_argument("--out", default="survey_map.ply")
parser.add_argument("--target-frame", default="base_link")
parser.add_argument("--voxel", type=float, default=0.005)
parser.add_argument("--max-range", type=float, default=2.0)
parser.add_argument(
"--urdf",
default="",
help="reconstruct transforms from /joint_states when legacy bags lack arm TF",
)
parser.add_argument("--self-check", action="store_true")
args, _ = parser.parse_known_args()
if args.self_check:
self_check()
return 0
if not args.bag:
parser.error("--bag is required")
if args.voxel < 0 or args.max_range < 0:
parser.error("--voxel and --max-range cannot be negative")
bag_path = Path(args.bag)
if not bag_path.is_dir():
print(f"ERROR: bag directory not found: {bag_path}", file=sys.stderr)
return 1
reader, topic_types = open_bag(bag_path)
if CLOUD_TOPIC not in topic_types:
print(f"ERROR: {CLOUD_TOPIC} is not present in the bag", file=sys.stderr)
return 1
fk_solver = None
if args.urdf:
urdf_path = Path(args.urdf)
if not urdf_path.is_file():
print(f"ERROR: URDF not found: {urdf_path}", file=sys.stderr)
return 1
try:
fk_solver = UrdfFkSolver(urdf_path)
except (OSError, ET.ParseError, ValueError) as exc:
print(f"ERROR: cannot load URDF: {exc}", file=sys.stderr)
return 1
tf_buffer = None if fk_solver else tf2_ros.Buffer(cache_time=Duration(seconds=3600))
latest_angles = None
clouds = []
colors = []
all_have_color = True
while reader.has_next():
topic, data, _ = reader.read_next()
if topic in {"/tf", "/tf_static"} and fk_solver is None:
message = rclpy.serialization.deserialize_message(data, TFMessage)
for transform in message.transforms:
if topic == "/tf_static":
tf_buffer.set_transform_static(transform, "bag")
else:
tf_buffer.set_transform(transform, "bag")
continue
if topic == "/joint_states" and fk_solver is not None:
message = rclpy.serialization.deserialize_message(data, JointState)
latest_angles = dict(zip(message.name, message.position))
continue
if topic != CLOUD_TOPIC:
continue
cloud = rclpy.serialization.deserialize_message(data, PointCloud2)
if fk_solver is not None:
if latest_angles is None:
print("WARN: skipping cloud recorded before any joint state")
continue
try:
matrix = fk_solver.lookup(
args.target_frame, cloud.header.frame_id, latest_angles
)
except ValueError as exc:
print(f"WARN: skipping cloud: {exc}")
continue
else:
try:
transform = tf_buffer.lookup_transform(
args.target_frame,
cloud.header.frame_id,
Time.from_msg(cloud.header.stamp),
timeout=Duration(seconds=0),
)
except tf2_ros.TransformException as exc:
print(f"WARN: skipping cloud without transform: {exc}")
continue
matrix = transform_matrix(transform)
xyz, rgb = extract_points(cloud)
if args.max_range > 0 and len(xyz):
keep = np.einsum("ij,ij->i", xyz, xyz) <= args.max_range**2
xyz = xyz[keep]
if rgb is not None:
rgb = rgb[keep]
if not len(xyz):
continue
xyz = xyz @ matrix[:3, :3].T + matrix[:3, 3]
raw_count = len(xyz)
xyz, rgb = voxel_downsample(xyz, rgb, args.voxel)
clouds.append(xyz)
if rgb is None:
all_have_color = False
colors.clear()
elif all_have_color:
colors.append(rgb)
print(f"Cloud {len(clouds)}: {raw_count:,} -> {len(xyz):,} points")
if not clouds:
print("ERROR: no clouds could be merged", file=sys.stderr)
return 1
xyz = np.vstack(clouds)
rgb = np.vstack(colors) if all_have_color else None
xyz, rgb = voxel_downsample(xyz, rgb, args.voxel)
write_ply(args.out, xyz, rgb)
print(f"Wrote {args.out}: {len(xyz):,} points from {len(clouds)} clouds")
return 0
if __name__ == "__main__":
sys.exit(main())