diff --git a/cli/rr_cam_swarm.py b/cli/rr_cam_swarm.py index d73c102..d523b4e 100644 --- a/cli/rr_cam_swarm.py +++ b/cli/rr_cam_swarm.py @@ -1,11 +1,11 @@ -import argparse import os from pathlib import Path -from typing import Union +from typing import Optional, Union import cv2 import numpy as np import torch +import typer from roboreg.core import NVDiffRastRenderer, Robot, RobotScene, VirtualCamera from roboreg.io import ( @@ -26,168 +26,7 @@ from .util.validate import validate_urdf_source - -def args_factory() -> argparse.Namespace: - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter - ) - parser.add_argument( - "--n-cameras", - type=int, - default=50, - help="The number of cameras / particles to optimize.", - ) - parser.add_argument( - "--min-distance", - type=float, - default=0.5, - help="The minimum distance of the camera from the object.", - ) - parser.add_argument( - "--max-distance", - type=float, - default=2.0, - help="The maximum distance of the camera from the object.", - ) - parser.add_argument( - "--angle-range", - type=float, - default=np.pi, - help="The initial angle range for the camera in [-angle_range/2, angle_range/2].", - ) - parser.add_argument( - "--w", - type=float, - default=0.7, - ) - parser.add_argument( - "--c1", - type=float, - default=1.5, - ) - parser.add_argument( - "--c2", - type=float, - default=1.5, - ) - parser.add_argument( - "--max-iterations", - type=int, - default=100, - help="The maximum number of iterations.", - ) - parser.add_argument( - "--min-fitness-change", - type=float, - default=2.0e-3, - help="The minimum fitness change for early convergence.", - ) - parser.add_argument( - "--max-iterations-below-min-fitness-change", - type=int, - default=20, - help="The maximum number of iterations below the minimum fitness change before early convergence.", - ) - parser.add_argument( - "--display-progress", - action="store_true", - help="Display optimization progress.", - ) - parser.add_argument( - "--urdf-path", - type=str, - default="test/assets/lbr_med7_r800/description/lbr_med7_r800.urdf", - help="Path to URDF file. Meshes resolved relative to this file. " - "Mutually exclusive with --ros-package/--xacro-path.", - ) - parser.add_argument( - "--ros-package", - type=str, - default=None, - help="ROS package containing robot description. " - "Requires --xacro-path. Mutually exclusive with --urdf-path.", - ) - parser.add_argument( - "--xacro-path", - type=str, - default=None, - help="Path to xacro file relative to --ros-package. " - "Requires --ros-package. Mutually exclusive with --urdf-path.", - ) - parser.add_argument( - "--root-link-name", - type=str, - default="", - help="Root link name. If unspecified, the first link with mesh will be used, which may cause errors.", - ) - parser.add_argument( - "--end-link-name", - type=str, - default="", - help="End link name. If unspecified, the last link with mesh will be used, which may cause errors.", - ) - parser.add_argument( - "--target-reduction", - type=float, - default=0.95, - help="Reduces the mesh vertex count for memory reduction. In [0, 1).", - ) - parser.add_argument( - "--scale", - type=float, - default=0.25, - help="Scale the camera resolution by this factor. Reduces memory usage.", - ) - parser.add_argument( - "--collision-meshes", - action="store_true", - help="If set, collision meshes will be used instead of visual meshes.", - ) - parser.add_argument( - "--camera-info-file", - type=str, - required=True, - help="Path to the camera parameters, /camera_info.yaml.", - ) - parser.add_argument("--path", type=str, required=True, help="Path to the data.") - parser.add_argument( - "--image-pattern", - type=str, - default="image_*.png", - help="Image file pattern. The images are only used to --display-progress.", - ) - parser.add_argument( - "--joint-states-pattern", - type=str, - default="joint_states_*.npy", - help="Joint state file pattern.", - ) - parser.add_argument( - "--mask-pattern", - type=str, - default="image_*_mask.png", - help="Mask file pattern.", - ) - parser.add_argument( - "--output-file", - type=str, - default="HT_cam_swarm.npy", - help="Output file name. Relative to --path.", - ) - parser.add_argument( - "--n-samples", - type=int, - default=5, - help="Number of samples to randomly select from the data for optimization.", - ) - parser.add_argument( - "--max-jobs", - type=int, - default=2, - help="Number of concurrent compilation jobs for nvdiffrast. Only relevant on first run.", - ) - validate_urdf_source(parser, parser.parse_args()) - return parser.parse_args() +app = typer.Typer(add_completion=False) def instantiate_particles( @@ -249,20 +88,100 @@ def instantiate_particles( return torch.cat([random_eyes, random_centers, random_angles], dim=-1) -def main() -> None: - args = args_factory() +@app.command() +def main( + camera_info_file: Path = typer.Option( + ..., help="Path to the camera parameters, /camera_info.yaml." + ), + path: Path = typer.Option(..., help="Path to the data."), + n_cameras: int = typer.Option( + 50, help="The number of cameras / particles to optimize." + ), + min_distance: float = typer.Option( + 0.5, help="The minimum distance of the camera from the object." + ), + max_distance: float = typer.Option( + 2.0, help="The maximum distance of the camera from the object." + ), + angle_range: float = typer.Option( + np.pi, + help="The initial angle range for the camera in [-angle_range/2, angle_range/2].", + ), + w: float = typer.Option(0.7), + c1: float = typer.Option(1.5), + c2: float = typer.Option(1.5), + max_iterations: int = typer.Option(100, help="The maximum number of iterations."), + min_fitness_change: float = typer.Option( + 2.0e-3, help="The minimum fitness change for early convergence." + ), + max_iterations_below_min_fitness_change: int = typer.Option( + 20, + help="The maximum number of iterations below the minimum fitness change before early convergence.", + ), + display_progress: bool = typer.Option(False, help="Display optimization progress."), + urdf_path: Optional[Path] = typer.Option( + "test/assets/lbr_med7_r800/description/lbr_med7_r800.urdf", + help="Path to URDF file. Meshes resolved relative to this file. " + "Mutually exclusive with --ros-package/--xacro-path.", + ), + ros_package: Optional[str] = typer.Option( + None, + help="ROS package containing robot description. " + "Requires --xacro-path. Mutually exclusive with --urdf-path.", + ), + xacro_path: Optional[str] = typer.Option( + None, + help="Path to xacro file relative to --ros-package. " + "Requires --ros-package. Mutually exclusive with --urdf-path.", + ), + root_link_name: str = typer.Option( + "", + help="Root link name. If unspecified, the first link with mesh will be used, which may cause errors.", + ), + end_link_name: str = typer.Option( + "", + help="End link name. If unspecified, the last link with mesh will be used, which may cause errors.", + ), + target_reduction: float = typer.Option( + 0.95, + help="Reduces the mesh vertex count for memory reduction. In [0, 1).", + ), + scale: float = typer.Option( + 0.25, help="Scale the camera resolution by this factor. Reduces memory usage." + ), + collision_meshes: bool = typer.Option( + False, help="If set, collision meshes will be used instead of visual meshes." + ), + image_pattern: str = typer.Option( + "image_*.png", + help="Image file pattern. The images are only used to --display-progress.", + ), + joint_states_pattern: str = typer.Option( + "joint_states_*.npy", help="Joint state file pattern." + ), + mask_pattern: str = typer.Option("image_*_mask.png", help="Mask file pattern."), + output_file: str = typer.Option( + "HT_cam_swarm.npy", help="Output file name. Relative to --path." + ), + n_samples: int = typer.Option( + 5, + help="Number of samples to randomly select from the data for optimization.", + ), + max_jobs: int = typer.Option( + 2, + help="Number of concurrent compilation jobs for nvdiffrast. Only relevant on first run.", + ), +) -> None: + r"""Particle swarm optimization for an initial camera pose estimate.""" + validate_urdf_source(urdf_path, ros_package, xacro_path) device = "cuda" if torch.cuda.is_available() else "cpu" - os.environ["MAX_JOBS"] = str(args.max_jobs) # limit number of concurrent jobs - path = Path(args.path) + os.environ["MAX_JOBS"] = str(max_jobs) # limit number of concurrent jobs # load data - height, width, intrinsics = parse_camera_info( - camera_info_file=args.camera_info_file - ) - image_files = find_files(path, args.image_pattern) - target_files = find_files(path, args.mask_pattern) - joint_states_files = find_files(path, args.joint_states_pattern) - n_samples = args.n_samples + height, width, intrinsics = parse_camera_info(camera_info_file=camera_info_file) + image_files = find_files(path, image_pattern) + target_files = find_files(path, mask_pattern) + joint_states_files = find_files(path, joint_states_pattern) 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) @@ -288,35 +207,35 @@ def main() -> None: masks = torch.tensor(np.array(masks), dtype=torch.float32, device=device) # scale image data (memory reduction) - height = int(height * args.scale) - width = int(width * args.scale) - intrinsics = intrinsics * args.scale + height = int(height * scale) + width = int(width * scale) + intrinsics = intrinsics * scale masks = torch.nn.functional.interpolate( masks.unsqueeze(1), size=(height, width), mode="nearest" ).squeeze(1) # prepare particles particles = instantiate_particles( - n_particles=args.n_cameras, + n_particles=n_cameras, height=height, width=width, focal_length_x=intrinsics[0, 0], focal_length_y=intrinsics[1, 1], - eye_min_dist=args.min_distance, - eye_max_dist=args.max_distance, - angle_interval=args.angle_range, + eye_min_dist=min_distance, + eye_max_dist=max_distance, + angle_interval=angle_range, device=device, ) particle_swarm = LinearParticleSwarm( particles=particles, - w=args.w, - c1=args.c1, - c2=args.c2, + w=w, + c1=c1, + c2=c2, ) # instantiate scene for fitness evaluation batch_size = ( - n_joint_states * args.n_cameras + n_joint_states * n_cameras ) # (each camera observes n_joint_states joint states) camera = VirtualCamera( resolution=(height, width), @@ -326,22 +245,22 @@ def main() -> None: ) # instantiate robot - if args.urdf_path is not None: + if urdf_path is not None: robot_data = load_robot_data_from_urdf_file( - urdf_path=args.urdf_path, - root_link_name=args.root_link_name, - end_link_name=args.end_link_name, - collision=args.collision_meshes, - target_reduction=args.target_reduction, + urdf_path=urdf_path, + root_link_name=root_link_name, + end_link_name=end_link_name, + collision=collision_meshes, + target_reduction=target_reduction, ) else: robot_data = load_robot_data_from_ros_xacro( - ros_package=args.ros_package, - xacro_path=args.xacro_path, - root_link_name=args.root_link_name, - end_link_name=args.end_link_name, - collision=args.collision_meshes, - target_reduction=args.target_reduction, + ros_package=ros_package, + xacro_path=xacro_path, + root_link_name=root_link_name, + end_link_name=end_link_name, + collision=collision_meshes, + target_reduction=target_reduction, ) robot = Robot.from_robot_data( robot_data=robot_data, batch_size=batch_size, device=device @@ -355,8 +274,8 @@ def main() -> None: ) # repeat joint states and masks for each camera - masks = masks.repeat(args.n_cameras, 1, 1) - joint_states = joint_states.repeat(args.n_cameras, 1) + masks = masks.repeat(n_cameras, 1, 1) + joint_states = joint_states.repeat(n_cameras, 1) if joint_states.shape[0] != batch_size: raise ValueError("Joint states of invalid shape.") scene.robot.configure(joint_states) @@ -372,11 +291,11 @@ def fitness_closure() -> torch.Tensor: renders = scene.observe_from(camera_name).squeeze() fitness = ( soft_dice_loss(renders.unsqueeze(-1), masks.unsqueeze(-1)) - .view(args.n_cameras, n_joint_states) + .view(n_cameras, n_joint_states) .mean(dim=1) ) # show the best particle of the current iteration - if args.display_progress: + if display_progress: offset = 0 current_best_idx = torch.argmin(fitness) current_best_render = ( @@ -408,9 +327,9 @@ def fitness_closure() -> torch.Tensor: # optimize best_particle, _ = particle_swarm_optimizer( fitness_function=fitness_closure, - max_iterations=args.max_iterations, - min_fitness_change=args.min_fitness_change, - max_iterations_below_min_fitness_change=args.max_iterations_below_min_fitness_change, + max_iterations=max_iterations, + min_fitness_change=min_fitness_change, + max_iterations_below_min_fitness_change=max_iterations_below_min_fitness_change, ) # save results @@ -420,8 +339,8 @@ def fitness_closure() -> torch.Tensor: HT_cam_swarm = look_at_from_angle( eye=best_eye, center=best_center, angle=best_angle ) - np.save(path / args.output_file, HT_cam_swarm.cpu().numpy()) + np.save(path / output_file, HT_cam_swarm.cpu().numpy()) if __name__ == "__main__": - main() + app() diff --git a/cli/rr_hydra.py b/cli/rr_hydra.py index 2775933..22fa403 100644 --- a/cli/rr_hydra.py +++ b/cli/rr_hydra.py @@ -1,9 +1,10 @@ -import argparse from pathlib import Path +from typing import Optional import numpy as np import rich import torch +import typer from roboreg.io import ( find_files, @@ -23,146 +24,7 @@ from .util.validate import validate_urdf_source - -def args_factory() -> argparse.Namespace: - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter - ) - parser.add_argument( - "--camera-info-file", - type=str, - required=True, - help="Path to the camera parameters, /camera_info.yaml.", - ) - parser.add_argument("--path", type=str, required=True, help="Path to the data.") - parser.add_argument( - "--mask-pattern", - type=str, - default="image_*_mask.png", - help="Mask file pattern.", - ) - parser.add_argument( - "--depth-pattern", - type=str, - default="depth_*.npy", - help="Depth file pattern. Note that depth values are expected in meters.", - ) - parser.add_argument( - "--joint-states-pattern", - type=str, - default="joint_states_*.npy", - help="Joint state file pattern.", - ) - parser.add_argument( - "--urdf-path", - type=str, - default="test/assets/lbr_med7_r800/description/lbr_med7_r800.urdf", - help="Path to URDF file. Meshes resolved relative to this file. " - "Mutually exclusive with --ros-package/--xacro-path.", - ) - parser.add_argument( - "--ros-package", - type=str, - default=None, - help="ROS package containing robot description. " - "Requires --xacro-path. Mutually exclusive with --urdf-path.", - ) - parser.add_argument( - "--xacro-path", - type=str, - default=None, - help="Path to xacro file relative to --ros-package. " - "Requires --ros-package. Mutually exclusive with --urdf-path.", - ) - parser.add_argument( - "--root-link-name", - type=str, - default="", - help="Root link name. If unspecified, the first link with mesh will be used, which may cause errors.", - ) - parser.add_argument( - "--end-link-name", - type=str, - default="", - help="End link name. If unspecified, the last link with mesh will be used, which may cause errors.", - ) - parser.add_argument( - "--collision-meshes", - action="store_true", - help="If set, collision meshes will be used instead of visual meshes.", - ) - parser.add_argument( - "--depth-conversion-factor", - type=float, - default=1.0, - help="Conversion factor for depth. Computes z = depth / conversion_factor e.g. to covert from millimeter to meter.", - ) - parser.add_argument( - "--z-min", - type=float, - default=0.01, - help="Minimum depth value.", - ) - parser.add_argument( - "--z-max", - type=float, - default=2.0, - help="Maximum depth value.", - ) - parser.add_argument( - "--number-of-points", - type=int, - default=5000, - help="Number of points to sample from robot mesh.", - ) - parser.add_argument( - "--max-distance", - type=float, - default=0.1, - help="Maximum distance between two points to be considered as a correspondence.", - ) - parser.add_argument( - "--outer-max-iter", - type=int, - default=50, - help="Maximum number of outer iterations.", - ) - parser.add_argument( - "--inner-max-iter", - type=int, - default=10, - help="Maximum number of inner iterations.", - ) - parser.add_argument( - "--output-file", - type=str, - default="HT_hydra_robust.npy", - help="Output file name. Relative to the path.", - ) - parser.add_argument( - "--no-boundary", - action="store_true", - help="Do not apply dilation / erosion to the mask.", - ) - parser.add_argument( - "--dilation-kernel-size", - type=int, - default=3, - help="Dilation kernel size for mask boundary. Larger value will result in larger boundary.", - ) - parser.add_argument( - "--erosion-kernel-size", - type=int, - default=10, - help="Erosion kernel size for mask boundary. Larger value will result in larger boundary. The closer the robot, the larger the recommended kernel size.", - ) - parser.add_argument( - "--display-results", - action="store_true", - help="Display point cloud registration results.", - ) - validate_urdf_source(parser, parser.parse_args()) - return parser.parse_args() +app = typer.Typer(add_completion=False) def visualize_hydra_result( @@ -185,55 +47,133 @@ def visualize_hydra_result( ) -def main(): - args = args_factory() +@app.command() +def main( + camera_info_file: Path = typer.Option( + ..., help="Path to the camera parameters, /camera_info.yaml." + ), + path: Path = typer.Option(..., help="Path to the data."), + mask_pattern: str = typer.Option("image_*_mask.png", help="Mask file pattern."), + depth_pattern: str = typer.Option( + "depth_*.npy", + help="Depth file pattern. Note that depth values are expected in meters.", + ), + joint_states_pattern: str = typer.Option( + "joint_states_*.npy", help="Joint state file pattern." + ), + urdf_path: Optional[Path] = typer.Option( + "test/assets/lbr_med7_r800/description/lbr_med7_r800.urdf", + help="Path to URDF file. Meshes resolved relative to this file. " + "Mutually exclusive with --ros-package/--xacro-path.", + ), + ros_package: Optional[str] = typer.Option( + None, + help="ROS package containing robot description. " + "Requires --xacro-path. Mutually exclusive with --urdf-path.", + ), + xacro_path: Optional[str] = typer.Option( + None, + help="Path to xacro file relative to --ros-package. " + "Requires --ros-package. Mutually exclusive with --urdf-path.", + ), + root_link_name: str = typer.Option( + "", + help="Root link name. If unspecified, the first link with mesh will be used, which may cause errors.", + ), + end_link_name: str = typer.Option( + "", + help="End link name. If unspecified, the last link with mesh will be used, which may cause errors.", + ), + collision_meshes: bool = typer.Option( + False, help="If set, collision meshes will be used instead of visual meshes." + ), + depth_conversion_factor: float = typer.Option( + 1.0, + help="Conversion factor for depth. Computes z = depth / conversion_factor e.g. to covert from millimeter to meter.", + ), + z_min: float = typer.Option(0.01, help="Minimum depth value."), + z_max: float = typer.Option(2.0, help="Maximum depth value."), + number_of_points: int = typer.Option( + 5000, help="Number of points to sample from robot mesh." + ), + max_distance: float = typer.Option( + 0.1, + help="Maximum distance between two points to be considered as a correspondence.", + ), + max_outer_iterations: int = typer.Option( + 50, help="Maximum number of outer iterations." + ), + max_inner_iterations: int = typer.Option( + 10, help="Maximum number of inner iterations." + ), + output_file: str = typer.Option( + "HT_hydra_robust.npy", help="Output file name. Relative to the path." + ), + no_boundary: bool = typer.Option( + False, help="Do not apply dilation / erosion to the mask." + ), + dilation_kernel_size: int = typer.Option( + 3, + help="Dilation kernel size for mask boundary. Larger value will result in larger boundary.", + ), + erosion_kernel_size: int = typer.Option( + 10, + help="Erosion kernel size for mask boundary. Larger value will result in larger boundary. The closer the robot, the larger the recommended kernel size.", + ), + display_results: bool = typer.Option( + False, help="Display point cloud registration results." + ), +) -> None: + r"""Hydra robust ICP: point-to-plane ICP registration on a Lie algebra.""" + validate_urdf_source(urdf_path, ros_package, xacro_path) device = "cuda" if torch.cuda.is_available() else "cpu" - path = Path(args.path) # load data observations = parse_hydra_observations( - joint_states_files=find_files(path, args.joint_states_pattern), - mask_files=find_files(path, args.mask_pattern), - depth_files=find_files(path, args.depth_pattern), + joint_states_files=find_files(path, joint_states_pattern), + mask_files=find_files(path, mask_pattern), + depth_files=find_files(path, depth_pattern), ) - _, _, intrinsics = parse_camera_info(args.camera_info_file) + _, _, intrinsics = parse_camera_info(camera_info_file) # load robot specifications - if args.urdf_path is not None: + if urdf_path is not None: robot_data = load_robot_data_from_urdf_file( - urdf_path=args.urdf_path, - root_link_name=args.root_link_name, - end_link_name=args.end_link_name, - collision=args.collision_meshes, + urdf_path=urdf_path, + root_link_name=root_link_name, + end_link_name=end_link_name, + collision=collision_meshes, ) else: robot_data = load_robot_data_from_ros_xacro( - ros_package=args.ros_package, - xacro_path=args.xacro_path, - root_link_name=args.root_link_name, - end_link_name=args.end_link_name, - collision=args.collision_meshes, + ros_package=ros_package, + xacro_path=xacro_path, + root_link_name=root_link_name, + end_link_name=end_link_name, + collision=collision_meshes, ) # register config = HydraRobustICPConfig( HydraConfig( - reference_points_per_mesh=args.number_of_points, + reference_points_per_mesh=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, + z_min=z_min, + z_max=z_max, + depth_conversion_factor=depth_conversion_factor, + use_mask_boundary=not no_boundary, + dilation_kernel_size=dilation_kernel_size, + erosion_kernel_size=erosion_kernel_size, ), - max_correspondence_distance=args.max_distance, - ) + max_correspondence_distance=max_distance, + ), + max_outer_iterations=max_outer_iterations, + max_inner_iterations=max_inner_iterations, ) hydra_robust_icp = HydraRobustICP( config=config, device=device, - on_after_registration=visualize_hydra_result if args.display_results else None, + on_after_registration=visualize_hydra_result if display_results else None, ) rich.print("Entering optimization...") result = hydra_robust_icp( @@ -250,8 +190,8 @@ def main(): # save extrinsics rich.print(f"Writing results to: '{path}'.") - np.save(path / args.output_file, result.extrinsics.cpu().numpy()) + np.save(path / output_file, result.extrinsics.cpu().numpy()) if __name__ == "__main__": - main() + app() diff --git a/cli/rr_mono_dr.py b/cli/rr_mono_dr.py index 82cbdf1..8a97fde 100644 --- a/cli/rr_mono_dr.py +++ b/cli/rr_mono_dr.py @@ -1,10 +1,11 @@ -import argparse import os from pathlib import Path +from typing import Optional import numpy as np import rich import torch +import typer from roboreg.io import ( find_files, @@ -33,197 +34,126 @@ from .util.validate import validate_urdf_source +app = typer.Typer(add_completion=False) -def args_factory() -> argparse.Namespace: - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter + +def print_optimization_state(state: OptimizationState) -> None: + rich.print( + f"Step [{state.iteration} / {state.max_iterations}], " + f"loss: {state.loss:.3f}, " + f"best loss: {state.best_loss:.3f}, " + f"lr: {state.learning_rate:.3e}" ) - parser.add_argument( - "--optimizer", - type=str, - default=DiffRenderingRegistrationConfig().optimizer, + + +@app.command() +def main( + camera_info_file: Path = typer.Option( + ..., + help="Full path to left camera parameters, /left_camera_info.yaml.", + ), + extrinsics_file: Path = typer.Option( + ..., + help="Full path to homogeneous transforms from base to left camera frame, /HT_hydra_robust.npy.", + ), + path: Path = typer.Option(..., help="Path to the data."), + optimizer: str = typer.Option( + DiffRenderingRegistrationConfig().optimizer, help="Optimizer to use, e.g. 'Adam' or 'SGD'. Imported from torch.optim.", - ) - parser.add_argument( - "--lr", - type=float, - default=DiffRenderingRegistrationConfig().lr, - help="Learning rate for the optimizer.", - ) - parser.add_argument( - "--max-iterations", - type=int, - default=ConvergenceConfig().max_iterations, + ), + lr: float = typer.Option( + DiffRenderingRegistrationConfig().lr, help="Learning rate for the optimizer." + ), + max_iterations: int = typer.Option( + ConvergenceConfig().max_iterations, help="Maximum number of epochs to optimize for.", - ) - parser.add_argument( - "--convergence-tolerance", - type=float, - default=ConvergenceConfig().tolerance, - ) - parser.add_argument( - "--convergence-patience", - type=int, - default=ConvergenceConfig().patience, - ) - parser.add_argument( - "--scheduler-factor", - type=float, - default=PlateauSchedulerConfig().factor, - ) - parser.add_argument( - "--scheduler-patience", - type=int, - default=PlateauSchedulerConfig().patience, - ) - parser.add_argument( - "--scheduler-threshold", - type=float, - default=PlateauSchedulerConfig().threshold, - ) - parser.add_argument( - "--rendering-objective", - type=RenderingObjectiveType, - choices=list(RenderingObjectiveType), - default=RenderingObjectiveType.DISTANCE_MAP, - help="Rendering objective.", - ) - parser.add_argument( - "--display-progress", - action="store_true", - help="Display optimization progress.", - ) - parser.add_argument( - "--urdf-path", - type=str, - default="test/assets/lbr_med7_r800/description/lbr_med7_r800.urdf", + ), + convergence_tolerance: float = typer.Option(ConvergenceConfig().tolerance), + convergence_patience: int = typer.Option(ConvergenceConfig().patience), + scheduler_factor: float = typer.Option(PlateauSchedulerConfig().factor), + scheduler_patience: int = typer.Option(PlateauSchedulerConfig().patience), + scheduler_threshold: float = typer.Option(PlateauSchedulerConfig().threshold), + rendering_objective: RenderingObjectiveType = typer.Option( + RenderingObjectiveType.DISTANCE_MAP, help="Rendering objective." + ), + display_progress: bool = typer.Option(False, help="Display optimization progress."), + urdf_path: Optional[Path] = typer.Option( + "test/assets/lbr_med7_r800/description/lbr_med7_r800.urdf", help="Path to URDF file. Meshes resolved relative to this file. " "Mutually exclusive with --ros-package/--xacro-path.", - ) - parser.add_argument( - "--ros-package", - type=str, - default=None, + ), + ros_package: Optional[str] = typer.Option( + None, help="ROS package containing robot description. " "Requires --xacro-path. Mutually exclusive with --urdf-path.", - ) - parser.add_argument( - "--xacro-path", - type=str, - default=None, + ), + xacro_path: Optional[str] = typer.Option( + None, help="Path to xacro file relative to --ros-package. " "Requires --ros-package. Mutually exclusive with --urdf-path.", - ) - parser.add_argument( - "--root-link-name", - type=str, - default="", + ), + root_link_name: str = typer.Option( + "", help="Root link name. If unspecified, the first link with mesh will be used, which may cause errors.", - ) - parser.add_argument( - "--end-link-name", - type=str, - default="", + ), + end_link_name: str = typer.Option( + "", help="End link name. If unspecified, the last link with mesh will be used, which may cause errors.", - ) - parser.add_argument( - "--collision-meshes", - action="store_true", - help="If set, collision meshes will be used instead of visual meshes.", - ) - parser.add_argument( - "--camera-info-file", - type=str, - required=True, - help="Full path to left camera parameters, /left_camera_info.yaml.", - ) - parser.add_argument( - "--extrinsics-file", - type=str, - required=True, - help="Full path to homogeneous transforms from base to left camera frame, /HT_hydra_robust.npy.", - ) - parser.add_argument("--path", type=str, required=True, help="Path to the data.") - parser.add_argument( - "--image-pattern", - type=str, - default="left_image_*.png", - help="Left image file pattern.", - ) - parser.add_argument( - "--joint-states-pattern", - type=str, - default="joint_states_*.npy", - help="Joint state file pattern.", - ) - parser.add_argument( - "--mask-pattern", - type=str, - default="left_mask_*.png", - help="Left mask file pattern.", - ) - parser.add_argument( - "--output-file", - type=str, - default="HT_left_dr.npy", - help="Left output file name. Relative to --path.", - ) - parser.add_argument( - "--max-jobs", - type=int, - default=2, + ), + collision_meshes: bool = typer.Option( + False, help="If set, collision meshes will be used instead of visual meshes." + ), + image_pattern: str = typer.Option( + "left_image_*.png", help="Left image file pattern." + ), + joint_states_pattern: str = typer.Option( + "joint_states_*.npy", help="Joint state file pattern." + ), + mask_pattern: str = typer.Option("left_mask_*.png", help="Left mask file pattern."), + output_file: str = typer.Option( + "HT_left_dr.npy", help="Left output file name. Relative to --path." + ), + max_jobs: int = typer.Option( + 2, help="Number of concurrent compilation jobs for nvdiffrast. Only relevant on first run.", - ) - validate_urdf_source(parser, parser.parse_args()) - return parser.parse_args() - - -def print_optimization_state(state: OptimizationState) -> None: - rich.print( - f"Step [{state.iteration} / {state.max_iterations}], " - f"loss: {state.loss:.3f}, " - f"best loss: {state.best_loss:.3f}, " - f"lr: {state.learning_rate:.3e}" - ) - - -def main() -> None: - args = args_factory() + ), +) -> None: + r"""Monocular differentiable rendering registration.""" + validate_urdf_source(urdf_path, ros_package, xacro_path) device = "cuda" if torch.cuda.is_available() else "cpu" - os.environ["MAX_JOBS"] = str(args.max_jobs) # limit number of concurrent jobs - path = Path(args.path) + os.environ["MAX_JOBS"] = str(max_jobs) # limit number of concurrent jobs # load data observations = parse_monocular_observations( - image_files=find_files(path, args.image_pattern), - joint_states_files=find_files(path, args.joint_states_pattern), - target_files=find_files(path, args.mask_pattern), + image_files=find_files(path, image_pattern), + joint_states_files=find_files(path, joint_states_pattern), + target_files=find_files(path, mask_pattern), ) - _, _, intrinsics = parse_camera_info(args.camera_info_file) - extrinsics = np.load(args.extrinsics_file) + _, _, intrinsics = parse_camera_info(camera_info_file) + extrinsics = np.load(extrinsics_file) # load robot specifications - if args.urdf_path is not None: + if urdf_path is not None: robot_data = load_robot_data_from_urdf_file( - urdf_path=args.urdf_path, - root_link_name=args.root_link_name, - end_link_name=args.end_link_name, - collision=args.collision_meshes, + urdf_path=urdf_path, + root_link_name=root_link_name, + end_link_name=end_link_name, + collision=collision_meshes, ) else: robot_data = load_robot_data_from_ros_xacro( - ros_package=args.ros_package, - xacro_path=args.xacro_path, - root_link_name=args.root_link_name, - end_link_name=args.end_link_name, - collision=args.collision_meshes, + ros_package=ros_package, + xacro_path=xacro_path, + root_link_name=root_link_name, + end_link_name=end_link_name, + collision=collision_meshes, ) # register on_iteration: list[OptimizationCallback] = [ print_optimization_state, ] - if args.display_progress: + if display_progress: on_iteration.append( RenderOverlayCallback( images={ @@ -236,21 +166,21 @@ def main() -> None: diff_rendering_registration = DiffRenderingRegistration( config=DiffRenderingRegistrationConfig( camera=CameraConfig(), - optimizer=args.optimizer, - lr=args.lr, + optimizer=optimizer, + lr=lr, convergence=ConvergenceConfig( - max_iterations=args.max_iterations, - tolerance=args.convergence_tolerance, - patience=args.convergence_patience, + max_iterations=max_iterations, + tolerance=convergence_tolerance, + patience=convergence_patience, ), plateau_scheduler=PlateauSchedulerConfig( mode="min", - factor=args.scheduler_factor, - patience=args.scheduler_patience, - threshold=args.scheduler_threshold, + factor=scheduler_factor, + patience=scheduler_patience, + threshold=scheduler_threshold, ), ), - objective=create_rendering_objective(objective_type=args.rendering_objective), + objective=create_rendering_objective(objective_type=rendering_objective), device=device, on_iteration=on_iteration, ) @@ -275,10 +205,10 @@ def main() -> None: # save extrinsics rich.print(f"Writing results to: '{path}'.") np.save( - path / args.output_file, + path / output_file, result.extrinsics.cpu().numpy(), ) if __name__ == "__main__": - main() + app() diff --git a/cli/rr_render.py b/cli/rr_render.py index a72252a..918fe8d 100644 --- a/cli/rr_render.py +++ b/cli/rr_render.py @@ -1,19 +1,16 @@ -import argparse import os +from enum import Enum from pathlib import Path +from typing import Optional import cv2 import numpy as np import torch +import typer from rich import progress from torch.utils.data import DataLoader -from roboreg.core import ( - NVDiffRastRenderer, - Robot, - RobotScene, - VirtualCamera, -) +from roboreg.core import NVDiffRastRenderer, Robot, RobotScene, VirtualCamera from roboreg.io import ( MonocularDataset, load_robot_data_from_ros_xacro, @@ -23,139 +20,98 @@ from .util.validate import validate_urdf_source +app = typer.Typer(add_completion=False) -def args_factory() -> argparse.Namespace: - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter - ) - parser.add_argument( - "--batch-size", - type=int, - default=1, + +class OverlayColor(str, Enum): + RED = "r" + GREEN = "g" + BLUE = "b" + + +@app.command() +def main( + camera_info_file: Path = typer.Option( + ..., help="Path to the camera parameters, /camera_info.yaml." + ), + extrinsics_file: Path = typer.Option( + ..., + help="Homogeneous transform from base to camera frame, /HT_hydra_robust.npy.", + ), + images_path: Path = typer.Option(..., help="Path to the images."), + joint_states_path: Path = typer.Option(..., help="Path to the joint states."), + output_path: Path = typer.Option(..., help="Output path."), + batch_size: int = typer.Option( + 1, help="Batch size for rendering. For batch_size > 1, the last batch may be dropped.", - ) - parser.add_argument( - "--num-workers", type=int, default=0, help="Number of workers for data loading." - ) - parser.add_argument( - "--urdf-path", - type=str, - default="test/assets/lbr_med7_r800/description/lbr_med7_r800.urdf", + ), + num_workers: int = typer.Option(0, help="Number of workers for data loading."), + urdf_path: Optional[Path] = typer.Option( + "test/assets/lbr_med7_r800/description/lbr_med7_r800.urdf", help="Path to URDF file. Meshes resolved relative to this file. " "Mutually exclusive with --ros-package/--xacro-path.", - ) - parser.add_argument( - "--ros-package", - type=str, - default=None, + ), + ros_package: Optional[str] = typer.Option( + None, help="ROS package containing robot description. " "Requires --xacro-path. Mutually exclusive with --urdf-path.", - ) - parser.add_argument( - "--xacro-path", - type=str, - default=None, + ), + xacro_path: Optional[str] = typer.Option( + None, help="Path to xacro file relative to --ros-package. " "Requires --ros-package. Mutually exclusive with --urdf-path.", - ) - parser.add_argument( - "--root-link-name", - type=str, - default="", + ), + root_link_name: str = typer.Option( + "", help="Root link name. If unspecified, the first link with mesh will be used, which may cause errors.", - ) - parser.add_argument( - "--end-link-name", - type=str, - default="", + ), + end_link_name: str = typer.Option( + "", help="End link name. If unspecified, the last link with mesh will be used, which may cause errors.", - ) - parser.add_argument( - "--collision-meshes", - action="store_true", - help="If set, collision meshes will be used instead of visual meshes.", - ) - parser.add_argument( - "--camera-info-file", - type=str, - required=True, - help="Path to the camera parameters, /camera_info.yaml.", - ) - parser.add_argument( - "--extrinsics-file", - type=str, - required=True, - help="Homogeneous transform from base to camera frame, /HT_hydra_robust.npy.", - ) - parser.add_argument( - "--images-path", type=str, required=True, help="Path to the images." - ) - parser.add_argument( - "--joint-states-path", type=str, required=True, help="Path to the joint states." - ) - parser.add_argument( - "--image-pattern", - type=str, - default="image_*.png", - help="Image file pattern.", - ) - parser.add_argument( - "--joint-states-pattern", - type=str, - default="joint_states_*.npy", - help="Joint state file pattern.", - ) - parser.add_argument( - "--output-path", - type=str, - required=True, - help="Output path.", - ) - parser.add_argument( - "--color", - type=str, - choices=["r", "g", "b"], - default="b", - help="Color channel to overlay the render.", - ) - parser.add_argument( - "--max-jobs", - type=int, - default=2, + ), + collision_meshes: bool = typer.Option( + False, help="If set, collision meshes will be used instead of visual meshes." + ), + image_pattern: str = typer.Option("image_*.png", help="Image file pattern."), + joint_states_pattern: str = typer.Option( + "joint_states_*.npy", help="Joint state file pattern." + ), + color: OverlayColor = typer.Option( + OverlayColor.BLUE, help="Color channel to overlay the render." + ), + max_jobs: int = typer.Option( + 2, help="Number of concurrent compilation jobs for nvdiffrast. Only relevant on first run.", - ) - validate_urdf_source(parser, parser.parse_args()) - return parser.parse_args() - - -def main(): - args = args_factory() + ), +) -> None: + r"""Render robot mesh overlays for a set of images given known extrinsics.""" + validate_urdf_source(urdf_path, ros_package, xacro_path) device = "cuda" if torch.cuda.is_available() else "cpu" - os.environ["MAX_JOBS"] = str(args.max_jobs) # limit number of concurrent jobs + os.environ["MAX_JOBS"] = str(max_jobs) # limit number of concurrent jobs camera = { "camera": VirtualCamera.from_camera_configs( - camera_info_file=args.camera_info_file, - extrinsics_file=args.extrinsics_file, + camera_info_file=camera_info_file, + extrinsics_file=extrinsics_file, device=device, ) } - if args.urdf_path is not None: + if urdf_path is not None: robot_data = load_robot_data_from_urdf_file( - urdf_path=args.urdf_path, - root_link_name=args.root_link_name, - end_link_name=args.end_link_name, - collision=args.collision_meshes, + urdf_path=urdf_path, + root_link_name=root_link_name, + end_link_name=end_link_name, + collision=collision_meshes, ) else: robot_data = load_robot_data_from_ros_xacro( - ros_package=args.ros_package, - xacro_path=args.xacro_path, - root_link_name=args.root_link_name, - end_link_name=args.end_link_name, - collision=args.collision_meshes, + ros_package=ros_package, + xacro_path=xacro_path, + root_link_name=root_link_name, + end_link_name=end_link_name, + collision=collision_meshes, ) robot = Robot.from_robot_data( - robot_data=robot_data, batch_size=args.batch_size, device=device + robot_data=robot_data, batch_size=batch_size, device=device ) scene = RobotScene( cameras=camera, @@ -163,20 +119,19 @@ def main(): renderer=NVDiffRastRenderer(device=device), ) dataset = MonocularDataset( - images_path=args.images_path, - image_pattern=args.image_pattern, - joint_states_path=args.joint_states_path, - joint_states_pattern=args.joint_states_pattern, + images_path=images_path, + image_pattern=image_pattern, + joint_states_path=joint_states_path, + joint_states_pattern=joint_states_pattern, ) dataloader = DataLoader( dataset, - batch_size=args.batch_size, + batch_size=batch_size, shuffle=False, drop_last=True, - num_workers=args.num_workers, + num_workers=num_workers, ) - output_path = Path(args.output_path) if not output_path.exists(): output_path.mkdir(parents=True) @@ -202,9 +157,9 @@ def main(): ) cv2.imwrite( output_file, - overlay_mask(image, render, args.color, scale=1.0), + overlay_mask(image, render, color.value, scale=1.0), ) if __name__ == "__main__": - main() + app() diff --git a/cli/rr_sam2.py b/cli/rr_sam2.py index 455969e..b587311 100644 --- a/cli/rr_sam2.py +++ b/cli/rr_sam2.py @@ -1,7 +1,8 @@ -import argparse +from pathlib import Path import cv2 import numpy as np +import typer from rich import progress from roboreg.detector import OpenCVDetector @@ -9,58 +10,37 @@ from roboreg.segmentor import Sam2Segmentor from roboreg.util import overlay_mask - -def args_factory() -> argparse.Namespace: - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter - ) - parser.add_argument("--path", type=str, required=True, help="Path to the images.") - parser.add_argument( - "--pattern", type=str, default="image_*.png", help="Image file pattern." - ) - parser.add_argument( - "--n-positive-samples", type=int, default=5, help="Number of positive samples." - ) - parser.add_argument( - "--n-negative-samples", type=int, default=5, help="Number of negative samples." - ) - parser.add_argument( - "--model-id", - type=str, - default="facebook/sam2-hiera-large", - help="Hugging Face model ID.", - ) - parser.add_argument( - "--device", - type=str, - default="cuda", - help="Device to run the model. Default: cuda", - ) - parser.add_argument( - "--pre-annotated", - action="store_true", - help="Try to read annotations.", - ) - return parser.parse_args() +app = typer.Typer(add_completion=False) -def main(): - args = args_factory() - image_files = find_files(args.path, args.pattern) +@app.command() +def main( + path: Path = typer.Option(..., help="Path to the images."), + pattern: str = typer.Option("image_*.png", help="Image file pattern."), + n_positive_samples: int = typer.Option(5, help="Number of positive samples."), + n_negative_samples: int = typer.Option(5, help="Number of negative samples."), + model_id: str = typer.Option( + "facebook/sam2-hiera-large", help="Hugging Face model ID." + ), + device: str = typer.Option("cuda", help="Device to run the model. Default: cuda"), + pre_annotated: bool = typer.Option(False, help="Try to read annotations."), +) -> None: + r"""Generate robot masks with SAM2, seeded by OpenCV-detected samples.""" + image_files = find_files(path, pattern) # detect detector = OpenCVDetector( - n_negative_samples=args.n_negative_samples, - n_positive_samples=args.n_positive_samples, + n_negative_samples=n_negative_samples, + n_positive_samples=n_positive_samples, ) # segment - segmentor = Sam2Segmentor(model_id=args.model_id, device=args.device) + segmentor = Sam2Segmentor(model_id=model_id, device=device) for image_file in progress.track(image_files, description="Generating masks..."): img = cv2.imread(image_file) annotations = False - if args.pre_annotated: + if pre_annotated: try: samples, labels = detector.read( path=image_file.parent / f"{image_file.stem}_samples.csv" @@ -90,4 +70,4 @@ def main(): if __name__ == "__main__": - main() + app() diff --git a/cli/rr_stereo_dr.py b/cli/rr_stereo_dr.py index 21ae4ce..724d440 100644 --- a/cli/rr_stereo_dr.py +++ b/cli/rr_stereo_dr.py @@ -1,10 +1,11 @@ -import argparse import os from pathlib import Path +from typing import Optional import numpy as np import rich import torch +import typer from roboreg.io import ( find_files, @@ -33,232 +34,150 @@ from .util.validate import validate_urdf_source +app = typer.Typer(add_completion=False) -def args_factory() -> argparse.Namespace: - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter + +def print_optimization_state(state: OptimizationState) -> None: + rich.print( + f"Step [{state.iteration} / {state.max_iterations}], " + f"loss: {state.loss:.3f}, " + f"best loss: {state.best_loss:.3f}, " + f"lr: {state.learning_rate:.3e}" ) - parser.add_argument( - "--optimizer", - type=str, - default=DiffRenderingRegistrationConfig().optimizer, + + +@app.command() +def main( + left_camera_info_file: Path = typer.Option( + ..., + help="Full path to left camera parameters, /left_camera_info.yaml.", + ), + right_camera_info_file: Path = typer.Option( + ..., + help="Full path to right camera parameters, /right_camera_info.yaml.", + ), + left_extrinsics_file: Path = typer.Option( + ..., + help="Full path to homogeneous transforms from base to left camera frame, /HT_hydra_robust.npy.", + ), + right_extrinsics_file: Path = typer.Option( + ..., + help="Full path to homogeneous transforms from base to right camera frame, /HT_right_to_left.npy.", + ), + path: Path = typer.Option(..., help="Path to the data."), + optimizer: str = typer.Option( + DiffRenderingRegistrationConfig().optimizer, help="Optimizer to use, e.g. 'Adam' or 'SGD'. Imported from torch.optim.", - ) - parser.add_argument( - "--lr", - type=float, - default=DiffRenderingRegistrationConfig().lr, - help="Learning rate for the optimizer.", - ) - parser.add_argument( - "--max-iterations", - type=int, - default=ConvergenceConfig().max_iterations, + ), + lr: float = typer.Option( + DiffRenderingRegistrationConfig().lr, help="Learning rate for the optimizer." + ), + max_iterations: int = typer.Option( + ConvergenceConfig().max_iterations, help="Maximum number of epochs to optimize for.", - ) - parser.add_argument( - "--convergence-tolerance", - type=float, - default=ConvergenceConfig().tolerance, - ) - parser.add_argument( - "--convergence-patience", - type=int, - default=ConvergenceConfig().patience, - ) - parser.add_argument( - "--scheduler-factor", - type=float, - default=PlateauSchedulerConfig().factor, - ) - parser.add_argument( - "--scheduler-patience", - type=int, - default=PlateauSchedulerConfig().patience, - ) - parser.add_argument( - "--scheduler-threshold", - type=float, - default=PlateauSchedulerConfig().threshold, - ) - parser.add_argument( - "--rendering-objective", - type=RenderingObjectiveType, - choices=list(RenderingObjectiveType), - default=RenderingObjectiveType.DISTANCE_MAP, - help="Rendering objective.", - ) - parser.add_argument( - "--display-progress", - action="store_true", - help="Display optimization progress.", - ) - parser.add_argument( - "--urdf-path", - type=str, - default="test/assets/lbr_med7_r800/description/lbr_med7_r800.urdf", + ), + convergence_tolerance: float = typer.Option(ConvergenceConfig().tolerance), + convergence_patience: int = typer.Option(ConvergenceConfig().patience), + scheduler_factor: float = typer.Option(PlateauSchedulerConfig().factor), + scheduler_patience: int = typer.Option(PlateauSchedulerConfig().patience), + scheduler_threshold: float = typer.Option(PlateauSchedulerConfig().threshold), + rendering_objective: RenderingObjectiveType = typer.Option( + RenderingObjectiveType.DISTANCE_MAP, help="Rendering objective." + ), + display_progress: bool = typer.Option(False, help="Display optimization progress."), + urdf_path: Optional[Path] = typer.Option( + "test/assets/lbr_med7_r800/description/lbr_med7_r800.urdf", help="Path to URDF file. Meshes resolved relative to this file. " "Mutually exclusive with --ros-package/--xacro-path.", - ) - parser.add_argument( - "--ros-package", - type=str, - default=None, + ), + ros_package: Optional[str] = typer.Option( + None, help="ROS package containing robot description. " "Requires --xacro-path. Mutually exclusive with --urdf-path.", - ) - parser.add_argument( - "--xacro-path", - type=str, - default=None, + ), + xacro_path: Optional[str] = typer.Option( + None, help="Path to xacro file relative to --ros-package. " "Requires --ros-package. Mutually exclusive with --urdf-path.", - ) - parser.add_argument( - "--root-link-name", - type=str, - default="", + ), + root_link_name: str = typer.Option( + "", help="Root link name. If unspecified, the first link with mesh will be used, which may cause errors.", - ) - parser.add_argument( - "--end-link-name", - type=str, - default="", + ), + end_link_name: str = typer.Option( + "", help="End link name. If unspecified, the last link with mesh will be used, which may cause errors.", - ) - parser.add_argument( - "--collision-meshes", - action="store_true", - help="If set, collision meshes will be used instead of visual meshes.", - ) - parser.add_argument( - "--left-camera-info-file", - type=str, - required=True, - help="Full path to left camera parameters, /left_camera_info.yaml.", - ) - parser.add_argument( - "--right-camera-info-file", - type=str, - required=True, - help="Full path to right camera parameters, /right_camera_info.yaml.", - ) - parser.add_argument( - "--left-extrinsics-file", - type=str, - required=True, - help="Full path to homogeneous transforms from base to left camera frame, /HT_hydra_robust.npy.", - ) - parser.add_argument( - "--right-extrinsics-file", - type=str, - required=True, - help="Full path to homogeneous transforms from base to right camera frame, /HT_right_to_left.npy.", - ) - parser.add_argument("--path", type=str, required=True, help="Path to the data.") - parser.add_argument( - "--left-image-pattern", - type=str, - default="left_image_*.png", - help="Left image file pattern.", - ) - parser.add_argument( - "--right-image-pattern", - type=str, - default="right_image_*.png", - help="Right image file pattern.", - ) - parser.add_argument( - "--joint-states-pattern", - type=str, - default="joint_states_*.npy", - help="Joint state file pattern.", - ) - parser.add_argument( - "--left-mask-pattern", - type=str, - default="left_mask_*.png", - help="Left mask file pattern.", - ) - parser.add_argument( - "--right-mask-pattern", - type=str, - default="right_mask_*.png", - help="Right mask file pattern.", - ) - parser.add_argument( - "--left-output-file", - type=str, - default="HT_left_dr.npy", - help="Left output file name. Relative to --path.", - ) - parser.add_argument( - "--right-output-file", - type=str, - default="HT_right_dr.npy", - help="Right output file name. Relative to --path.", - ) - parser.add_argument( - "--max-jobs", - type=int, - default=2, + ), + collision_meshes: bool = typer.Option( + False, help="If set, collision meshes will be used instead of visual meshes." + ), + left_image_pattern: str = typer.Option( + "left_image_*.png", help="Left image file pattern." + ), + right_image_pattern: str = typer.Option( + "right_image_*.png", help="Right image file pattern." + ), + joint_states_pattern: str = typer.Option( + "joint_states_*.npy", help="Joint state file pattern." + ), + left_mask_pattern: str = typer.Option( + "left_mask_*.png", help="Left mask file pattern." + ), + right_mask_pattern: str = typer.Option( + "right_mask_*.png", help="Right mask file pattern." + ), + left_output_file: str = typer.Option( + "HT_left_dr.npy", help="Left output file name. Relative to --path." + ), + right_output_file: str = typer.Option( + "HT_right_dr.npy", help="Right output file name. Relative to --path." + ), + max_jobs: int = typer.Option( + 2, help="Number of concurrent compilation jobs for nvdiffrast. Only relevant on first run.", - ) - validate_urdf_source(parser, parser.parse_args()) - return parser.parse_args() - - -def print_optimization_state(state: OptimizationState) -> None: - rich.print( - f"Step [{state.iteration} / {state.max_iterations}], " - f"loss: {state.loss:.3f}, " - f"best loss: {state.best_loss:.3f}, " - f"lr: {state.learning_rate:.3e}" - ) - - -def main() -> None: - args = args_factory() + ), +) -> None: + r"""Stereo differentiable rendering registration.""" + validate_urdf_source(urdf_path, ros_package, xacro_path) device = "cuda" if torch.cuda.is_available() else "cpu" - os.environ["MAX_JOBS"] = str(args.max_jobs) # limit number of concurrent jobs - path = Path(args.path) + os.environ["MAX_JOBS"] = str(max_jobs) # limit number of concurrent jobs # load data observations = parse_stereo_observations( - left_image_files=find_files(path, args.left_image_pattern), - right_image_files=find_files(path, args.right_image_pattern), - joint_states_files=find_files(path, args.joint_states_pattern), - left_target_files=find_files(path, args.left_mask_pattern), - right_target_files=find_files(path, args.right_mask_pattern), + left_image_files=find_files(path, left_image_pattern), + right_image_files=find_files(path, right_image_pattern), + joint_states_files=find_files(path, joint_states_pattern), + left_target_files=find_files(path, left_mask_pattern), + right_target_files=find_files(path, right_mask_pattern), ) - _, _, left_intrinsics = parse_camera_info(args.left_camera_info_file) - _, _, right_intrinsics = parse_camera_info(args.right_camera_info_file) - extrinsics = np.load(args.left_extrinsics_file) - right_extrinsics = np.load(args.right_extrinsics_file) + _, _, left_intrinsics = parse_camera_info(left_camera_info_file) + _, _, right_intrinsics = parse_camera_info(right_camera_info_file) + extrinsics = np.load(left_extrinsics_file) + right_extrinsics = np.load(right_extrinsics_file) # load robot specifications - if args.urdf_path is not None: + if urdf_path is not None: robot_data = load_robot_data_from_urdf_file( - urdf_path=args.urdf_path, - root_link_name=args.root_link_name, - end_link_name=args.end_link_name, - collision=args.collision_meshes, + urdf_path=urdf_path, + root_link_name=root_link_name, + end_link_name=end_link_name, + collision=collision_meshes, ) else: robot_data = load_robot_data_from_ros_xacro( - ros_package=args.ros_package, - xacro_path=args.xacro_path, - root_link_name=args.root_link_name, - end_link_name=args.end_link_name, - collision=args.collision_meshes, + ros_package=ros_package, + xacro_path=xacro_path, + root_link_name=root_link_name, + end_link_name=end_link_name, + collision=collision_meshes, ) # register on_iteration: list[OptimizationCallback] = [ print_optimization_state, ] - if args.display_progress: + if display_progress: on_iteration.append( RenderOverlayCallback( images={ @@ -271,21 +190,21 @@ def main() -> None: diff_rendering_registration = DiffRenderingRegistration( config=DiffRenderingRegistrationConfig( camera=CameraConfig(), - optimizer=args.optimizer, - lr=args.lr, + optimizer=optimizer, + lr=lr, convergence=ConvergenceConfig( - max_iterations=args.max_iterations, - tolerance=args.convergence_tolerance, - patience=args.convergence_patience, + max_iterations=max_iterations, + tolerance=convergence_tolerance, + patience=convergence_patience, ), plateau_scheduler=PlateauSchedulerConfig( mode="min", - factor=args.scheduler_factor, - patience=args.scheduler_patience, - threshold=args.scheduler_threshold, + factor=scheduler_factor, + patience=scheduler_patience, + threshold=scheduler_threshold, ), ), - objective=create_rendering_objective(objective_type=args.rendering_objective), + objective=create_rendering_objective(objective_type=rendering_objective), device=device, on_iteration=[print_optimization_state], ) @@ -314,14 +233,14 @@ def main() -> None: # save extrinsics rich.print(f"Writing results to: '{path}'.") np.save( - path / args.left_output_file, + path / left_output_file, result.extrinsics.cpu().numpy(), ) np.save( - path / args.right_output_file, + path / right_output_file, result.extrinsics.cpu().numpy() @ right_extrinsics, ) if __name__ == "__main__": - main() + app() diff --git a/cli/util/validate.py b/cli/util/validate.py index 03477af..ae30f7d 100644 --- a/cli/util/validate.py +++ b/cli/util/validate.py @@ -1,36 +1,44 @@ -import argparse from pathlib import Path +from typing import Optional +import typer -def validate_urdf_source(parser: argparse.ArgumentParser, args: argparse.Namespace): + +def validate_urdf_source( + urdf_path: Optional[Path], + ros_package: Optional[str], + xacro_path: Optional[str], +) -> None: r"""Validate mutually exclusive URDF source options.""" - urdf_provided = args.urdf_path is not None - ros_provided = args.ros_package is not None - xacro_provided = args.xacro_path is not None + urdf_provided = urdf_path is not None + ros_provided = ros_package is not None + xacro_provided = xacro_path is not None + + def _error(message: str) -> None: + typer.echo(f"Error: {message}", err=True) + raise typer.Exit(code=2) # check if both methods provided if urdf_provided and (ros_provided or xacro_provided): - parser.error( + _error( "Cannot specify --urdf-path together with --ros-package or --xacro-path. " "Use either --urdf-path OR (--ros-package + --xacro-path)." ) # check if ROS method incomplete if ros_provided and not xacro_provided: - parser.error("--ros-package requires --xacro-path") + _error("--ros-package requires --xacro-path") if xacro_provided and not ros_provided: - parser.error("--xacro-path requires --ros-package") + _error("--xacro-path requires --ros-package") # check if nothing provided if not urdf_provided and not ros_provided: - parser.error( + _error( "Must specify URDF source: either --urdf-path OR " "(--ros-package + --xacro-path)" ) # validate file exists if using urdf-path - if urdf_provided: - urdf_path = Path(args.urdf_path) - if not urdf_path.exists(): - parser.error(f"URDF file not found: {urdf_path}") + if urdf_provided and not urdf_path.exists(): + _error(f"URDF file not found: {urdf_path}") diff --git a/pyproject.toml b/pyproject.toml index 208e796..b80cb3e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,7 @@ dependencies = [ "torch", "transformations", "trimesh", + "typer", ] classifiers = [ "Programming Language :: Python :: 3", @@ -44,12 +45,12 @@ Homepage = "https://github.com/lbr-stack/roboreg" Issues = "https://github.com/lbr-stack/roboreg/issues" [project.scripts] -rr-cam-swarm = "cli.rr_cam_swarm:main" -rr-hydra = "cli.rr_hydra:main" -rr-mono-dr = "cli.rr_mono_dr:main" -rr-render = "cli.rr_render:main" -rr-sam2 = "cli.rr_sam2:main" -rr-stereo-dr = "cli.rr_stereo_dr:main" +rr-cam-swarm = "cli.rr_cam_swarm:app" +rr-hydra = "cli.rr_hydra:app" +rr-mono-dr = "cli.rr_mono_dr:app" +rr-render = "cli.rr_render:app" +rr-sam2 = "cli.rr_sam2:app" +rr-stereo-dr = "cli.rr_stereo_dr:app" [tool.setuptools.packages.find] where = ["."]