Merge pull request #106 from lbr-stack/feature/roboreg-103/mesh_origin_refactor

LGTM

Author: mhubii

README.md

@@ -183,7 +183,7 @@ rr-cam-swarm \
     --xacro-path urdf/med7/med7.xacro \
     --root-link-name lbr_link_0 \
     --end-link-name lbr_link_7 \
-    --target-reduction 0.95 \
+    --target-reduction 0.8 \
     --scale 0.1 \
     --n-samples 1 \
     --camera-info-file test/assets/lbr_med7/zed2i/left_camera_info.yaml \

cli/rr_cam_swarm.py

@@ -6,7 +6,7 @@
 import torch
 
 from roboreg import differentiable as rrd
-from roboreg.io import find_files, parse_camera_info, parse_mono_data
+from roboreg.io import URDFParser, find_files, parse_camera_info, parse_mono_data
 from roboreg.losses import soft_dice_loss
 from roboreg.optim import LinearParticleSwarm, ParticleSwarmOptimizer
 from roboreg.util import (
@@ -249,7 +249,6 @@ def main() -> None:
     mask_files = np.array(mask_files)[random_indices].tolist()
     joint_states_files = np.array(joint_states_files)[random_indices].tolist()
     images, joint_states, masks = parse_mono_data(
-        path=args.path,
         image_files=image_files,
         mask_files=mask_files,
         joint_states_files=joint_states_files,
@@ -302,9 +301,10 @@ def main() -> None:
         device=device,
     )
 
-    urdf_parser = rrd.URDFParser()
-    urdf_parser.from_ros_xacro(ros_package=args.ros_package, xacro_path=args.xacro_path)
-    robot = rrd.Robot(
+    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,
         root_link_name=args.root_link_name,
         end_link_name=args.end_link_name,

cli/rr_hydra.py

@@ -155,16 +155,16 @@ def main():
     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(
-        path=args.path,
         joint_states_files=joint_states_files,
         mask_files=mask_files,
         depth_files=depth_files,
     )
     height, width, intrinsics = parse_camera_info(args.camera_info_file)
 
     # instantiate kinematics
-    urdf_parser = URDFParser()
-    urdf_parser.from_ros_xacro(ros_package=args.ros_package, xacro_path=args.xacro_path)
+    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 == "":
@@ -184,7 +184,7 @@ def main():
 
     # instantiate robot
     batch_size = len(joint_states)
-    robot = Robot(
+    robot = Robot.from_urdf_parser(
         urdf_parser=urdf_parser,
         root_link_name=root_link_name,
         end_link_name=end_link_name,
@@ -226,7 +226,9 @@ def main():
     mesh_normals = []
     for i in range(batch_size):
         mesh_normals.append(
-            compute_vertex_normals(vertices=mesh_vertices[i], faces=robot.faces)
+            compute_vertex_normals(
+                vertices=mesh_vertices[i], faces=robot.mesh_container.faces
+            )
         )
 
     # clean observed vertices and turn into tensor

cli/rr_mono_dr.py

@@ -153,7 +153,6 @@ def main() -> None:
     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(
-        path=args.path,
         image_files=image_files,
         joint_states_files=joint_states_files,
         mask_files=mask_files,

cli/rr_render.py

@@ -1,6 +1,6 @@
 import argparse
 import os
-import pathlib
+from pathlib import Path
 
 import cv2
 import numpy as np
@@ -142,7 +142,7 @@ def main():
         num_workers=args.num_workers,
     )
 
-    output_path = pathlib.Path(args.output_path)
+    output_path = Path(args.output_path)
     if not output_path.exists():
         output_path.mkdir(parents=True)
 
@@ -162,13 +162,12 @@ def main():
         images = images.numpy()
         renders = (renders * 255.0).squeeze(-1).cpu().numpy().astype(np.uint8)
         for render, image, image_file in zip(renders, images, image_files):
-            image_stem = pathlib.Path(image_file).stem
-            image_suffix = pathlib.Path(image_file).suffix
+            image_file = Path(image_file)
+            output_file = (
+                output_path / f"overlay_render_{image_file.stem + image_file.suffix}"
+            )
             cv2.imwrite(
-                os.path.join(
-                    str(output_path.absolute()),
-                    f"overlay_render_{image_stem + image_suffix}",
-                ),
+                output_file,
                 overlay_mask(image, render, args.color, scale=1.0),
             )
 

cli/rr_sam2.py

@@ -1,6 +1,4 @@
 import argparse
-import os
-import pathlib
 
 import cv2
 import numpy as np
@@ -48,8 +46,7 @@ def args_factory() -> argparse.Namespace:
 
 def main():
     args = args_factory()
-    path = pathlib.Path(args.path)
-    image_names = find_files(path.absolute(), args.pattern)
+    image_files = find_files(args.path, args.pattern)
 
     # detect
     detector = OpenCVDetector(
@@ -60,23 +57,21 @@ def main():
     # segment
     segmentor = Sam2Segmentor(model_id=args.model_id, device=args.device)
 
-    for image_name in progress.track(image_names, description="Generating masks..."):
-        image_stem = pathlib.Path(image_name).stem
-        image_suffix = pathlib.Path(image_name).suffix
-        img = cv2.imread(os.path.join(path.absolute(), image_name))
+    for image_file in progress.track(image_files, description="Generating masks..."):
+        img = cv2.imread(image_file)
         annotations = False
         if args.pre_annotated:
             try:
                 samples, labels = detector.read(
-                    path=os.path.join(path.absolute(), f"{image_stem}_samples.csv")
+                    path=image_file.parent / f"{image_file.stem}_samples.csv"
                 )
                 annotations = True
             except FileNotFoundError:
                 pass
         if not annotations:
             samples, labels = detector.detect(img)
             detector.write(
-                path=os.path.join(path.absolute(), f"{image_stem}_samples.csv"),
+                path=image_file.parent / f"{image_file.stem}_samples.csv",
                 samples=samples,
                 labels=labels,
             )
@@ -86,15 +81,9 @@ def main():
         overlay = overlay_mask(img, mask, mode="g", scale=1.0)
 
         # write probability and mask
-        probability_path = os.path.join(
-            path.absolute(), f"probability_sam2_{image_stem + image_suffix}"
-        )
-        mask_path = os.path.join(
-            path.absolute(), f"mask_sam2_{image_stem + image_suffix}"
-        )
-        overlay_path = os.path.join(
-            path.absolute(), f"overlay_sam2_{image_stem + image_suffix}"
-        )
+        probability_path = image_file.parent / f"probability_sam2_{image_file.name}"
+        mask_path = image_file.parent / f"mask_sam2_{image_file.name}"
+        overlay_path = image_file.parent / f"overlay_sam2_{image_file.name}"
         cv2.imwrite(probability_path, (probability * 255.0).astype(np.uint8))
         cv2.imwrite(mask_path, mask)
         cv2.imwrite(overlay_path, overlay)

cli/rr_stereo_dr.py

@@ -186,7 +186,6 @@ def main() -> None:
     right_mask_files = find_files(args.path, args.right_mask_pattern)
     left_images, right_images, joint_states, left_masks, right_masks = (
         parse_stereo_data(
-            path=args.path,
             left_image_files=left_image_files,
             right_image_files=right_image_files,
             joint_states_files=joint_states_files,

roboreg/differentiable/kinematics.py

@@ -3,43 +3,31 @@
 import pytorch_kinematics as pk
 import torch
 
-from roboreg.io import URDFParser
-
 
 class TorchKinematics:
     __slots__ = [
         "_root_link_name",
         "_end_link_name",
         "_chain",
-        "_mesh_origins_lookup",
         "_device",
     ]
 
     def __init__(
         self,
-        urdf_parser: URDFParser,
+        urdf: str,
         root_link_name: str,
         end_link_name: str,
-        device: torch.device = "cuda",
+        device: torch.device = torch.device("cuda"),
     ) -> None:
         self._root_link_name = root_link_name
         self._end_link_name = end_link_name
         self._chain = self._build_serial_chain_from_urdf(
-            urdf_parser.urdf,
+            urdf=urdf,
             root_link_name=self._root_link_name,
             end_link_name=self._end_link_name,
         )
-
-        self._mesh_origins_lookup = urdf_parser.mesh_origins(
-            root_link_name=root_link_name, end_link_name=end_link_name
-        )
-        self._mesh_origins_lookup = {
-            key: torch.from_numpy(value).to(device=device, dtype=torch.float32)
-            for key, value in self._mesh_origins_lookup.items()
-        }
-
-        # default move to device
-        self.to(device=device)
+        self._device = device
+        self.to(device=self._device)
 
     def _build_serial_chain_from_urdf(
         self, urdf: str, root_link_name: str, end_link_name: str
@@ -50,18 +38,11 @@ def _build_serial_chain_from_urdf(
 
     def to(self, device: torch.device) -> None:
         self._chain.to(device=device)
-        for link_name in self._mesh_origins_lookup:
-            self._mesh_origins_lookup[link_name] = self._mesh_origins_lookup[
-                link_name
-            ].to(device=device)
         self._device = device
 
-    def mesh_forward_kinematics(self, q: torch.Tensor) -> Dict[str, torch.Tensor]:
-        r"""Computes forward kinematics and returns corresponding homogeneous transformations.
-        Corrects for mesh offsets. Meshes that are tranformed by the returned transformation appear physically correct.
-        """
+    def forward_kinematics(self, q: torch.Tensor) -> Dict[str, torch.Tensor]:
         ht_lookup = {
-            key: value.get_matrix() @ self._mesh_origins_lookup[key]
+            key: value.get_matrix()
             for key, value in self._chain.forward_kinematics(q, end_only=False).items()
         }
         return ht_lookup

roboreg/differentiable/rendering.py

@@ -13,7 +13,7 @@ class NVDiffRastRenderer:
 
     __slots__ = ["_device", "_ctx"]
 
-    def __init__(self, device: torch.device = "cuda") -> None:
+    def __init__(self, device: torch.device = torch.device("cuda")) -> None:
         super().__init__()
         if not torch.cuda.is_available():
             raise ValueError("CUDA is not available.")