diff --git a/cli/rr_cam_swarm.py b/cli/rr_cam_swarm.py index 1a52e93..ae811e0 100644 --- a/cli/rr_cam_swarm.py +++ b/cli/rr_cam_swarm.py @@ -10,8 +10,6 @@ NVDiffRastRenderer, Robot, RobotScene, - TorchKinematics, - TorchMeshContainer, VirtualCamera, ) from roboreg.io import ( @@ -19,7 +17,7 @@ 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 @@ -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) @@ -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 @@ -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, ) diff --git a/cli/rr_hydra.py b/cli/rr_hydra.py index 955e632..7c93285 100644 --- a/cli/rr_hydra.py +++ b/cli/rr_hydra.py @@ -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 @@ -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" @@ -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, @@ -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__": diff --git a/cli/rr_mono_dr.py b/cli/rr_mono_dr.py index 9d6dc73..d9b1882 100644 --- a/cli/rr_mono_dr.py +++ b/cli/rr_mono_dr.py @@ -10,19 +10,12 @@ import rich.progress import torch -from roboreg.core import ( - NVDiffRastRenderer, - Robot, - RobotScene, - TorchKinematics, - TorchMeshContainer, - VirtualCamera, -) +from roboreg.core import NVDiffRastRenderer, Robot, RobotScene, VirtualCamera from roboreg.io import ( find_files, load_robot_data_from_ros_xacro, load_robot_data_from_urdf_file, - parse_mono_data, + parse_monocular_observations, ) from roboreg.losses import soft_dice_loss from roboreg.util import mask_distance_transform, mask_exponential_decay, overlay_mask @@ -175,21 +168,21 @@ def main() -> None: # load data image_files = find_files(args.path, args.image_pattern) joint_states_files = find_files(args.path, args.joint_states_pattern) - mask_files = find_files(args.path, args.mask_pattern) - images, joint_states, masks = parse_mono_data( + target_files = find_files(args.path, args.mask_pattern) + observations = parse_monocular_observations( image_files=image_files, joint_states_files=joint_states_files, - mask_files=mask_files, + target_files=target_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 ) if mode == REGISTRATION_MODE.DISTANCE_FUNCTION: - targets = [mask_distance_transform(mask) for mask in masks] + targets = [mask_distance_transform(mask) for mask in observations.targets] elif mode == REGISTRATION_MODE.SEGMENTATION: - targets = [mask_exponential_decay(mask) for mask in masks] + targets = [mask_exponential_decay(mask) for mask in observations.targets] else: raise ValueError("Invalid registration mode.") targets = torch.tensor( @@ -220,20 +213,8 @@ def main() -> None: end_link_name=args.end_link_name, collision=args.collision_meshes, ) - mesh_container = TorchMeshContainer( - meshes=robot_data.meshes, - batch_size=joint_states.shape[0], - 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=joint_states.shape[0], device=device ) # instantiate scene @@ -298,7 +279,7 @@ def main() -> None: # display optimization progress if args.display_progress: render = renders["camera"][0].squeeze().detach().cpu().numpy() - image = images[0] + image = observations.images[0] render_overlay = overlay_mask( image, (render * 255.0).astype(np.uint8), @@ -307,7 +288,7 @@ def main() -> None: # difference left / right render / mask difference = ( cv2.cvtColor( - np.abs(render - masks[0].astype(np.float32) / 255.0), + np.abs(render - observations.targets[0].astype(np.float32) / 255.0), cv2.COLOR_GRAY2BGR, ) * 255.0 @@ -315,7 +296,7 @@ def main() -> None: # overlay segmentation mask segmentation_overlay = overlay_mask( image, - masks[0], + observations.targets[0], mode="b", scale=1.0, ) @@ -343,8 +324,10 @@ def main() -> None: for i, render in enumerate(renders): render = render.squeeze().cpu().numpy() - overlay = overlay_mask(images[i], (render * 255.0).astype(np.uint8), scale=1.0) - difference = np.abs(render - masks[i].astype(np.float32) / 255.0) + overlay = overlay_mask( + observations.images[i], (render * 255.0).astype(np.uint8), scale=1.0 + ) + difference = np.abs(render - observations.targets[i].astype(np.float32) / 255.0) cv2.imwrite(os.path.join(args.path, f"dr_overlay_{i}.png"), overlay) cv2.imwrite( diff --git a/cli/rr_render.py b/cli/rr_render.py index 65be955..a72252a 100644 --- a/cli/rr_render.py +++ b/cli/rr_render.py @@ -12,8 +12,6 @@ NVDiffRastRenderer, Robot, RobotScene, - TorchKinematics, - TorchMeshContainer, VirtualCamera, ) from roboreg.io import ( @@ -156,20 +154,8 @@ def main(): end_link_name=args.end_link_name, collision=args.collision_meshes, ) - mesh_container = TorchMeshContainer( - meshes=robot_data.meshes, - batch_size=args.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=args.batch_size, device=device ) scene = RobotScene( cameras=camera, diff --git a/cli/rr_stereo_dr.py b/cli/rr_stereo_dr.py index bb158ed..30de37f 100644 --- a/cli/rr_stereo_dr.py +++ b/cli/rr_stereo_dr.py @@ -14,15 +14,13 @@ 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_stereo_data, + parse_stereo_observations, ) from roboreg.losses import soft_dice_loss from roboreg.util import mask_distance_transform, mask_exponential_decay, overlay_mask @@ -206,28 +204,34 @@ def main() -> None: left_image_files = find_files(args.path, args.left_image_pattern) right_image_files = find_files(args.path, args.right_image_pattern) joint_states_files = find_files(args.path, args.joint_states_pattern) - left_mask_files = find_files(args.path, args.left_mask_pattern) - right_mask_files = find_files(args.path, args.right_mask_pattern) - left_images, right_images, joint_states, left_masks, right_masks = ( - parse_stereo_data( - left_image_files=left_image_files, - right_image_files=right_image_files, - joint_states_files=joint_states_files, - left_mask_files=left_mask_files, - right_mask_files=right_mask_files, - ) + left_target_files = find_files(args.path, args.left_mask_pattern) + right_target_files = find_files(args.path, args.right_mask_pattern) + observations = parse_stereo_observations( + left_image_files=left_image_files, + right_image_files=right_image_files, + joint_states_files=joint_states_files, + left_target_files=left_target_files, + right_target_files=right_target_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 ) if mode == REGISTRATION_MODE.DISTANCE_FUNCTION: - left_targets = [mask_distance_transform(mask) for mask in left_masks] - right_targets = [mask_distance_transform(mask) for mask in right_masks] + left_targets = [ + mask_distance_transform(mask) for mask in observations.left_targets + ] + right_targets = [ + mask_distance_transform(mask) for mask in observations.right_targets + ] elif mode == REGISTRATION_MODE.SEGMENTATION: - left_targets = [mask_exponential_decay(mask) for mask in left_masks] - right_targets = [mask_exponential_decay(mask) for mask in right_masks] + left_targets = [ + mask_exponential_decay(mask) for mask in observations.left_targets + ] + right_targets = [ + mask_exponential_decay(mask) for mask in observations.right_targets + ] else: raise ValueError("Invalid registration mode.") left_targets = torch.tensor( @@ -268,20 +272,8 @@ def main() -> None: end_link_name=args.end_link_name, collision=args.collision_meshes, ) - mesh_container = TorchMeshContainer( - meshes=robot_data.meshes, - batch_size=joint_states.shape[0], - 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=joint_states.shape[0], device=device ) # instantiate scene @@ -352,7 +344,7 @@ def main() -> None: if args.display_progress: render_overlays = [] left_render = renders["left"][0].squeeze().detach().cpu().numpy() - left_image = left_images[0] + left_image = observations.left_images[0] render_overlays.append( overlay_mask( left_image, @@ -361,7 +353,7 @@ def main() -> None: ) ) right_render = renders["right"][0].squeeze().detach().cpu().numpy() - right_image = right_images[0] + right_image = observations.right_images[0] render_overlays.append( overlay_mask( right_image, @@ -374,7 +366,10 @@ def main() -> None: differences.append( ( cv2.cvtColor( - np.abs(left_render - left_masks[0].astype(np.float32) / 255.0), + np.abs( + left_render + - observations.left_targets[0].astype(np.float32) / 255.0 + ), cv2.COLOR_GRAY2BGR, ) * 255.0 @@ -384,7 +379,8 @@ def main() -> None: ( cv2.cvtColor( np.abs( - right_render - right_masks[0].astype(np.float32) / 255.0 + right_render + - observations.right_targets[0].astype(np.float32) / 255.0 ), cv2.COLOR_GRAY2BGR, ) @@ -396,7 +392,7 @@ def main() -> None: segmentation_overlays.append( overlay_mask( left_image, - left_masks[0], + observations.left_targets[0], mode="b", scale=1.0, ) @@ -404,7 +400,7 @@ def main() -> None: segmentation_overlays.append( overlay_mask( right_image, - right_masks[0], + observations.right_targets[0], mode="b", scale=1.0, ) @@ -440,14 +436,20 @@ def main() -> None: left_render = left_render.squeeze().cpu().numpy() right_render = right_render.squeeze().cpu().numpy() left_overlay = overlay_mask( - left_images[i], (left_render * 255.0).astype(np.uint8), scale=1.0 + observations.left_images[i], + (left_render * 255.0).astype(np.uint8), + scale=1.0, ) right_overlay = overlay_mask( - right_images[i], (right_render * 255.0).astype(np.uint8), scale=1.0 + observations.right_images[i], + (right_render * 255.0).astype(np.uint8), + scale=1.0, + ) + left_difference = np.abs( + left_render - observations.left_targets[i].astype(np.float32) / 255.0 ) - left_difference = np.abs(left_render - left_masks[i].astype(np.float32) / 255.0) right_difference = np.abs( - right_render - right_masks[i].astype(np.float32) / 255.0 + right_render - observations.right_targets[i].astype(np.float32) / 255.0 ) cv2.imwrite(os.path.join(args.path, f"left_dr_overlay_{i}.png"), left_overlay) diff --git a/roboreg/core/robot.py b/roboreg/core/robot.py index b2ab4ad..f21bccd 100644 --- a/roboreg/core/robot.py +++ b/roboreg/core/robot.py @@ -1,9 +1,20 @@ -from typing import Union +from dataclasses import dataclass +from typing import Dict, Union import torch from .kinematics import TorchKinematics -from .structs import TorchMeshContainer +from .structs import Mesh, TorchMeshContainer + + +@dataclass +class RobotData: + r"""Data needed to construct a Robot.""" + + meshes: Dict[str, Mesh] + urdf: str + root_link_name: str + end_link_name: str class Robot: @@ -23,6 +34,27 @@ def __init__( ) self._device = mesh_container.device + @classmethod + def from_robot_data( + cls, + robot_data: RobotData, + batch_size: int, + device: Union[torch.device, str] = "cuda", + ) -> "Robot": + return Robot( + 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, + ), + ) + def configure( self, q: torch.FloatTensor, ht_root: torch.FloatTensor = None ) -> None: diff --git a/roboreg/core/structs.py b/roboreg/core/structs.py index cc95da1..803a375 100644 --- a/roboreg/core/structs.py +++ b/roboreg/core/structs.py @@ -1,12 +1,19 @@ import abc from collections import OrderedDict +from dataclasses import dataclass from pathlib import Path from typing import Dict, List, Optional, Tuple, Union import numpy as np import torch -from roboreg.io import Mesh + +@dataclass +class Mesh: + r"""Dataclass to hold mesh data.""" + + vertices: np.ndarray + faces: np.ndarray class TorchMeshContainer: diff --git a/roboreg/hydra_icp.py b/roboreg/hydra_icp.py deleted file mode 100644 index 616d6e9..0000000 --- a/roboreg/hydra_icp.py +++ /dev/null @@ -1,309 +0,0 @@ -from typing import List, Tuple - -import torch -from rich import print -from rich.progress import track - - -def kabsch_register( - input: torch.Tensor, target: torch.Tensor -) -> Tuple[torch.Tensor, torch.Tensor]: - r"""Kabsch algorithm: https://en.wikipedia.org/wiki/Kabsch_algorithm. - Computes rotation and translation such that input @ R + t = target. - - Args: - input (torch.Tensor): input of shape (..., M, 3). - target(torch.Tensor): target of shape (..., M, 3). - - Returns: - Tuple[torch.Tensor,torch.Tensor]: - - Rotation matrix of shape (..., 3, 3). - - Translation vector of shape (..., 3). - """ - # compute centroids - input_centroid = torch.mean(input, dim=-2) - target_centroid = torch.mean(target, dim=-2) - - # compute centered points - input_centered = input - input_centroid - target_centered = target - target_centroid - - # compute covariance matrix - H = target_centered.transpose(-1, -2) @ input_centered - - # compute SVD - U, _, V = torch.svd(H) - - E = torch.eye(3, dtype=U.dtype, device=U.device) - E[-1, -1] = torch.det(V @ U.transpose(-1, -2)) - - # compute rotation - R = V @ E @ U.transpose(-1, -2) - - # compute translation - t = target_centroid - input_centroid @ R - return R, t - - -def hydra_correspondence_indices( - input: torch.Tensor, target: torch.Tensor, max_distance: float = 0.1 -) -> Tuple[torch.Tensor, torch.Tensor]: - r"""For each point in input, find nearest neighbor index in target. - - Args: - input (torch.Tensor): Input of shape (M, 3) or (B, M, 3). - target (torch.Tensor): Target of shape (N, 3) or (B, N, 3). - max_distance (float): Maximum distance between point correspondences. - - Returns: - Tuple[torch.Tensor,torch.Tensor]: - - Match-indices of shape (M) or (B, M), where mi is the index of the nearest neighbor in target. - - Mask of shape (M) or (B, M). - """ - if input.shape[-1] != 3 or target.shape[-1] != 3: - raise ValueError("Input and target must have shape (..., 3).") - if max_distance < 0: - raise ValueError("Max distance must be positive.") - distances = torch.cdist(input, target, p=2) # (M, N) - min_distance, matchindices = torch.min(distances, dim=-1) # (M) - mask = min_distance < max_distance - return matchindices, mask - - -def hydra_centroid_alignment( - Xs: List[torch.Tensor], - Ys: List[torch.Tensor], -) -> torch.Tensor: - r"""Aligns centroids of Xs and Ys as an initial guess. - - Args: - Xs (List[torch.Tensor]): List of poinclouds of shape (Mi, 3). - Ys (List[torch.Tensor]): List of pointclouds of shape (Ni, 3). - - Returns: - torch.Tensor: Homogeneous transformation of shape (4, 4). HT @ Xs = Ys. - """ - # for each cloud compute centroid - Xs_centroids = [torch.mean(observation, dim=-2) for observation in Xs] - Ys_centroids = [torch.mean(mesh, dim=-2) for mesh in Ys] - - # estimate transform - R, t = kabsch_register( - torch.stack(Xs_centroids).unsqueeze(0), - torch.stack(Ys_centroids).unsqueeze(0), - ) - - HT = torch.eye(4, dtype=R.dtype, device=R.device) - R = R.squeeze(0) - t = t.squeeze(0) - HT[:3, :3] = R.T - HT[:3, 3] = t - return HT - - -def hydra_icp( - HT_init: torch.Tensor, - observations: List[torch.Tensor], - meshes: List[torch.Tensor], - max_distance: float = 0.1, - max_iter: int = 100, - rmse_change: float = 1e-6, - exit_early: bool = True, -) -> torch.Tensor: - r"""Hydra iterative closest point algorithm. - - Args: - HT_init: Initial guess. HT_init @ observations = meshes. - observations: List of observations of shape (Mi, 3). - meshes: List of meshes of shape (Ni, 3). - max_distance: Maximum distance between point correspondences. - max_iter: Maximum number of iterations. - rmse_change: Minimum change in rmse to continue iterating. - - Returns: - torch.Tensor: Homogeneous transformation of shape (4, 4). HT @ observations = meshes. - """ - HT = HT_init - # registration - prev_rmse = float("inf") - for _ in track(range(max_iter), description=f"Running Hydra ICP..."): - observation_corr = [] - mesh_corr = [] - for i in range(len(meshes)): - # search correspondences - observations_tf = observations[i] @ HT[:3, :3].T + HT[:3, 3] - matchindices, mask = hydra_correspondence_indices( - observations_tf, meshes[i], max_distance - ) - - observation_corr.append(observations[i][mask]) - mesh_corr.append(meshes[i][matchindices[mask]].squeeze()) - - observation_corr = torch.concatenate(observation_corr).unsqueeze(0) - mesh_corr = torch.concatenate(mesh_corr).unsqueeze(0) - - ( - R, - t, - ) = kabsch_register( - observation_corr, - mesh_corr, - ) - R = R.squeeze(0) - t = t.squeeze(0) - HT[:3, :3] = R.T - HT[:3, 3] = t - - # compute rmse between observation and mesh_corr - rmse = torch.sqrt( - torch.mean( - torch.sum( - torch.pow( - mesh_corr - observation_corr, - 2, - ), - dim=-1, - ) - ) - ) - - if abs(prev_rmse - rmse.item()) < rmse_change and exit_early: - print("Converged early. Exiting.") - break - - prev_rmse = rmse.item() - - print("HT estimate:\n", HT) - return HT - - -def hydra_robust_icp( - HT_init: torch.Tensor, - observations: List[torch.Tensor], - meshes: List[torch.Tensor], - mesh_normals: List[torch.Tensor], - max_distance: float = 0.1, - outer_max_iter: int = 100, - inner_max_iter: int = 3, - rmse_change: float = 1e-6, -) -> torch.Tensor: - r"""Lie-algebra point-to-plane ICP with robust loss, refer to section 1 - https://drive.google.com/file/d/1iIUqKchAbcYzwyS2D6jNI1J6KotReD1h/view?usp=sharing. - - Args: - HT_init: Initial guess. HT_init @ observations = meshes. - observations: List of observations of shape (Mi, 3). - meshes: List of meshes of shape (Ni, 3). - mesh_normals: List of mesh normals of shape (Ni, 3). - max_distance: Maximum distance between point correspondences. - outer_max_iter: Maximum number of outer iterations. - inner_max_iter: Maximum number of inner iterations. - rmse_change: Minimum change in rmse to continue iterating. - - Returns: - torch.Tensor: Homogeneous transformation of shape (4, 4). HT @ observations = meshes. - """ - HT = HT_init # HT @ observation = mesh - - observations_cross_mat = [] - for i in range(len(observations)): - # build observation cross product matrix, refer eq. 4 (gets created once) - observations_cross_mat.append( - torch.stack( - [ - torch.zeros_like(observations[i][:, 0]), - -observations[i][:, 2], - observations[i][:, 1], - observations[i][:, 2], - torch.zeros_like(observations[i][:, 0]), - -observations[i][:, 0], - -observations[i][:, 1], - observations[i][:, 0], - torch.zeros_like(observations[i][:, 0]), - ], - dim=-1, - ).reshape(-1, 3, 3) - ) - - # implementation of algorithm 1 - prev_rmse = float("inf") - dTh = torch.zeros_like(HT) - for _ in track(range(outer_max_iter), description=f"Running Hydra robust ICP..."): - observations_corr = [] - observations_cross_mat_corr = [] - meshes_corr = [] - meshes_normals_corr = [] - - for i in range(len(observations)): - if len(observations) != len(meshes): - raise ValueError("Length of observations and meshes must be the same.") - # search correspondences - observations_tf = observations[i] @ HT[:3, :3].T + HT[:3, 3] - matchindices, mask = hydra_correspondence_indices( - observations_tf, meshes[i], max_distance - ) - - observations_corr.append(observations[i][mask]) - observations_cross_mat_corr.append(observations_cross_mat[i][mask]) - meshes_corr.append(meshes[i][matchindices[mask].squeeze()]) - meshes_normals_corr.append(mesh_normals[i][matchindices[mask].squeeze()]) - - observations_corr = torch.cat(observations_corr) - observations_cross_mat_corr = torch.cat(observations_cross_mat_corr) - meshes_corr = torch.cat(meshes_corr) - meshes_normals_corr = torch.cat(meshes_normals_corr) - - for _ in range(inner_max_iter): - # ||A @ dTh - B||^2, refer eq. 14 - Al = meshes_normals_corr @ HT[:3, :3] # eq. 18 - Au = -Al.unsqueeze(1) @ observations_cross_mat_corr # eq. 19 - A = torch.cat((Au.squeeze(), Al.squeeze()), dim=-1) - B = torch.linalg.vecdot( - meshes_normals_corr, - meshes_corr - (observations_corr @ HT[:3, :3].T + HT[:3, 3]), - ) - # weight associated with Huber loss - kappa = ( - 1.345 * torch.median(torch.abs(B - torch.median(B))) / 0.6745 - ) # eq. 26 - W = torch.where( - torch.abs(B) < kappa, - torch.ones_like(B), - torch.full_like(B, kappa) / torch.abs(B), - ) - - dTh_vec, resid, rank, singvals = torch.linalg.lstsq(W[:, None] * A, W * B) - dTh[0, 1] = -dTh_vec[2] - dTh[0, 2] = dTh_vec[1] - dTh[1, 0] = dTh_vec[2] - dTh[1, 2] = -dTh_vec[0] - dTh[2, 0] = -dTh_vec[1] - dTh[2, 1] = dTh_vec[0] - - dTh[0, 3] = dTh_vec[3] - dTh[1, 3] = dTh_vec[4] - dTh[2, 3] = dTh_vec[5] - - HT = HT @ torch.linalg.matrix_exp(dTh) - - # compute rmse between observation and mesh_corr - rmse = torch.sqrt( - torch.mean( - torch.sum( - torch.pow( - meshes_corr - observations_corr, - 2, - ), - dim=-1, - ) - ) - ) - - if abs(prev_rmse - rmse.item()) < rmse_change: - print("Converged early. Exiting.") - break - - prev_rmse = rmse.item() - - print("HT estimate:\n", HT) - return HT diff --git a/roboreg/io/meshes.py b/roboreg/io/meshes.py index cfafc2c..27b3c3b 100644 --- a/roboreg/io/meshes.py +++ b/roboreg/io/meshes.py @@ -1,4 +1,3 @@ -from dataclasses import dataclass from pathlib import Path from typing import Dict, Union @@ -6,13 +5,7 @@ import numpy as np import trimesh - -@dataclass -class Mesh: - r"""Dataclass to hold mesh data.""" - - vertices: np.ndarray - faces: np.ndarray +from roboreg.core.structs import Mesh def load_mesh(path: Union[Path, str]) -> Mesh: diff --git a/roboreg/io/parsers.py b/roboreg/io/parsers.py index 50705ab..5fe1e12 100644 --- a/roboreg/io/parsers.py +++ b/roboreg/io/parsers.py @@ -7,6 +7,9 @@ import yaml from pytorch_kinematics import urdf_parser_py +from roboreg.registration.image.request import MonocularObservations, StereoObservations +from roboreg.registration.point_cloud.request import HydraObservations + class URDFParser: __slots__ = ["_urdf", "_robot"] @@ -312,11 +315,11 @@ def parse_camera_info( return height, width, intrinsic_matrix -def parse_hydra_data( +def parse_hydra_observations( joint_states_files: List[Path], mask_files: List[Path], depth_files: List[Path], -) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray]]: +) -> HydraObservations: r"""Parse data for Hydra registration. Args: @@ -325,10 +328,7 @@ def parse_hydra_data( depth_files (List[Path]): Depth files. Note that depth values are expected in meters. Returns: - Tuple[List[np.ndarray],List[np.ndarray],List[np.ndarray]]: - - Joint states. - - Masks of shape HxW. - - Point clouds of shape HxWx3. + HydraObservations: Data for Hydra registration. """ if len(joint_states_files) == 0 or len(mask_files) == 0 or len(depth_files) == 0: raise ValueError("No files found.") @@ -348,104 +348,74 @@ def parse_hydra_data( joint_states = [np.load(f) for f in joint_states_files] masks = [cv2.imread(f, cv2.IMREAD_GRAYSCALE) for f in mask_files] depths = [np.load(f) for f in depth_files] - if not all([mask.dtype == np.uint8 for mask in masks]): - raise ValueError("Masks must be of type np.uint8.") - if not all([np.all(mask >= 0) and np.all(mask <= 255) for mask in masks]): - raise ValueError("Masks must be in the range [0, 255].") if not all( [mask.shape[:2] == depth.shape[:2] for mask, depth in zip(masks, depths)] ): raise ValueError("Mask and depth shapes do not match.") - if not all(mask.ndim == 2 for mask in masks): - raise ValueError("Masks must be 2D.") - if not all(depth.ndim == 2 for depth in depths): - raise ValueError("Depths must be 2D.") - return joint_states, masks, depths + return HydraObservations(joint_states=joint_states, masks=masks, depths=depths) -def parse_mono_data( +def parse_monocular_observations( image_files: List[Path], joint_states_files: List[Path], - mask_files: List[Path], -) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray]]: - r"""Parse monocular data. + target_files: List[Path], +) -> MonocularObservations: + r"""Parse monocular observations. Args: image_files (List[Path]): Image files. joint_states_files (List[Path]): Joint states files. - mask_files (List[Path]): Mask files. + target_files (List[Path]): Target files. Returns: - Tuple[List[np.ndarray],List[np.ndarray],List[np.ndarray]]: - - Images of shape HxWx3. - - Joint states. - - Masks of shape HxW. + MonocularObservations: Data for monocular registration. """ if len(image_files) != len(joint_states_files) or len(image_files) != len( - mask_files + target_files ): raise ValueError("Number of images, joint states, masks do not match.") rich.print("Parsing the following files:") rich.print(f"Images: {[f.name for f in image_files]}") rich.print(f"Joint states: {[f.name for f in joint_states_files]}") - rich.print(f"Masks: {[f.name for f in mask_files]}") + rich.print(f"Targets: {[f.name for f in target_files]}") images = [cv2.imread(f) for f in image_files] joint_states = [np.load(f) for f in joint_states_files] - masks = [cv2.imread(f, cv2.IMREAD_GRAYSCALE) for f in mask_files] - if not all([mask.dtype == np.uint8 for mask in masks]): - raise ValueError("Masks must be of type np.uint8.") - if not all([np.all(mask >= 0) and np.all(mask <= 255) for mask in masks]): - raise ValueError("Masks must be in the range [0, 255].") + masks = [cv2.imread(f, cv2.IMREAD_GRAYSCALE) for f in target_files] if not all( [mask.shape[:2] == image.shape[:2] for mask, image in zip(masks, images)] ): raise ValueError("Mask and image shapes do not match.") - if not all(mask.ndim == 2 for mask in masks): - raise ValueError("Masks must be 2D.") - if not all(image.ndim == 3 for image in images): - raise ValueError("Images must be 3D.") - if not all(image.shape[-1] == 3 for image in images): - raise ValueError("Images must have 3 channels") - return images, joint_states, masks + return MonocularObservations( + images=images, joint_states=joint_states, targets=masks + ) -def parse_stereo_data( +def parse_stereo_observations( left_image_files: List[Path], right_image_files: List[Path], joint_states_files: List[Path], - left_mask_files: List[Path], - right_mask_files: List[Path], -) -> Tuple[ - List[np.ndarray], - List[np.ndarray], - List[np.ndarray], - List[np.ndarray], - List[np.ndarray], -]: - r"""Parse stereo data. + left_target_files: List[Path], + right_target_files: List[Path], +) -> StereoObservations: + r"""Parse stereo observations. Args: left_image_files (List[Path]): Left image files. right_image_files (List[Path]): Right image files. joint_states_files (List[Path]): Joint states files. - left_mask_files (List[Path]): Left mask files. - right_mask_files (List[Path]): Right mask files. + left_target_files (List[Path]): Left target files. + right_target_files (List[Path]): Right target files. Returns: - Tuple[List[np.ndarray],List[np.ndarray],List[np.ndarray],List[np.ndarray],List[np.ndarray]]: - - Left images of shape HxWx3. - - Right images of shape HxWx3. - - Joint states. - - Left masks of shape HxW. - - Right masks of shape HxW. + StereoObservations: Data for stereo registration. """ if ( len(left_image_files) != len(right_image_files) or len(left_image_files) != len(joint_states_files) - or len(left_image_files) != len(left_mask_files) - or len(left_image_files) != len(right_mask_files) + or len(left_image_files) != len(left_target_files) + or len(left_image_files) != len(right_target_files) ): raise ValueError( "Number of left / right images, joint states, left / right masks do not match." @@ -455,32 +425,18 @@ def parse_stereo_data( rich.print(f"Left images: {[f.name for f in left_image_files]}") rich.print(f"Right images: {[f.name for f in right_image_files]}") rich.print(f"Joint states: {[f.name for f in joint_states_files]}") - rich.print(f"Left masks: {[f.name for f in left_mask_files]}") - rich.print(f"Right masks: {[f.name for f in right_mask_files]}") + rich.print(f"Left targets: {[f.name for f in left_target_files]}") + rich.print(f"Right targets: {[f.name for f in right_target_files]}") left_images = [cv2.imread(f) for f in left_image_files] right_images = [cv2.imread(f) for f in right_image_files] joint_states = [np.load(f) for f in joint_states_files] - left_masks = [cv2.imread(f, cv2.IMREAD_GRAYSCALE) for f in left_mask_files] - right_masks = [cv2.imread(f, cv2.IMREAD_GRAYSCALE) for f in right_mask_files] - if not all([mask.dtype == np.uint8 for mask in left_masks]): - raise ValueError("Left masks must be of type np.uint8.") - if not all([np.all(mask >= 0) and np.all(mask <= 255) for mask in left_masks]): - raise ValueError("Left masks must be in the range [0, 255].") - if not all([mask.dtype == np.uint8 for mask in right_masks]): - raise ValueError("Left masks must be of type np.uint8.") - if not all([np.all(mask >= 0) and np.all(mask <= 255) for mask in right_masks]): - raise ValueError("Left masks must be in the range [0, 255].") - if not all(mask.ndim == 2 for mask in left_masks): - raise ValueError("Left masks must be 2D.") - if not all(image.ndim == 3 for image in left_images): - raise ValueError("Left images must be 3D.") - if not all(image.shape[-1] == 3 for image in left_images): - raise ValueError("Left images must have 3 channels") - if not all(mask.ndim == 2 for mask in right_masks): - raise ValueError("Right masks must be 2D.") - if not all(image.ndim == 3 for image in right_images): - raise ValueError("Right images must be 3D.") - if not all(image.shape[-1] == 3 for image in right_images): - raise ValueError("Right images must have 3 channels") - return left_images, right_images, joint_states, left_masks, right_masks + left_targets = [cv2.imread(f, cv2.IMREAD_GRAYSCALE) for f in left_target_files] + right_targets = [cv2.imread(f, cv2.IMREAD_GRAYSCALE) for f in right_target_files] + return StereoObservations( + left_images=left_images, + right_images=right_images, + joint_states=joint_states, + left_targets=left_targets, + right_targets=right_targets, + ) diff --git a/roboreg/io/robot_data.py b/roboreg/io/robot_data.py index 6e49c84..39e20b5 100644 --- a/roboreg/io/robot_data.py +++ b/roboreg/io/robot_data.py @@ -1,26 +1,10 @@ -from dataclasses import dataclass from pathlib import Path -from typing import Dict, Union +from typing import Union import rich -from roboreg.io import ( - Mesh, - URDFParser, - apply_mesh_origins, - load_meshes, - simplify_meshes, -) - - -@dataclass -class RobotData: - r"""Data needed to construct a Robot.""" - - meshes: Dict[str, Mesh] - urdf: str - root_link_name: str - end_link_name: str +from roboreg.core.robot import RobotData +from roboreg.io import URDFParser, apply_mesh_origins, load_meshes, simplify_meshes def load_robot_data_from_ros_xacro( diff --git a/roboreg/registration/__init__.py b/roboreg/registration/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/roboreg/registration/_validation.py b/roboreg/registration/_validation.py new file mode 100644 index 0000000..e66ce3d --- /dev/null +++ b/roboreg/registration/_validation.py @@ -0,0 +1,41 @@ +from typing import List + +import numpy as np + + +def validate_intrinsics(intrinsics: np.ndarray) -> None: + if intrinsics.shape != (3, 3): + raise ValueError(f"Intrinsics must have shape (3, 3), got {intrinsics.shape}.") + + +def validate_extrinsics(extrinsics: np.ndarray) -> None: + if extrinsics.shape != (4, 4): + raise ValueError(f"Extrinsics must have shape (4, 4), got {extrinsics.shape}.") + + +def validate_images(images: List[np.ndarray], name: str) -> None: + for index, image in enumerate(images): + if image.ndim != 3: + raise ValueError(f"{name}[{index}] must be 3D, got shape {image.shape}.") + + if image.shape[-1] != 3: + raise ValueError( + f"{name}[{index}] must have 3 channels, got shape {image.shape}." + ) + + +def validate_masks(masks: List[np.ndarray], name: str) -> None: + for index, mask in enumerate(masks): + if mask.ndim != 2: + raise ValueError(f"{name}[{index}] must be 2D, got shape {mask.shape}.") + + if mask.dtype != np.uint8: + raise ValueError( + f"{name}[{index}] must have dtype np.uint8, got {mask.dtype}." + ) + + +def validate_targets(targets: List[np.ndarray], name: str) -> None: + for index, target in enumerate(targets): + if target.ndim != 2: + raise ValueError(f"{name}[{index}] must be 2D, got shape {target.shape}.") diff --git a/roboreg/registration/image/__init__.py b/roboreg/registration/image/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/roboreg/registration/image/config.py b/roboreg/registration/image/config.py new file mode 100644 index 0000000..4fbccbf --- /dev/null +++ b/roboreg/registration/image/config.py @@ -0,0 +1,75 @@ +import math +from dataclasses import dataclass, field + + +@dataclass(frozen=True) +class ConvergenceConfig: + max_iterations: int = 400 + tolerance: float = 1.0e-3 + patience: int = 50 + + def __post_init__(self) -> None: + if self.max_iterations <= 0: + raise ValueError("max_iterations must be positive.") + if self.tolerance < 0: + raise ValueError("tolerance must be non-negative.") + if self.patience < 0: + raise ValueError("patience must be non-negative.") + + +@dataclass(frozen=True) +class PlateauSchedulerConfig: + mode: str = "min" + factor: float = 0.1 + patience: int = 50 + threshold: float = 1.0e-4 + + def __post_init__(self) -> None: + if self.factor <= 0 or self.factor >= 1: + raise ValueError("factor must be in the range (0, 1).") + if self.patience < 0: + raise ValueError("patience must be non-negative.") + if self.threshold < 0: + raise ValueError("threshold must be non-negative.") + + +@dataclass(frozen=True) +class DRRegConfig: + optimizer: str = "AdamW" + lr: float = 3.0e-2 + + plateau_scheduler: PlateauSchedulerConfig = field( + default_factory=PlateauSchedulerConfig + ) + convergence: ConvergenceConfig = field(default_factory=ConvergenceConfig) + + def __post_init__(self) -> None: + if self.lr <= 0: + raise ValueError("lr must be positive.") + + +@dataclass(frozen=True) +class CSRegConfig: + n_cameras: int = 50 + min_distance: float = 0.5 + max_distance: float = 2.0 + angle_range: float = math.pi + + inertia_weight: float = 0.7 + cognitive_coefficient: float = 1.5 + social_coefficient: float = 1.5 + + convergence: ConvergenceConfig = field(default_factory=ConvergenceConfig) + + def __post_init__(self) -> None: + if self.n_cameras <= 0: + raise ValueError("n_cameras must be positive.") + + if self.min_distance <= 0: + raise ValueError("min_distance must be positive.") + + if self.max_distance <= self.min_distance: + raise ValueError("max_distance must be greater than min_distance.") + + if self.angle_range <= 0: + raise ValueError("angle_range must be positive.") diff --git a/roboreg/registration/image/request.py b/roboreg/registration/image/request.py new file mode 100644 index 0000000..832c336 --- /dev/null +++ b/roboreg/registration/image/request.py @@ -0,0 +1,96 @@ +from dataclasses import dataclass +from typing import List + +import numpy as np + +from roboreg.core.robot import RobotData +from roboreg.registration._validation import ( + validate_extrinsics, + validate_images, + validate_intrinsics, + validate_targets, +) + + +@dataclass(frozen=True) +class MonocularObservations: + images: List[np.ndarray] + joint_states: List[np.ndarray] + targets: List[np.ndarray] + + def __post_init__(self) -> None: + lengths = { + "images": len(self.images), + "joint_states": len(self.joint_states), + "targets": len(self.targets), + } + + if len(set(lengths.values())) != 1: + raise ValueError( + f"All observation fields must have the same length, got {lengths}." + ) + + if not self.images: + raise ValueError("Expected at least one observation.") + + validate_images(self.images, "images") + validate_targets(self.targets, "targets") + + +@dataclass(frozen=True) +class StereoObservations: + left_images: List[np.ndarray] + right_images: List[np.ndarray] + joint_states: List[np.ndarray] + left_targets: List[np.ndarray] + right_targets: List[np.ndarray] + + def __post_init__(self) -> None: + lengths = { + "left_images": len(self.left_images), + "right_images": len(self.right_images), + "joint_states": len(self.joint_states), + "left_targets": len(self.left_targets), + "right_targets": len(self.right_targets), + } + + if len(set(lengths.values())) != 1: + raise ValueError( + f"All observation fields must have the same length, got {lengths}." + ) + + if not self.left_images: + raise ValueError("Expected at least one observation.") + + validate_images(self.left_images, "left_images") + validate_images(self.right_images, "right_images") + validate_targets(self.left_targets, "left_targets") + validate_targets(self.right_targets, "right_targets") + + +@dataclass(frozen=True) +class MonocularRequest: + intrinsics: np.ndarray + robot_data: RobotData + observations: MonocularObservations + initial_extrinsics: np.ndarray + + def __post_init__(self) -> None: + validate_intrinsics(self.intrinsics) + validate_extrinsics(self.initial_extrinsics) + + +@dataclass(frozen=True) +class StereoRequest: + left_intrinsics: np.ndarray + right_intrinsics: np.ndarray + initial_left_extrinsics: np.ndarray + left_to_right_extrinsics: np.ndarray + robot_data: RobotData + observations: StereoObservations + + def __post_init__(self) -> None: + validate_intrinsics(self.left_intrinsics) + validate_intrinsics(self.right_intrinsics) + validate_extrinsics(self.initial_left_extrinsics) + validate_extrinsics(self.left_to_right_extrinsics) diff --git a/roboreg/registration/point_cloud/__init__.py b/roboreg/registration/point_cloud/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/roboreg/registration/point_cloud/config.py b/roboreg/registration/point_cloud/config.py new file mode 100644 index 0000000..79122a0 --- /dev/null +++ b/roboreg/registration/point_cloud/config.py @@ -0,0 +1,37 @@ +from dataclasses import dataclass, field + + +@dataclass(frozen=True) +class DepthToPointCloudConfig: + z_min: float = 0.01 + z_max: float = 2.0 + depth_conversion_factor: float = 1.0 + + use_mask_boundary: bool = True + dilation_kernel_size: int = 3 + erosion_kernel_size: int = 10 + + +@dataclass(frozen=True) +class HydraConfig: + reference_points_per_mesh: int = 5000 + + depth_to_point_cloud: DepthToPointCloudConfig = field( + default_factory=DepthToPointCloudConfig + ) + + max_correspondence_distance: float = 0.1 + rmse_change_tolerance: float = 1e-6 + + +@dataclass(frozen=True) +class HydraICPConfig: + hydra: HydraConfig = field(default_factory=HydraConfig) + max_iterations: int = 100 + + +@dataclass(frozen=True) +class HydraRobustICPConfig: + hydra: HydraConfig = field(default_factory=HydraConfig) + max_outer_iterations: int = 50 + max_inner_iterations: int = 10 diff --git a/roboreg/registration/point_cloud/hydra.py b/roboreg/registration/point_cloud/hydra.py new file mode 100644 index 0000000..c2a4d71 --- /dev/null +++ b/roboreg/registration/point_cloud/hydra.py @@ -0,0 +1,342 @@ +from typing import List, Tuple + +import torch +from rich.progress import track + +from roboreg.registration.result import RegistrationResult, TerminationReason + + +def kabsch_register( + observed_vertices: torch.Tensor, reference_vertices: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + r"""Kabsch algorithm: https://en.wikipedia.org/wiki/Kabsch_algorithm. + Computes rotation and translation such that observed_vertices @ R + t = reference_vertices. + + Args: + observed_vertices (torch.Tensor): Observed vertices of shape (..., M, 3). + reference_vertices (torch.Tensor): Reference vertices of shape (..., M, 3). + + Returns: + Tuple[torch.Tensor,torch.Tensor]: + - Rotation matrix of shape (..., 3, 3). + - Translation vector of shape (..., 3). + """ + # compute centroids + observed_centroid = torch.mean(observed_vertices, dim=-2) + reference_centroid = torch.mean(reference_vertices, dim=-2) + + # compute centered points + observed_centered = observed_vertices - observed_centroid + reference_centered = reference_vertices - reference_centroid + + # compute covariance matrix + H = reference_centered.transpose(-1, -2) @ observed_centered + + # compute SVD + U, _, V = torch.svd(H) + + E = torch.eye(3, dtype=U.dtype, device=U.device) + E[-1, -1] = torch.det(V @ U.transpose(-1, -2)) + + # compute rotation + R = V @ E @ U.transpose(-1, -2) + + # compute translation + t = reference_centroid - observed_centroid @ R + return R, t + + +def correspondence_indices( + observed_vertices: torch.Tensor, + reference_vertices: torch.Tensor, + max_correspondence_distance: float = 0.1, +) -> Tuple[torch.Tensor, torch.Tensor]: + r"""For each point in input, find nearest neighbor index in target. + + Args: + observed_vertices (torch.Tensor): Observed vertices of shape (M, 3) or (B, M, 3). + reference_vertices (torch.Tensor): Reference vertices of shape (N, 3) or (B, N, 3). + max_correspondence_distance (float): Maximum distance between point correspondences. + + Returns: + Tuple[torch.Tensor,torch.Tensor]: + - Match-indices of shape (M) or (B, M), where mi is the index of the nearest neighbor in target. + - Mask of shape (M) or (B, M). + """ + if observed_vertices.shape[-1] != 3 or reference_vertices.shape[-1] != 3: + raise ValueError("Input and target must have shape (..., 3).") + if max_correspondence_distance < 0: + raise ValueError("Max distance must be positive.") + distances = torch.cdist(observed_vertices, reference_vertices, p=2) # (M, N) + min_distance, match_indices = torch.min(distances, dim=-1) # (M) + mask = min_distance < max_correspondence_distance + return match_indices, mask + + +def centroid_alignment( + observed_vertices: List[torch.Tensor], + reference_vertices: List[torch.Tensor], +) -> torch.Tensor: + r"""Aligns centroids of observed_vertices and Ys as an initial guess. + + Args: + observed_vertices (List[torch.Tensor]): List of poinclouds of shape (Mi, 3). + reference_vertices (List[torch.Tensor]): List of pointclouds of shape (Ni, 3). + + Returns: + torch.Tensor: Homogeneous transformation of shape (4, 4). HT @ observed_vertices = reference_vertices. + """ + # for each cloud compute centroid + observed_centroids = [ + torch.mean(observation, dim=-2) for observation in observed_vertices + ] + reference_centroids = [torch.mean(mesh, dim=-2) for mesh in reference_vertices] + + # estimate transform + R, t = kabsch_register( + observed_vertices=torch.stack(observed_centroids).unsqueeze(0), + reference_vertices=torch.stack(reference_centroids).unsqueeze(0), + ) + + HT = torch.eye(4, dtype=R.dtype, device=R.device) + R = R.squeeze(0) + t = t.squeeze(0) + HT[:3, :3] = R.T + HT[:3, 3] = t + return HT + + +def point_to_point_icp( + HT_init: torch.Tensor, + observed_vertices: List[torch.Tensor], + reference_vertices: List[torch.Tensor], + max_correspondence_distance: float = 0.1, + max_iterations: int = 100, + rmse_change_tolerance: float = 1e-6, +) -> RegistrationResult: + r"""Hydra iterative closest point algorithm. + + Args: + HT_init: Initial guess. HT_init @ observed_vertices = reference_vertices. + observed_vertices: List of observed vertices of shape (Mi, 3). + reference_vertices: List of reference vertices of shape (Ni, 3). + max_correspondence_distance: Maximum distance between point correspondences. + max_iterations: Maximum number of iterations. + rmse_change_tolerance: Minimum change in rmse to continue iterating. + + Returns: + RegistrationResult: Result with homogeneous transformation of shape (4, 4). HT @ observed_vertices = reference_vertices. + """ + HT = HT_init + # registration + previous_rmse = float("inf") + for iteration in track( + range(max_iterations), description=f"Running point to point ICP..." + ): + observed_correspondences = [] + reference_correspondences = [] + for i in range(len(reference_vertices)): + # search correspondences + observations_tf = observed_vertices[i] @ HT[:3, :3].T + HT[:3, 3] + match_indices, mask = correspondence_indices( + observations_tf, reference_vertices[i], max_correspondence_distance + ) + + observed_correspondences.append(observed_vertices[i][mask]) + reference_correspondences.append( + reference_vertices[i][match_indices[mask]].squeeze() + ) + + observed_correspondences = torch.concatenate( + observed_correspondences + ).unsqueeze(0) + reference_correspondences = torch.concatenate( + reference_correspondences + ).unsqueeze(0) + + ( + R, + t, + ) = kabsch_register( + observed_correspondences, + reference_correspondences, + ) + R = R.squeeze(0) + t = t.squeeze(0) + HT[:3, :3] = R.T + HT[:3, 3] = t + + # compute rmse between observed_correspondences and reference_correspondences + rmse = torch.sqrt( + torch.mean( + torch.sum( + torch.pow( + reference_correspondences - observed_correspondences, + 2, + ), + dim=-1, + ) + ) + ) + + if abs(previous_rmse - rmse.item()) < rmse_change_tolerance: + return RegistrationResult( + extrinsics=HT, + iterations=iteration, + termination_reason=TerminationReason.CONVERGED, + ) + + previous_rmse = rmse.item() + + return RegistrationResult( + extrinsics=HT, + iterations=max_iterations, + termination_reason=TerminationReason.MAX_ITERATIONS, + ) + + +def point_to_plane_robust_icp( + HT_init: torch.Tensor, + observed_vertices: List[torch.Tensor], + reference_vertices: List[torch.Tensor], + reference_normals: List[torch.Tensor], + max_correspondence_distance: float = 0.1, + max_outer_iterations: int = 100, + max_inner_iterations: int = 3, + rmse_change_tolerance: float = 1e-6, +) -> RegistrationResult: + r"""Lie-algebra point-to-plane ICP with robust loss, refer to section 1 + https://drive.google.com/file/d/1iIUqKchAbcYzwyS2D6jNI1J6KotReD1h/view?usp=sharing. + + Args: + HT_init: Initial guess. HT_init @ observed_vertices = reference_vertices. + observed_vertices: List of observed vertices of shape (Mi, 3). + reference_vertices: List of reference vertices of shape (Ni, 3). + reference_normals: List of reference normals of shape (Ni, 3). + max_correspondence_distance: Maximum distance between point correspondences. + max_outer_iterations: Maximum number of outer iterations. + max_inner_iterations: Maximum number of inner iterations. + rmse_change_tolerance: Minimum change in rmse to continue iterating. + + Returns: + RegistrationResult: Result with homogeneous transformation of shape (4, 4). HT @ observed_vertices = reference_vertices. + """ + HT = HT_init # HT @ observed_vertices = reference_vertices + + observed_cross_mat = [] + for i in range(len(observed_vertices)): + # build observation cross product matrix, refer eq. 4 (gets created once) + observed_cross_mat.append( + torch.stack( + [ + torch.zeros_like(observed_vertices[i][:, 0]), + -observed_vertices[i][:, 2], + observed_vertices[i][:, 1], + observed_vertices[i][:, 2], + torch.zeros_like(observed_vertices[i][:, 0]), + -observed_vertices[i][:, 0], + -observed_vertices[i][:, 1], + observed_vertices[i][:, 0], + torch.zeros_like(observed_vertices[i][:, 0]), + ], + dim=-1, + ).reshape(-1, 3, 3) + ) + + # implementation of algorithm 1 + previous_rmse = float("inf") + dTh = torch.zeros_like(HT) + for outer_iteration in track( + range(max_outer_iterations), description=f"Running point to plane robust ICP..." + ): + observed_correspondences = [] + observed_cross_mat_correspondences = [] + reference_correspondences = [] + reference_normals_correspondences = [] + + for i in range(len(observed_vertices)): + if len(observed_vertices) != len(reference_vertices): + raise ValueError("Length of observations and meshes must be the same.") + # search correspondences + observed_vertices_tf = observed_vertices[i] @ HT[:3, :3].T + HT[:3, 3] + match_indices, mask = correspondence_indices( + observed_vertices_tf, reference_vertices[i], max_correspondence_distance + ) + + observed_correspondences.append(observed_vertices[i][mask]) + observed_cross_mat_correspondences.append(observed_cross_mat[i][mask]) + reference_correspondences.append( + reference_vertices[i][match_indices[mask].squeeze()] + ) + reference_normals_correspondences.append( + reference_normals[i][match_indices[mask].squeeze()] + ) + + observed_correspondences = torch.cat(observed_correspondences) + observed_cross_mat_correspondences = torch.cat( + observed_cross_mat_correspondences + ) + reference_correspondences = torch.cat(reference_correspondences) + reference_normals_correspondences = torch.cat(reference_normals_correspondences) + + for _ in range(max_inner_iterations): + # ||A @ dTh - B||^2, refer eq. 14 + Al = reference_normals_correspondences @ HT[:3, :3] # eq. 18 + Au = -Al.unsqueeze(1) @ observed_cross_mat_correspondences # eq. 19 + A = torch.cat((Au.squeeze(), Al.squeeze()), dim=-1) + B = torch.linalg.vecdot( + reference_normals_correspondences, + reference_correspondences + - (observed_correspondences @ HT[:3, :3].T + HT[:3, 3]), + ) + # weight associated with Huber loss + kappa = ( + 1.345 * torch.median(torch.abs(B - torch.median(B))) / 0.6745 + ) # eq. 26 + W = torch.where( + torch.abs(B) < kappa, + torch.ones_like(B), + torch.full_like(B, kappa) / torch.abs(B), + ) + + dTh_vec, resid, rank, singvals = torch.linalg.lstsq(W[:, None] * A, W * B) + dTh[0, 1] = -dTh_vec[2] + dTh[0, 2] = dTh_vec[1] + dTh[1, 0] = dTh_vec[2] + dTh[1, 2] = -dTh_vec[0] + dTh[2, 0] = -dTh_vec[1] + dTh[2, 1] = dTh_vec[0] + + dTh[0, 3] = dTh_vec[3] + dTh[1, 3] = dTh_vec[4] + dTh[2, 3] = dTh_vec[5] + + HT = HT @ torch.linalg.matrix_exp(dTh) + + # compute rmse between observation and mesh_correspondences + rmse = torch.sqrt( + torch.mean( + torch.sum( + torch.pow( + reference_correspondences - observed_correspondences, + 2, + ), + dim=-1, + ) + ) + ) + + if abs(previous_rmse - rmse.item()) < rmse_change_tolerance: + return RegistrationResult( + extrinsics=HT, + iterations=outer_iteration, + termination_reason=TerminationReason.CONVERGED, + ) + + previous_rmse = rmse.item() + + return RegistrationResult( + extrinsics=HT, + iterations=max_outer_iterations, + termination_reason=TerminationReason.MAX_ITERATIONS, + ) diff --git a/roboreg/registration/point_cloud/request.py b/roboreg/registration/point_cloud/request.py new file mode 100644 index 0000000..09966e7 --- /dev/null +++ b/roboreg/registration/point_cloud/request.py @@ -0,0 +1,57 @@ +from dataclasses import dataclass +from typing import List, Tuple + +import numpy as np + +from roboreg.core.robot import RobotData +from roboreg.registration._validation import ( + validate_intrinsics, + validate_masks, + validate_targets, +) + + +@dataclass(frozen=True) +class HydraObservations: + joint_states: List[np.ndarray] + masks: List[np.ndarray] + depths: List[np.ndarray] + + def __post_init__(self) -> None: + lengths = { + "joint_states": len(self.joint_states), + "masks": len(self.masks), + "depths": len(self.depths), + } + + if len(set(lengths.values())) != 1: + raise ValueError( + f"All observation fields must have the same length, got {lengths}." + ) + + if not self.joint_states: + raise ValueError("Expected at least one observation.") + + validate_masks(self.masks, "masks") + validate_targets(self.depths, "depths") + + for i, (mask, depth) in enumerate(zip(self.masks, self.depths)): + if mask.shape != depth.shape: + raise ValueError( + f"masks[{i}] and depths[{i}] have incompatible shapes: " + f"{mask.shape} and {depth.shape}." + ) + + @property + def shape(self) -> Tuple[int, int]: + return self.depths[0].shape + + +@dataclass(frozen=True) +class HydraRequest: + intrinsics: np.ndarray + robot_data: RobotData + observations: HydraObservations + + def __post_init__(self) -> None: + validate_intrinsics(self.intrinsics) diff --git a/roboreg/registration/point_cloud/solver.py b/roboreg/registration/point_cloud/solver.py new file mode 100644 index 0000000..0a36515 --- /dev/null +++ b/roboreg/registration/point_cloud/solver.py @@ -0,0 +1,217 @@ +from dataclasses import dataclass +from typing import Callable, List, Optional + +import numpy as np +import torch + +from roboreg.core.robot import Robot +from roboreg.registration.result import RegistrationResult +from roboreg.util.mask import mask_extract_extended_boundary +from roboreg.util.points import ( + clean_xyz, + compute_vertex_normals, + from_homogeneous, + to_homogeneous, +) +from roboreg.util.transform import depth_to_xyz, generate_ht_optical + +from .config import HydraConfig, HydraICPConfig, HydraRobustICPConfig +from .hydra import centroid_alignment, point_to_plane_robust_icp, point_to_point_icp +from .request import HydraRequest + + +@dataclass(frozen=True) +class HydraProblem: + observed_vertices: List[torch.Tensor] + reference_vertices: List[torch.Tensor] + reference_normals: Optional[List[torch.Tensor]] = None + + +HydraCallback = Callable[ + ["HydraProblem", RegistrationResult], + None, +] + + +def _prepare_hydra_problem( + request: HydraRequest, + config: HydraConfig, + device: torch.device, + compute_normals: bool = True, +) -> HydraProblem: + # 1) construct robot on request + robot = Robot.from_robot_data( + robot_data=request.robot_data, + batch_size=len(request.observations.joint_states), + device=device, + ) + + # 2) to tensor + joint_states = torch.tensor( + np.stack(request.observations.joint_states), dtype=torch.float32, device=device + ) + intrinsics = torch.tensor(request.intrinsics, dtype=torch.float32, device=device) + depths = torch.tensor( + np.stack(request.observations.depths), dtype=torch.float32, device=device + ) + + # 3) perform forward kinematics + robot.configure(joint_states) + + # 4) process depths + xyzs = depth_to_xyz( + depth=depths, + intrinsics=intrinsics, + z_min=config.depth_to_point_cloud.z_min, + z_max=config.depth_to_point_cloud.z_max, + conversion_factor=config.depth_to_point_cloud.depth_conversion_factor, + ) + height, width = request.observations.shape + xyzs = xyzs.view(-1, height * width, 3) # flatten BxHxWx3 -> Bx(H*W)x3 + 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) + xyzs = xyzs.view(-1, height, width, 3) + xyzs = [xyz.squeeze() for xyz in xyzs.cpu().numpy()] + + # 5) clean observed vertices and turn into tensor + observed_vertices = [ + torch.tensor( + clean_xyz( + xyz=xyz, + mask=( + mask_extract_extended_boundary( + mask, + dilation_kernel=np.ones( + [ + config.depth_to_point_cloud.dilation_kernel_size, + config.depth_to_point_cloud.dilation_kernel_size, + ] + ), + erosion_kernel=np.ones( + [ + config.depth_to_point_cloud.erosion_kernel_size, + config.depth_to_point_cloud.erosion_kernel_size, + ] + ), + ) + if config.depth_to_point_cloud.use_mask_boundary + else mask + ), + ), + dtype=torch.float32, + device=device, + ) + for xyz, mask in zip(xyzs, request.observations.masks) + ] + + # mesh vertices to list + batch_size = len(request.observations.joint_states) + + mesh_vertices = from_homogeneous(robot.configured_vertices) + mesh_vertices = [mesh_vertices[i].contiguous() for i in range(batch_size)] + + mesh_normals: list[torch.Tensor] | None = None + + if compute_normals: + mesh_normals = [ + compute_vertex_normals( + vertices=mesh_vertices[i], + faces=robot.mesh_container.faces, + ) + for i in range(batch_size) + ] + + # sample N points per mesh + for i in range(batch_size): + n_points = min( + config.reference_points_per_mesh, + mesh_vertices[i].shape[0], + ) + + idx = torch.randperm( + mesh_vertices[i].shape[0], + device=mesh_vertices[i].device, + )[:n_points] + + mesh_vertices[i] = mesh_vertices[i][idx] + + if mesh_normals is not None: + mesh_normals[i] = mesh_normals[i][idx] + + return HydraProblem( + observed_vertices=observed_vertices, + reference_vertices=mesh_vertices, + reference_normals=mesh_normals, + ) + + +class HydraICP: + def __init__( + self, + config: HydraICPConfig | None = None, + device: torch.device | str = "cuda", + callback: HydraCallback | None = None, + ) -> None: + self._config = config or HydraICPConfig() + self._device = torch.device(device) + self._callback = callback + + def __call__(self, request: HydraRequest) -> RegistrationResult: + hydra_problem = _prepare_hydra_problem( + request=request, + config=self._config.hydra, + device=self._device, + compute_normals=False, + ) + HT_init = centroid_alignment( + hydra_problem.observed_vertices, hydra_problem.reference_vertices + ) + result = point_to_point_icp( + HT_init, + hydra_problem.observed_vertices, + hydra_problem.reference_vertices, + max_correspondence_distance=self._config.hydra.max_correspondence_distance, + max_iterations=self._config.max_iterations, + rmse_change_tolerance=self._config.hydra.rmse_change_tolerance, + ) + if self._callback is not None: + self._callback(hydra_problem, result) + return result + + +class HydraRobustICP: + def __init__( + self, + config: HydraRobustICPConfig | None = None, + device: torch.device | str = "cuda", + callback: HydraCallback | None = None, + ) -> None: + self._config = config or HydraRobustICPConfig() + self._device = torch.device(device) + self._callback = callback + + def __call__(self, request: HydraRequest) -> RegistrationResult: + hydra_problem = _prepare_hydra_problem( + request=request, + config=self._config.hydra, + device=self._device, + compute_normals=True, + ) + HT_init = centroid_alignment( + hydra_problem.observed_vertices, hydra_problem.reference_vertices + ) + result = point_to_plane_robust_icp( + HT_init, + hydra_problem.observed_vertices, + hydra_problem.reference_vertices, + hydra_problem.reference_normals, + max_correspondence_distance=self._config.hydra.max_correspondence_distance, + max_outer_iterations=self._config.max_outer_iterations, + max_inner_iterations=self._config.max_inner_iterations, + rmse_change_tolerance=self._config.hydra.rmse_change_tolerance, + ) + if self._callback is not None: + self._callback(hydra_problem, result) + return result diff --git a/roboreg/registration/result.py b/roboreg/registration/result.py new file mode 100644 index 0000000..f3b27df --- /dev/null +++ b/roboreg/registration/result.py @@ -0,0 +1,22 @@ +from dataclasses import dataclass +from enum import Enum + +import torch + + +class TerminationReason(str, Enum): + CONVERGED = "converged" + MAX_ITERATIONS = "max_iterations" + FAILED = "failed" + + +@dataclass +class RegistrationResult: + extrinsics: torch.Tensor + iterations: int + termination_reason: TerminationReason + message: str | None = None + + @property + def converged(self) -> bool: + return self.termination_reason == TerminationReason.CONVERGED diff --git a/test/core/test_robot.py b/test/core/test_robot.py index 4a4c4e0..4fa3c9d 100644 --- a/test/core/test_robot.py +++ b/test/core/test_robot.py @@ -1,6 +1,6 @@ import torch -from roboreg.core import Robot, TorchKinematics, TorchMeshContainer +from roboreg.core import Robot from roboreg.io import load_robot_data_from_urdf_file @@ -12,20 +12,8 @@ def test_robot() -> None: collision=True, ) - 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 ) assert robot.device == torch.device(device), "Robot device mismatch." diff --git a/test/core/test_scene.py b/test/core/test_scene.py index 88e3b0a..6a86cab 100644 --- a/test/core/test_scene.py +++ b/test/core/test_scene.py @@ -10,8 +10,6 @@ NVDiffRastRenderer, Robot, RobotScene, - TorchKinematics, - TorchMeshContainer, VirtualCamera, ) from roboreg.io import find_files, load_robot_data_from_urdf_file @@ -103,20 +101,8 @@ def __init__( root_link_name=root_link_name, end_link_name=end_link_name, ) - mesh_container = TorchMeshContainer( - meshes=robot_data.meshes, - batch_size=self.joint_states.shape[0], - 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=self.joint_states.shape[0], device=device ) # instantiate scene @@ -238,20 +224,8 @@ def test_single_camera_multiple_poses() -> None: root_link_name="lbr_link_0", end_link_name="lbr_link_7", ) - 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 diff --git a/test/io/test_parsers.py b/test/io/test_parsers.py index b3a13ea..f813fec 100644 --- a/test/io/test_parsers.py +++ b/test/io/test_parsers.py @@ -7,9 +7,9 @@ URDFParser, find_files, parse_camera_info, - parse_hydra_data, - parse_mono_data, - parse_stereo_data, + parse_hydra_observations, + parse_monocular_observations, + parse_stereo_observations, ) @@ -95,97 +95,110 @@ def test_parse_camera_info() -> None: assert intrinsic_matrix.shape == (3, 3), "Intrinsic matrix should be of shape 3x3." -def test_parse_hydra_data() -> None: +def test_parse_hydra_observations() -> None: path = "test/assets/lbr_med7_r800/samples" - joint_states, masks, depths = parse_hydra_data( + observations = parse_hydra_observations( joint_states_files=find_files(path, "joint_states_*.npy"), mask_files=find_files(path, "mask_sam2_left_*.png"), depth_files=find_files(path, "depth_*.npy"), ) assert ( - len(joint_states) == len(masks) == len(depths) + len(observations.joint_states) + == len(observations.masks) + == len(observations.depths) ), "Expected same number of joint states / masks / depths." - assert len(joint_states) >= 1, "Should at least have one sample." - assert masks[0].ndim == 2, "Expected 2D mask." - assert masks[0].dtype == np.uint8, "Expected unsigned integers for mask." - assert np.all(masks[0] >= 0) and np.all( - masks[0] <= 255 + assert len(observations.joint_states) >= 1, "Should at least have one sample." + assert observations.masks[0].ndim == 2, "Expected 2D mask." + assert ( + observations.masks[0].dtype == np.uint8 + ), "Expected unsigned integers for mask." + assert np.all(observations.masks[0] >= 0) and np.all( + observations.masks[0] <= 255 ), "Expected mask in range [0, 255]." - assert depths[0].ndim == 2, "Expected 2D depth map." + assert observations.depths[0].ndim == 2, "Expected 2D depth map." -def test_parse_mono_data() -> None: +def test_parse_monocular_observations() -> None: path = "test/assets/lbr_med7_r800/samples" - images, joint_states, masks = parse_mono_data( + observations = parse_monocular_observations( image_files=find_files(path, "left_image_*.png"), joint_states_files=find_files(path, "joint_states_*.npy"), - mask_files=find_files(path, "mask_sam2_left_*.png"), + target_files=find_files(path, "mask_sam2_left_*.png"), ) assert ( - len(images) == len(joint_states) == len(masks) + len(observations.images) + == len(observations.joint_states) + == len(observations.targets) ), "Expected same number of images / joint states / masks." - assert len(images) >= 1, "Should at least have one sample." - assert images[0].ndim == 3, "Expected 3D image (HxWx3)." - assert images[0].shape[-1] == 3, "Expected 3 color channels." - assert masks[0].ndim == 2, "Expected 2D mask." - assert masks[0].dtype == np.uint8, "Expected unsigned integers for mask." - assert np.all(masks[0] >= 0) and np.all( - masks[0] <= 255 + assert len(observations.images) >= 1, "Should at least have one sample." + assert observations.images[0].ndim == 3, "Expected 3D image (HxWx3)." + assert observations.images[0].shape[-1] == 3, "Expected 3 color channels." + assert observations.targets[0].ndim == 2, "Expected 2D mask." + assert ( + observations.targets[0].dtype == np.uint8 + ), "Expected unsigned integers for mask." + assert np.all(observations.targets[0] >= 0) and np.all( + observations.targets[0] <= 255 ), "Expected mask in range [0, 255]." assert ( - masks[0].shape[:2] == images[0].shape[:2] + observations.targets[0].shape[:2] == observations.images[0].shape[:2] ), "Mask and image dimensions should match." -def test_parse_stereo_data() -> None: +def test_parse_stereo_observations() -> None: path = "test/assets/lbr_med7_r800/samples" - left_images, right_images, joint_states, left_masks, right_masks = ( - parse_stereo_data( - left_image_files=find_files(path, "left_image_*.png"), - right_image_files=find_files(path, "right_image_*.png"), - joint_states_files=find_files(path, "joint_states_*.npy"), - left_mask_files=find_files(path, "mask_sam2_left_*.png"), - right_mask_files=find_files(path, "mask_sam2_right_*.png"), - ) + observations = parse_stereo_observations( + left_image_files=find_files(path, "left_image_*.png"), + right_image_files=find_files(path, "right_image_*.png"), + joint_states_files=find_files(path, "joint_states_*.npy"), + left_target_files=find_files(path, "mask_sam2_left_*.png"), + right_target_files=find_files(path, "mask_sam2_right_*.png"), ) assert ( - len(left_images) - == len(right_images) - == len(joint_states) - == len(left_masks) - == len(right_masks) + len(observations.left_images) + == len(observations.right_images) + == len(observations.joint_states) + == len(observations.left_targets) + == len(observations.right_targets) ), "Expected same number of left/right images, joint states, and left/right masks." - assert len(left_images) >= 1, "Should at least have one sample." + assert len(observations.left_images) >= 1, "Should at least have one sample." # Test left data - assert left_images[0].ndim == 3, "Expected 3D left image (HxWx3)." - assert left_images[0].shape[-1] == 3, "Expected 3 color channels for left image." - assert left_masks[0].ndim == 2, "Expected 2D left mask." - assert left_masks[0].dtype == np.uint8, "Expected unsigned integers for left mask." - assert np.all(left_masks[0] >= 0) and np.all( - left_masks[0] <= 255 + assert observations.left_images[0].ndim == 3, "Expected 3D left image (HxWx3)." + assert ( + observations.left_images[0].shape[-1] == 3 + ), "Expected 3 color channels for left image." + assert observations.left_targets[0].ndim == 2, "Expected 2D left mask." + assert ( + observations.left_targets[0].dtype == np.uint8 + ), "Expected unsigned integers for left mask." + assert np.all(observations.left_targets[0] >= 0) and np.all( + observations.left_targets[0] <= 255 ), "Expected left mask in range [0, 255]." # Test right data - assert right_images[0].ndim == 3, "Expected 3D right image (HxWx3)." - assert right_images[0].shape[-1] == 3, "Expected 3 color channels for right image." - assert right_masks[0].ndim == 2, "Expected 2D right mask." + assert observations.right_images[0].ndim == 3, "Expected 3D right image (HxWx3)." + assert ( + observations.right_images[0].shape[-1] == 3 + ), "Expected 3 color channels for right image." + assert observations.right_targets[0].ndim == 2, "Expected 2D right mask." assert ( - right_masks[0].dtype == np.uint8 + observations.right_targets[0].dtype == np.uint8 ), "Expected unsigned integers for right mask." - assert np.all(right_masks[0] >= 0) and np.all( - right_masks[0] <= 255 + assert np.all(observations.right_targets[0] >= 0) and np.all( + observations.right_targets[0] <= 255 ), "Expected right mask in range [0, 255]." # Test dimensions match assert ( - left_masks[0].shape[:2] == left_images[0].shape[:2] + observations.left_targets[0].shape[:2] == observations.left_images[0].shape[:2] ), "Left mask and image dimensions should match." assert ( - right_masks[0].shape[:2] == right_images[0].shape[:2] + observations.right_targets[0].shape[:2] + == observations.right_images[0].shape[:2] ), "Right mask and image dimensions should match." @@ -200,6 +213,6 @@ def test_parse_stereo_data() -> None: test_urdf_parser_from_ros_xacro() test_find_files() test_parse_camera_info() - test_parse_hydra_data() - test_parse_mono_data() - test_parse_stereo_data() + test_parse_hydra_observations() + test_parse_monocular_observations() + test_parse_stereo_observations() diff --git a/test/test_hydra_icp.py b/test/test_hydra_icp.py index 97c9cd4..42a9b79 100644 --- a/test/test_hydra_icp.py +++ b/test/test_hydra_icp.py @@ -6,19 +6,19 @@ import transformations as tf from roboreg.core import TorchKinematics, TorchMeshContainer -from roboreg.hydra_icp import ( - hydra_centroid_alignment, - hydra_correspondence_indices, - hydra_icp, - hydra_robust_icp, -) from roboreg.io import ( URDFParser, + apply_mesh_origins, find_files, load_meshes, parse_camera_info, - apply_mesh_origins, - parse_hydra_data, + parse_hydra_observations, +) +from roboreg.registration.point_cloud.hydra import ( + centroid_alignment, + correspondence_indices, + point_to_point_icp, + point_to_plane_robust_icp, ) from roboreg.util import ( RegistrationVisualizer, @@ -48,7 +48,7 @@ def test_hydra_centroid_alignment(): for mesh_centroid in mesh_centroids ] - HT = hydra_centroid_alignment(mesh_centroids, observed_centroids) + HT = centroid_alignment(mesh_centroids, observed_centroids) assert torch.allclose(HT, HT_random) @@ -76,19 +76,23 @@ def test_index_shape( raise ValueError("Indices contain negative indices.") # single input - input = torch.rand(M, dim) - target = torch.rand(N, dim) # e.g. the mesh vertices - matchindices, mask = hydra_correspondence_indices( - input, target, max_distance=np.sqrt(dim) / 2.0 # remove some elements randomly + observed_vertices = torch.rand(M, dim) + reference_vertices = torch.rand(N, dim) # e.g. the mesh vertices + matchindices, mask = correspondence_indices( + observed_vertices, + reference_vertices, + max_correspondence_distance=np.sqrt(dim) / 2.0, # remove some elements randomly ) test_index_shape(matchindices, mask, torch.Size([M]), N) # batched input batch_size = 2 - input = torch.rand(batch_size, M, dim) - target = torch.rand(batch_size, N, dim) - matchindices, mask = hydra_correspondence_indices( - input, target, max_distance=np.sqrt(dim) / 2.0 + observed_vertices = torch.rand(batch_size, M, dim) + reference_vertices = torch.rand(batch_size, N, dim) + matchindices, mask = correspondence_indices( + observed_vertices, + reference_vertices, + max_correspondence_distance=np.sqrt(dim) / 2.0, ) test_index_shape(matchindices, mask, torch.Size([batch_size, M]), N) @@ -96,16 +100,18 @@ def test_index_shape( M = 10 N = 100 - input = torch.rand(M, dim) - target = torch.rand(N, dim) - matchindices, mask = hydra_correspondence_indices( - input, target, max_distance=np.sqrt(dim) / 2.0 + observed_vertices = torch.rand(M, dim) + reference_vertices = torch.rand(N, dim) + matchindices, mask = correspondence_indices( + observed_vertices, + reference_vertices, + max_correspondence_distance=np.sqrt(dim) / 2.0, ) test_index_shape(matchindices, mask, torch.Size([M]), N) @pytest.mark.skip(reason="To be fixed.") -def test_hydra_icp(): +def test_hydra_point_to_point_icp(): device = "cuda" if torch.cuda.is_available() else "cpu" ros_package = "lbr_description" xacro_path = "urdf/med7/med7.xacro" @@ -118,7 +124,7 @@ def test_hydra_icp(): depth_pattern = "depth_*.npy" # load data - joint_states, masks, depths = parse_hydra_data( + observations = parse_hydra_observations( joint_states_files=find_files(path, joint_states_pattern), mask_files=find_files(path, mask_pattern), depth_files=find_files(path, depth_pattern), @@ -139,7 +145,7 @@ def test_hydra_icp(): ) # instantiate mesh - batch_size = len(joint_states) + batch_size = len(observations.joint_states) meshes = TorchMeshContainer( meshes=apply_mesh_origins( meshes=load_meshes( @@ -156,19 +162,19 @@ def test_hydra_icp(): ) # perform forward kinematics - mesh_vertices = meshes.vertices.clone() + reference_vertices = meshes.vertices.clone() joint_states = torch.tensor( - np.array(joint_states), dtype=torch.float32, device=device + np.array(observations.joint_states), dtype=torch.float32, device=device ) - ht_lookup = kinematics.mesh_forward_kinematics(joint_states) + ht_lookup = kinematics.forward_kinematics(joint_states) for link_name, ht in ht_lookup.items(): - mesh_vertices[ + reference_vertices[ :, meshes.lower_vertex_index_lookup[ link_name ] : meshes.upper_vertex_index_lookup[link_name], ] = torch.matmul( - mesh_vertices[ + reference_vertices[ :, meshes.lower_vertex_index_lookup[ link_name @@ -176,11 +182,13 @@ def test_hydra_icp(): ], ht.transpose(-1, -2), ) - mesh_vertices = from_homogeneous(mesh_vertices) + reference_vertices = from_homogeneous(reference_vertices) # turn depths into xyzs intrinsics = torch.tensor(intrinsics, dtype=torch.float32, device=device) - depths = torch.tensor(np.array(depths), dtype=torch.float32, device=device) + depths = torch.tensor( + np.array(observations.depths), dtype=torch.float32, device=device + ) xyzs = depth_to_xyz(depth=depths, intrinsics=intrinsics, z_max=1.5) # flatten BxHxWx3 -> Bx(H*W)x3 @@ -194,8 +202,8 @@ def test_hydra_icp(): xyzs = xyzs.view(-1, height, width, 3) xyzs = [xyz.squeeze() for xyz in xyzs.cpu().numpy()] - # mesh vertices to list - mesh_vertices = [mesh_vertices[i].contiguous() for i in range(batch_size)] + # reference vertices to list + reference_vertices = [reference_vertices[i].contiguous() for i in range(batch_size)] # clean observed vertices and turn into tensor observed_vertices = [ @@ -204,39 +212,41 @@ def test_hydra_icp(): dtype=torch.float32, device=device, ) - for xyz, mask in zip(xyzs, masks) + for xyz, mask in zip(xyzs, observations.masks) ] # sample 5000 points per mesh for i in range(batch_size): - idx = torch.randperm(mesh_vertices[i].shape[0])[:5000] - mesh_vertices[i] = mesh_vertices[i][idx] + idx = torch.randperm(reference_vertices[i].shape[0])[:5000] + reference_vertices[i] = reference_vertices[i][idx] - HT_init = hydra_centroid_alignment(observed_vertices, mesh_vertices) - HT = hydra_icp( + HT_init = centroid_alignment(observed_vertices, reference_vertices) + registration_result = point_to_point_icp( HT_init, observed_vertices, - mesh_vertices, - max_distance=0.1, - max_iter=int(1e3), - rmse_change=1e-8, + reference_vertices, + max_correspondence_distance=0.1, + max_iterations=int(1e3), + rmse_change_tolerance=1e-8, ) # visualize visualizer = RegistrationVisualizer() - visualizer(mesh_vertices=mesh_vertices, observed_vertices=observed_vertices) + visualizer(mesh_vertices=reference_vertices, observed_vertices=observed_vertices) visualizer( - mesh_vertices=mesh_vertices, + mesh_vertices=reference_vertices, observed_vertices=observed_vertices, - HT=torch.linalg.inv(HT), + HT=torch.linalg.inv(registration_result.extrinsics), ) # to numpy - np.save(os.path.join(path, "HT_hydra.npy"), HT.cpu().numpy()) + np.save( + os.path.join(path, "HT_hydra.npy"), registration_result.extrinsics.cpu().numpy() + ) @pytest.mark.skip(reason="To be fixed.") -def test_hydra_robust_icp() -> None: +def test_hydra_point_to_plane_robust_icp() -> None: device = "cuda" if torch.cuda.is_available() else "cpu" ros_package = "lbr_description" xacro_path = "urdf/med7/med7.xacro" @@ -249,7 +259,7 @@ def test_hydra_robust_icp() -> None: depth_pattern = "depth_*.npy" # load data - joint_states, masks, depths = parse_hydra_data( + observations = parse_hydra_observations( joint_states_files=find_files(path, joint_states_pattern), mask_files=find_files(path, mask_pattern), depth_files=find_files(path, depth_pattern), @@ -270,7 +280,7 @@ def test_hydra_robust_icp() -> None: ) # instantiate mesh - batch_size = len(joint_states) + batch_size = len(observations.joint_states) meshes = TorchMeshContainer( meshes=apply_mesh_origins( meshes=load_meshes( @@ -287,19 +297,19 @@ def test_hydra_robust_icp() -> None: ) # perform forward kinematics - mesh_vertices = meshes.vertices.clone() + reference_vertices = meshes.vertices.clone() joint_states = torch.tensor( - np.array(joint_states), dtype=torch.float32, device=device + np.array(observations.joint_states), dtype=torch.float32, device=device ) ht_lookup = kinematics.forward_kinematics(joint_states) for link_name, ht in ht_lookup.items(): - mesh_vertices[ + reference_vertices[ :, meshes.lower_vertex_index_lookup[ link_name ] : meshes.upper_vertex_index_lookup[link_name], ] = torch.matmul( - mesh_vertices[ + reference_vertices[ :, meshes.lower_vertex_index_lookup[ link_name @@ -310,7 +320,9 @@ def test_hydra_robust_icp() -> None: # turn depths into xyzs intrinsics = torch.tensor(intrinsics, dtype=torch.float32, device=device) - depths = torch.tensor(np.array(depths), dtype=torch.float32, device=device) + depths = torch.tensor( + np.array(observations.depths), dtype=torch.float32, device=device + ) xyzs = depth_to_xyz(depth=depths, intrinsics=intrinsics, z_max=1.5) # flatten BxHxWx3 -> Bx(H*W)x3 @@ -325,12 +337,12 @@ def test_hydra_robust_icp() -> None: xyzs = [xyz.squeeze() for xyz in xyzs.cpu().numpy()] # mesh vertices to list - mesh_vertices = from_homogeneous(mesh_vertices) - mesh_vertices = [mesh_vertices[i].contiguous() for i in range(batch_size)] - mesh_normals = [] + reference_vertices = from_homogeneous(reference_vertices) + reference_vertices = [reference_vertices[i].contiguous() for i in range(batch_size)] + reference_normals = [] for i in range(batch_size): - mesh_normals.append( - compute_vertex_normals(vertices=mesh_vertices[i], faces=meshes.faces) + reference_normals.append( + compute_vertex_normals(vertices=reference_vertices[i], faces=meshes.faces) ) # clean observed vertices and turn into tensor @@ -340,38 +352,40 @@ def test_hydra_robust_icp() -> None: dtype=torch.float32, device=device, ) - for xyz, mask in zip(xyzs, masks) + for xyz, mask in zip(xyzs, observations.masks) ] # sample 5000 points per mesh for i in range(batch_size): - idx = torch.randperm(mesh_vertices[i].shape[0])[:5000] - mesh_vertices[i] = mesh_vertices[i][idx] - mesh_normals[i] = mesh_normals[i][idx] + idx = torch.randperm(reference_vertices[i].shape[0])[:5000] + reference_vertices[i] = reference_vertices[i][idx] + reference_normals[i] = reference_normals[i][idx] - HT_init = hydra_centroid_alignment(observed_vertices, mesh_vertices) - HT = hydra_robust_icp( + HT_init = centroid_alignment(observed_vertices, reference_vertices) + registration_result = point_to_plane_robust_icp( HT_init, observed_vertices, - mesh_vertices, - mesh_normals, - max_distance=0.1, - outer_max_iter=int(50), - inner_max_iter=10, + reference_vertices, + reference_normals, + max_correspondence_distance=0.1, + max_outer_iterations=50, + max_inner_iterations=10, ) # visualize visualizer = RegistrationVisualizer() - visualizer(mesh_vertices=mesh_vertices, observed_vertices=observed_vertices) + visualizer(mesh_vertices=reference_vertices, observed_vertices=observed_vertices) visualizer( - mesh_vertices=mesh_vertices, + mesh_vertices=reference_vertices, observed_vertices=observed_vertices, - HT=torch.linalg.inv(HT), + HT=torch.linalg.inv(registration_result.extrinsics), ) # to numpy - HT = HT.cpu().numpy() - np.save(os.path.join(path, "HT_hydra_robust.npy"), HT) + np.save( + os.path.join(path, "HT_hydra_robust.npy"), + registration_result.extrinsics.cpu().numpy(), + ) if __name__ == "__main__": @@ -383,5 +397,5 @@ def test_hydra_robust_icp() -> None: # test_hydra_centroid_alignment() # test_hydra_correspondence_indices() - # test_hydra_icp() - test_hydra_robust_icp() + # test_hydra_point_to_point_icp() + test_hydra_point_to_plane_robust_icp()