successfully returned target points and scene points

This commit is contained in:
0nhc 2024-10-13 01:35:53 -05:00
parent 620010d05a
commit 110726fd13

View File

@ -19,7 +19,7 @@ class ActivePerceptionPolicy(MultiViewPolicy):
super().activate(bbox, view_sphere) super().activate(bbox, view_sphere)
def update(self, img, seg, target_id, x, q): 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(): # if len(self.views) > self.max_views or self.best_grasp_prediction_is_stable():
# self.done = True # self.done = True
# else: # else:
@ -42,6 +42,9 @@ class ActivePerceptionPolicy(MultiViewPolicy):
# self.x_d = nbv # self.x_d = nbv
def depth_image_to_ap_input(self, depth_img, seg_img, target_id): def depth_image_to_ap_input(self, depth_img, seg_img, target_id):
target_points = []
scene_points = []
K = self.intrinsic.K K = self.intrinsic.K
depth_shape = depth_img.shape depth_shape = depth_img.shape
seg_shape = seg_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") print("Depth image shape and segmentation image shape are not the same")
return None return None
# Depth image to PCD # Convert depth image to 3D points
u_indices , v_indices = np.meshgrid(np.arange(img_shape[0]), np.arange(img_shape[1])) 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] x_factors = (u_indices - K[0, 2]) / K[0, 0]
y_factors = (v_indices - K[1, 2]) / K[1, 1] 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 i in range(img_shape[0]):
for j in range(img_shape[1]): for j in range(img_shape[1]):
seg_id = seg_img[i, j] 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)): 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): def best_grasp_prediction_is_stable(self):