From f25281e9302ba5bcc7860c3019a7073876d8b89d Mon Sep 17 00:00:00 2001
From: Balakumar Sundaralingam <balakumars@nvidia.com>
Date: Fri, 23 Feb 2024 10:07:36 -0800
Subject: [PATCH] minor fix to robot segmentation api

---
 src/curobo/wrap/model/robot_segmenter.py | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/src/curobo/wrap/model/robot_segmenter.py b/src/curobo/wrap/model/robot_segmenter.py
index 05b450a..8efa7fb 100644
--- a/src/curobo/wrap/model/robot_segmenter.py
+++ b/src/curobo/wrap/model/robot_segmenter.py
@@ -17,10 +17,8 @@ import torch
 from torch.profiler import record_function
 
 # CuRobo
-from curobo.geom.cv import (
-    get_projection_rays,
-    project_depth_using_rays,
-)
+from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel
+from curobo.geom.cv import get_projection_rays, project_depth_using_rays
 from curobo.geom.types import PointCloud
 from curobo.types.base import TensorDeviceType
 from curobo.types.camera import CameraObservation
@@ -189,6 +187,18 @@ class RobotSegmenter:
 
         return mask, filtered_image
 
+    @property
+    def kinematics(self) -> CudaRobotModel:
+        return self._robot_world.kinematics
+
+    @property
+    def robot_world(self) -> RobotWorld:
+        return self._robot_world
+
+    @property
+    def base_link(self) -> str:
+        return self.kinematics.base_link
+
 
 @torch.jit.script
 def mask_image(