hofee/nbv_sim
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

successfully returned target points and scene points

Browse Source
This commit is contained in:
0nhc 2024-10-13 01:35:53 -05:00
parent 620010d05a
commit 110726fd13
1 changed files with 24 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
src/active_grasp/active_perception_policy.py
View File

@ -19,7 +19,7 @@ class ActivePerceptionPolicy(MultiViewPolicy):
super().activate(bbox, view_sphere)
def update(self, img, seg, target_id, x, q):
self.depth_image_to_ap_input(img, seg, target_id)
target_points, scene_points = self.depth_image_to_ap_input(img, seg, target_id)
# if len(self.views) > self.max_views or self.best_grasp_prediction_is_stable():
# self.done = True
# else:
@ -42,6 +42,9 @@ class ActivePerceptionPolicy(MultiViewPolicy):
# self.x_d = nbv
def depth_image_to_ap_input(self, depth_img, seg_img, target_id):
target_points = []
scene_points = []
K = self.intrinsic.K
depth_shape = depth_img.shape
seg_shape = seg_img.shape
@ -51,17 +54,32 @@ class ActivePerceptionPolicy(MultiViewPolicy):
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]))
# Convert depth image to 3D points
u_indices , v_indices = np.meshgrid(np.arange(img_shape[1]), np.arange(img_shape[0]))
x_factors = (u_indices - K[0, 2]) / K[0, 0]
y_factors = (v_indices - K[1, 2]) / K[1, 1]
z_mat = depth_img
x_mat = x_factors * z_mat
y_mat = y_factors * z_mat
for i in range(img_shape[0]):
for j in range(img_shape[1]):
seg_id = seg_img[i, j]
# print(seg_id)
x = x_mat[i][j]
y = y_mat[i][j]
z = z_mat[i][j]
if(int(seg_id) == int(target_id)):
print("woahhhhhhhhhhhhhhhhh!")
# This pixel belongs to the target object to be grasped
target_points.append([x,y,z])
else:
# This pixel belongs to the scene
scene_points.append([x,y,z])
target_points = np.asarray(target_points)
target_points = target_points.reshape(1, target_points.shape[0], 3)
scene_points = np.asarray(scene_points)
scene_points = scene_points.reshape(1, scene_points.shape[0], 3)
return target_points, scene_points
def best_grasp_prediction_is_stable(self):