From 82815304bddfe86a69fe204f2ef0506007a88180 Mon Sep 17 00:00:00 2001
From: 0nhc <thefoxfoxfox@outlook.com>
Date: Sun, 13 Oct 2024 01:02:57 -0500
Subject: [PATCH] add new srv

---
 CMakeLists.txt                               |  1 +
 scripts/bt_sim_node.py                       |  6 +++++
 src/active_grasp/active_perception_policy.py | 23 +++++++++++++++++---
 srv/TargetID.srv                             |  2 ++
 4 files changed, 29 insertions(+), 3 deletions(-)
 create mode 100644 srv/TargetID.srv

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 962fb78..f96c09c 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -19,6 +19,7 @@ add_service_files(
   FILES
   Reset.srv
   Seed.srv
+  TargetID.srv
 )
 
 generate_messages(
diff --git a/scripts/bt_sim_node.py b/scripts/bt_sim_node.py
index a8742c9..91e48c0 100755
--- a/scripts/bt_sim_node.py
+++ b/scripts/bt_sim_node.py
@@ -73,6 +73,7 @@ class BtSimNode:
             SwitchController,
             self.switch_controller,
         )
+        rospy.Service("get_target_seg_id", Trigger, self.get_target_seg_id)
 
     def seed(self, req):
         self.sim.seed(req.seed)
@@ -86,6 +87,11 @@ class BtSimNode:
         bbox = self.sim.reset()
         self.activate_plugins()
         return ResetResponse(to_bbox_msg(bbox))
+    
+    def seed(self, req):
+        self.sim.seed(req.seed)
+        rospy.loginfo(f"Seeded the rng with {req.seed}.")
+        return SeedResponse()
 
     def switch_controller(self, req):
         for controller in req.stop_controllers:
diff --git a/src/active_grasp/active_perception_policy.py b/src/active_grasp/active_perception_policy.py
index b86887a..7afe884 100644
--- a/src/active_grasp/active_perception_policy.py
+++ b/src/active_grasp/active_perception_policy.py
@@ -19,7 +19,7 @@ class ActivePerceptionPolicy(MultiViewPolicy):
         super().activate(bbox, view_sphere)
 
     def update(self, img, seg, x, q):
-        self.depth_image_to_ap_input(img)
+        self.depth_image_to_ap_input(img, seg)
         # if len(self.views) > self.max_views or self.best_grasp_prediction_is_stable():
         #     self.done = True
         # else:
@@ -41,8 +41,25 @@ class ActivePerceptionPolicy(MultiViewPolicy):
 
         #     self.x_d = nbv
     
-    def depth_image_to_ap_input(self, depth_img):
-        print(self.intrinsic.K)
+    def depth_image_to_ap_input(self, depth_img, seg_img):
+        K = self.intrinsic.K
+        depth_shape = depth_img.shape
+        seg_shape = seg_img.shape
+        if(depth_shape == seg_shape):
+            img_shape = depth_shape
+        else:
+            print("Depth image shape and segmentation image shape are not the same")
+            return None
+        
+        # Depth image to PCD
+        u_indices , v_indices = np.meshgrid(np.arange(img_shape[0]), np.arange(img_shape[1]))
+        x_factors = (u_indices - K[0, 2]) / K[0, 0]
+        y_factors = (v_indices - K[1, 2]) / K[1, 1]
+
+        for i in range(img_shape[0]):
+            for j in range(img_shape[1]):
+                seg_id = seg_img[i, j]
+
 
     def best_grasp_prediction_is_stable(self):
         if self.best_grasp:
diff --git a/srv/TargetID.srv b/srv/TargetID.srv
new file mode 100644
index 0000000..50e5a0c
--- /dev/null
+++ b/srv/TargetID.srv
@@ -0,0 +1,2 @@
+---
+int16 id