robot/scene instantiation refactor (#110)

Author: mhubii

cli/rr_cam_swarm.py

@@ -6,8 +6,20 @@
 import numpy as np
 import torch
 
-from roboreg import differentiable as rrd
-from roboreg.io import URDFParser, find_files, parse_camera_info, parse_mono_data
+from roboreg.core import (
+    NVDiffRastRenderer,
+    Robot,
+    RobotScene,
+    TorchKinematics,
+    TorchMeshContainer,
+    VirtualCamera,
+)
+from roboreg.io import (
+    find_files,
+    load_robot_data_from_ros_xacro,
+    parse_camera_info,
+    parse_mono_data,
+)
 from roboreg.losses import soft_dice_loss
 from roboreg.optim import LinearParticleSwarm, ParticleSwarmOptimizer
 from roboreg.util import (
@@ -295,31 +307,43 @@ def main() -> None:
         n_joint_states * args.n_cameras
     )  # (each camera observes n_joint_states joint states)
     camera_name = "camera"
-    camera = rrd.VirtualCamera(
+    camera = VirtualCamera(
         resolution=(height, width),
         intrinsics=intrinsics,
         extrinsics=torch.eye(4, device=device).unsqueeze(0).expand(batch_size, -1, -1),
         device=device,
     )
 
-    urdf_parser = URDFParser.from_ros_xacro(
-        ros_package=args.ros_package, xacro_path=args.xacro_path
-    )
-    robot = rrd.Robot.from_urdf_parser(
-        urdf_parser=urdf_parser,
+    # instantiate robot
+    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,
+    )
+    mesh_container = TorchMeshContainer(
+        meshes=robot_data.meshes,
         batch_size=batch_size,
         device=device,
-        target_reduction=args.target_reduction,  # reduce mesh vertex count for memory reduction
+    )
+    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,
     )
 
-    renderer = rrd.NVDiffRastRenderer(device=device)
-    scene = rrd.RobotScene(
+    # instantiate scene
+    scene = RobotScene(
         cameras={camera_name: camera},
         robot=robot,
-        renderer=renderer,
+        renderer=NVDiffRastRenderer(device=device),
     )
 
     # repeat joint states and masks for each camera

cli/rr_hydra.py

@@ -2,12 +2,16 @@
 import os
 
 import numpy as np
-import rich
 import torch
 
-from roboreg.differentiable import Robot
+from roboreg.core import Robot, TorchKinematics, TorchMeshContainer
 from roboreg.hydra_icp import hydra_centroid_alignment, hydra_robust_icp
-from roboreg.io import URDFParser, find_files, parse_camera_info, parse_hydra_data
+from roboreg.io import (
+    find_files,
+    load_robot_data_from_ros_xacro,
+    parse_camera_info,
+    parse_hydra_data,
+)
 from roboreg.util import (
     RegistrationVisualizer,
     clean_xyz,
@@ -161,35 +165,29 @@ def main():
     )
     height, width, intrinsics = parse_camera_info(args.camera_info_file)
 
-    # instantiate kinematics
-    urdf_parser = URDFParser.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
-    if root_link_name == "":
-        root_link_name = urdf_parser.link_names_with_meshes(
-            collision=args.collision_meshes
-        )[0]
-        rich.print(
-            f"Root link name not provided. Using the first link with mesh: '{root_link_name}'."
-        )
-    if end_link_name == "":
-        end_link_name = urdf_parser.link_names_with_meshes(
-            collision=args.collision_meshes
-        )[-1]
-        rich.print(
-            f"End link name not provided. Using the last link with mesh: '{end_link_name}'."
-        )
-
     # instantiate robot
     batch_size = len(joint_states)
-    robot = Robot.from_urdf_parser(
-        urdf_parser=urdf_parser,
-        root_link_name=root_link_name,
-        end_link_name=end_link_name,
+    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,
-        batch_size=batch_size,
+    )
+    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

cli/rr_mono_dr.py

@@ -10,11 +10,17 @@
 import rich.progress
 import torch
 
-from roboreg.differentiable import VirtualCamera
-from roboreg.io import find_files, parse_mono_data
+from roboreg.core import (
+    NVDiffRastRenderer,
+    Robot,
+    RobotScene,
+    TorchKinematics,
+    TorchMeshContainer,
+    VirtualCamera,
+)
+from roboreg.io import find_files, load_robot_data_from_ros_xacro, parse_mono_data
 from roboreg.losses import soft_dice_loss
 from roboreg.util import mask_distance_transform, mask_exponential_decay, overlay_mask
-from roboreg.util.factories import create_robot_scene
 
 
 class REGISTRATION_MODE(Enum):
@@ -181,17 +187,36 @@ def main() -> None:
         )
     }
 
-    # instantiate robot scene
-    scene = create_robot_scene(
-        batch_size=joint_states.shape[0],
+    # instantiate robot
+    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,
-        cameras=camera,
+    )
+    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,
+    )
+
+    # instantiate scene
+    scene = RobotScene(
+        cameras=camera,
+        robot=robot,
+        renderer=NVDiffRastRenderer(device=device),
+    )
 
     # load extrinsics estimate
     extrinsics = torch.tensor(

cli/rr_render.py

@@ -8,10 +8,16 @@
 from rich import progress
 from torch.utils.data import DataLoader
 
-from roboreg.differentiable import VirtualCamera
-from roboreg.io import MonocularDataset
+from roboreg.core import (
+    NVDiffRastRenderer,
+    Robot,
+    RobotScene,
+    TorchKinematics,
+    TorchMeshContainer,
+    VirtualCamera,
+)
+from roboreg.io import MonocularDataset, load_robot_data_from_ros_xacro
 from roboreg.util import overlay_mask
-from roboreg.util.factories import create_robot_scene
 
 
 def args_factory() -> argparse.Namespace:
@@ -119,16 +125,33 @@ def main():
             device=device,
         )
     }
-    scene = create_robot_scene(
-        batch_size=args.batch_size,
+    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,
-        cameras=camera,
-        device=device,
         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,
+    )
+    scene = RobotScene(
+        cameras=camera,
+        robot=robot,
+        renderer=NVDiffRastRenderer(device=device),
+    )
     dataset = MonocularDataset(
         images_path=args.images_path,
         image_pattern=args.image_pattern,

cli/rr_stereo_dr.py

@@ -10,11 +10,17 @@
 import rich.progress
 import torch
 
-from roboreg.differentiable import VirtualCamera
-from roboreg.io import find_files, parse_stereo_data
+from roboreg.core import (
+    NVDiffRastRenderer,
+    Robot,
+    RobotScene,
+    TorchKinematics,
+    TorchMeshContainer,
+    VirtualCamera,
+)
+from roboreg.io import find_files, load_robot_data_from_ros_xacro, parse_stereo_data
 from roboreg.losses import soft_dice_loss
 from roboreg.util import mask_distance_transform, mask_exponential_decay, overlay_mask
-from roboreg.util.factories import create_robot_scene
 
 
 class REGISTRATION_MODE(Enum):
@@ -229,17 +235,36 @@ def main() -> None:
         ),
     }
 
-    # instantiate robot scene
-    scene = create_robot_scene(
-        batch_size=joint_states.shape[0],
+    # instantiate robot
+    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,
-        cameras=cameras,
-        device=device,
         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,
+    )
+
+    # instantiate scene
+    scene = RobotScene(
+        cameras=cameras,
+        robot=robot,
+        renderer=NVDiffRastRenderer(device=device),
+    )
 
     # load extrinscis estimate......
     left_extrinsics = torch.tensor(