Skip to content

Reusable registration routines #134

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
mhubii wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

Reusable registration routines #134

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
2d6a2d5
data structure refactor
mhubii Jul 24, 2026
bd0dccf
missing typing
mhubii Jul 24, 2026
da68c91
further refactors
mhubii Jul 27, 2026
e983b46
finished the hydra refactor
mhubii Jul 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 15 additions & 25 deletions cli/rr_cam_swarm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,16 +10,14 @@
NVDiffRastRenderer,
Robot,
RobotScene,
TorchKinematics,
TorchMeshContainer,
VirtualCamera,
)
from roboreg.io import (
find_files,
load_robot_data_from_ros_xacro,
load_robot_data_from_urdf_file,
parse_camera_info,
parse_mono_data,
parse_monocular_observations,
)
from roboreg.losses import soft_dice_loss
from roboreg.optim import LinearParticleSwarm, ParticleSwarmOptimizer
Expand Down Expand Up @@ -265,27 +263,27 @@ def main() -> None:
camera_info_file=args.camera_info_file
)
image_files = find_files(args.path, args.image_pattern)
mask_files = find_files(args.path, args.mask_pattern)
target_files = find_files(args.path, args.mask_pattern)
joint_states_files = find_files(args.path, args.joint_states_pattern)
n_samples = args.n_samples
if n_samples > len(image_files): # randomly sample n_samples
n_samples = len(image_files)
random_indices = np.random.choice(len(image_files), n_samples, replace=False)
image_files = np.array(image_files)[random_indices].tolist()
mask_files = np.array(mask_files)[random_indices].tolist()
target_files = np.array(target_files)[random_indices].tolist()
joint_states_files = np.array(joint_states_files)[random_indices].tolist()
images, joint_states, masks = parse_mono_data(
observations = parse_monocular_observations(
image_files=image_files,
mask_files=mask_files,
target_files=target_files,
joint_states_files=joint_states_files,
)

# pre-process data
joint_states = torch.tensor(
np.array(joint_states), dtype=torch.float32, device=device
np.array(observations.joint_states), dtype=torch.float32, device=device
)
n_joint_states = joint_states.shape[0]
masks = [mask_exponential_decay(mask) for mask in masks]
masks = [mask_exponential_decay(mask) for mask in observations.targets]
masks = torch.tensor(np.array(masks), dtype=torch.float32, device=device)

# scale image data (memory reduction)
Expand Down Expand Up @@ -345,20 +343,8 @@ def main() -> None:
collision=args.collision_meshes,
target_reduction=args.target_reduction,
)
mesh_container = TorchMeshContainer(
meshes=robot_data.meshes,
batch_size=batch_size,
device=device,
)
kinematics = TorchKinematics(
urdf=robot_data.urdf,
root_link_name=robot_data.root_link_name,
end_link_name=robot_data.end_link_name,
device=device,
)
robot = Robot(
mesh_container=mesh_container,
kinematics=kinematics,
robot = Robot.from_robot_data(
robot_data=robot_data, batch_size=batch_size, device=device
)

# instantiate scene
Expand Down Expand Up @@ -399,10 +385,14 @@ def fitness_closure() -> torch.Tensor:
).astype(np.uint8)
# upscale render
current_best_render = cv2.resize(
current_best_render, (images[offset].shape[1], images[offset].shape[0])
current_best_render,
(
observations.images[offset].shape[1],
observations.images[offset].shape[0],
),
)
overlay = overlay_mask(
images[offset],
observations.images[offset],
current_best_render,
scale=1.0,
)
Expand Down
173 changes: 54 additions & 119 deletions cli/rr_hydra.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,24 +4,21 @@
import numpy as np
import torch

from roboreg.core import Robot, TorchKinematics, TorchMeshContainer
from roboreg.hydra_icp import hydra_centroid_alignment, hydra_robust_icp
from roboreg.io import (
find_files,
load_robot_data_from_ros_xacro,
load_robot_data_from_urdf_file,
parse_camera_info,
parse_hydra_data,
parse_hydra_observations,
)
from roboreg.util import (
clean_xyz,
compute_vertex_normals,
depth_to_xyz,
from_homogeneous,
generate_ht_optical,
mask_extract_extended_boundary,
to_homogeneous,
from roboreg.registration.point_cloud.config import (
DepthToPointCloudConfig,
HydraConfig,
HydraRobustICPConfig,
)
from roboreg.registration.point_cloud.request import HydraRequest
from roboreg.registration.point_cloud.solver import HydraProblem, HydraRobustICP
from roboreg.registration.result import RegistrationResult

from .util.validate import validate_urdf_source

Expand Down Expand Up @@ -167,6 +164,26 @@ def args_factory() -> argparse.Namespace:
return parser.parse_args()


def visualize_hydra_result(
problem: HydraProblem,
result: RegistrationResult,
) -> None:
from roboreg.util import RegistrationVisualizer

visualizer = RegistrationVisualizer()

visualizer(
mesh_vertices=problem.reference_vertices,
observed_vertices=problem.observed_vertices,
)

visualizer(
mesh_vertices=problem.reference_vertices,
observed_vertices=problem.observed_vertices,
HT=torch.linalg.inv(result.extrinsics),
)


def main():
args = args_factory()
device = "cuda" if torch.cuda.is_available() else "cpu"
Expand All @@ -175,15 +192,14 @@ def main():
joint_states_files = find_files(args.path, args.joint_states_pattern)
mask_files = find_files(args.path, args.mask_pattern)
depth_files = find_files(args.path, args.depth_pattern)
joint_states, masks, depths = parse_hydra_data(
observations = parse_hydra_observations(
joint_states_files=joint_states_files,
mask_files=mask_files,
depth_files=depth_files,
)
height, width, intrinsics = parse_camera_info(args.camera_info_file)
_, _, intrinsics = parse_camera_info(args.camera_info_file)

# instantiate robot
batch_size = len(joint_states)
if args.urdf_path is not None:
robot_data = load_robot_data_from_urdf_file(
urdf_path=args.urdf_path,
Expand All @@ -199,118 +215,37 @@ def main():
end_link_name=args.end_link_name,
collision=args.collision_meshes,
)
mesh_container = TorchMeshContainer(
meshes=robot_data.meshes,
batch_size=len(joint_states),
device=device,
)
kinematics = TorchKinematics(
urdf=robot_data.urdf,
root_link_name=robot_data.root_link_name,
end_link_name=robot_data.end_link_name,
device=device,
)
robot = Robot(
mesh_container=mesh_container,
kinematics=kinematics,
)

# perform forward kinematics
joint_states = torch.tensor(
np.array(joint_states), dtype=torch.float32, device=device
)
robot.configure(joint_states)

# turn depths into xyzs
intrinsics = torch.tensor(intrinsics, dtype=torch.float32, device=device)
depths = torch.tensor(np.array(depths), dtype=torch.float32, device=device)
xyzs = depth_to_xyz(
depth=depths,
intrinsics=intrinsics,
z_min=args.z_min,
z_max=args.z_max,
conversion_factor=args.depth_conversion_factor,
)

# flatten BxHxWx3 -> Bx(H*W)x3
xyzs = xyzs.view(-1, height * width, 3)
xyzs = to_homogeneous(xyzs)
ht_optical = generate_ht_optical(xyzs.shape[0], dtype=torch.float32, device=device)
xyzs = torch.matmul(xyzs, ht_optical.transpose(-1, -2))
xyzs = from_homogeneous(xyzs)

# unflatten
xyzs = xyzs.view(-1, height, width, 3)
xyzs = [xyz.squeeze() for xyz in xyzs.cpu().numpy()]

# mesh vertices to list
mesh_vertices = from_homogeneous(robot.configured_vertices)
mesh_vertices = [mesh_vertices[i].contiguous() for i in range(batch_size)]
mesh_normals = []
for i in range(batch_size):
mesh_normals.append(
compute_vertex_normals(
vertices=mesh_vertices[i], faces=robot.mesh_container.faces
)
)

# clean observed vertices and turn into tensor
observed_vertices = [
torch.tensor(
clean_xyz(
xyz=xyz,
mask=(
mask
if args.no_boundary
else mask_extract_extended_boundary(
mask,
dilation_kernel=np.ones(
[args.dilation_kernel_size, args.dilation_kernel_size]
),
erosion_kernel=np.ones(
[args.erosion_kernel_size, args.erosion_kernel_size]
),
)
),
# register
config = HydraRobustICPConfig(
HydraConfig(
reference_points_per_mesh=args.number_of_points,
depth_to_point_cloud=DepthToPointCloudConfig(
z_min=args.z_min,
z_max=args.z_max,
depth_conversion_factor=args.depth_conversion_factor,
use_mask_boundary=not args.no_boundary,
dilation_kernel_size=args.dilation_kernel_size,
erosion_kernel_size=args.erosion_kernel_size,
),
dtype=torch.float32,
device=device,
max_correspondence_distance=args.max_distance,
)
for xyz, mask in zip(xyzs, masks)
]

# sample N points per mesh
for i in range(batch_size):
idx = torch.randperm(mesh_vertices[i].shape[0])[: args.number_of_points]
mesh_vertices[i] = mesh_vertices[i][idx]
mesh_normals[i] = mesh_normals[i][idx]

HT_init = hydra_centroid_alignment(observed_vertices, mesh_vertices)
HT = hydra_robust_icp(
HT_init,
observed_vertices,
mesh_vertices,
mesh_normals,
max_distance=args.max_distance,
outer_max_iter=args.outer_max_iter,
inner_max_iter=args.inner_max_iter,
)

# visualize
if args.display_results:
from roboreg.util import RegistrationVisualizer

visualizer = RegistrationVisualizer()
visualizer(mesh_vertices=mesh_vertices, observed_vertices=observed_vertices)
visualizer(
mesh_vertices=mesh_vertices,
observed_vertices=observed_vertices,
HT=torch.linalg.inv(HT),
hydra_robust_icp = HydraRobustICP(
config=config,
device=device,
callback=visualize_hydra_result if args.display_results else None,
)
result = hydra_robust_icp(
request=HydraRequest(
intrinsics=intrinsics,
robot_data=robot_data,
observations=observations,
)
)

# to numpy
HT = HT.cpu().numpy()
np.save(os.path.join(args.path, args.output_file), HT)
np.save(os.path.join(args.path, args.output_file), result.extrinsics.cpu().numpy())


if __name__ == "__main__":
Expand Down
Loading
Loading