def sample_grasp(self, env, mesh_obj, vis=True):
# Setup sensor.
#camera_intr = CameraIntrinsics.load(camera_intr_filename)
# Read images.
# get depth image from camera
inpaint_rescale_factor = 1.0
segmask_filename = None
_, center, extents, obj_xquat, bbox_xpos = self.get_bbox_properties(
env, mesh_obj)
camera_pos = copy.deepcopy(center)
camera_pos[2] += 0.5 * extents[2] + 0.2
env.set_camera(camera_pos,
np.array([0.7071068, 0, 0, -0.7071068]),
camera_name=f"ext_camera_0")
rgb, depth = env.render_from_camera(int(self.config.camera_img_height),
int(self.config.camera_img_width),
camera_name=f"ext_camera_0")
# enforce zoom out
from scipy.interpolate import interp2d
center_x = self.config.camera_img_height / 2 + 1
center_y = self.config.camera_img_width / 2 + 1
img_height = self.config.camera_img_height
img_width = self.config.camera_img_width
xdense = np.linspace(0, img_height - 1, img_height)
ydense = np.linspace(0, img_width - 1, img_width)
xintr = (xdense - center_x) * (1.0 / self.rescale_factor) + center_x
yintr = (ydense - center_y) * (1.0 / self.rescale_factor) + center_y
xintr[xintr < 0] = 0
xintr[xintr > (img_height - 1)] = img_height - 1
yintr[yintr < 0] = 0
yintr[yintr > (img_width - 1)] = img_width - 1
fr = interp2d(xdense, ydense, rgb[:, :, 0], kind="linear")
rgb_r_new = fr(xintr, yintr)
fg = interp2d(xdense, ydense, rgb[:, :, 1], kind="linear")
rgb_g_new = fg(xintr, yintr)
fb = interp2d(xdense, ydense, rgb[:, :, 2], kind="linear")
rgb_b_new = fb(xintr, yintr)
rgb_new = np.stack([rgb_r_new, rgb_g_new, rgb_b_new], axis=2)
fd = interp2d(xdense, ydense, depth, kind="linear")
depth_new = fd(xintr, yintr)
#from skimage.transform import resize
#rgb22, depth2 = env.render_from_camera(int(self.config.camera_img_height) , int(self.config.camera_img_width), camera_name=f"ext_camera_0")
#import ipdb; ipdb.set_trace()
# visualize the interpolation
#import imageio
#imageio.imwrite(f"tmp/rgb_{self.iter_id}.png", rgb)
#imageio.imwrite(f"tmp/rgb2_{self.iter_id}.png", rgb_new)
#imageio.imwrite(f"tmp/depth_{self.iter_id}.png", depth)
#imageio.imwrite(f"tmp/depth2_{self.iter_id}.png", depth_new)
#import ipdb; ipdb.set_trace()
rgb = rgb_new
depth = depth_new
depth = depth * self.rescale_factor
# rgb: 128 x 128 x 1
# depth: 128 x 128 x 1
scaled_camera_fov_y = self.config.camera_fov_y
aspect = 1
scaled_fovx = 2 * np.arctan(
np.tan(np.deg2rad(scaled_camera_fov_y) * 0.5) * aspect)
scaled_fovx = np.rad2deg(scaled_fovx)
scaled_fovy = scaled_camera_fov_y
cx = self.config.camera_img_width * 0.5
cy = self.config.camera_img_height * 0.5
scaled_fx = cx / np.tan(np.deg2rad(
scaled_fovx / 2.)) * (self.rescale_factor)
scaled_fy = cy / np.tan(np.deg2rad(
scaled_fovy / 2.)) * (self.rescale_factor)
camera_intr = CameraIntrinsics(frame='phoxi',
fx=scaled_fx,
fy=scaled_fy,
cx=self.config.camera_img_width * 0.5,
cy=self.config.camera_img_height * 0.5,
height=self.config.camera_img_height,
width=self.config.camera_img_width)
depth_im = DepthImage(depth, frame=camera_intr.frame)
color_im = ColorImage(np.zeros([depth_im.height, depth_im.width,
3]).astype(np.uint8),
frame=camera_intr.frame)
# Optionally read a segmask.
valid_px_mask = depth_im.invalid_pixel_mask().inverse()
segmask = valid_px_mask
# Inpaint.
depth_im = depth_im.inpaint(rescale_factor=inpaint_rescale_factor)
# Create state.
rgbd_im = RgbdImage.from_color_and_depth(color_im, depth_im)
state = RgbdImageState(rgbd_im, camera_intr, segmask=segmask)
# Query policy.
policy_start = time.time()
action = self.policy(state)
print("Planning took %.3f sec" % (time.time() - policy_start))
# numpy array with 2 values
grasp_center = action.grasp.center._data[:, 0] #(width, depth)
grasp_depth = action.grasp.depth * (1 / self.rescale_factor)
grasp_angle = action.grasp.angle #np.pi*0.3
if self.config.data_collection_mode:
self.current_grasp = action.grasp
depth_im = state.rgbd_im.depth
scale = 1.0
depth_im_scaled = depth_im.resize(scale)
translation = scale * np.array([
depth_im.center[0] - grasp_center[1],
depth_im.center[1] - grasp_center[0]
])
im_tf = depth_im_scaled
im_tf = depth_im_scaled.transform(translation, grasp_angle)
im_tf = im_tf.crop(self.gqcnn_image_size, self.gqcnn_image_size)
# get the patch
self.current_patch = im_tf.raw_data
XYZ_origin, gripper_quat = self.compute_grasp_pts_from_grasp_sample(
grasp_center, grasp_depth, grasp_angle, env)
return XYZ_origin[:, 0], gripper_quat
# Vis final grasp.
if vis:
from visualization import Visualizer2D as vis
vis.figure(size=(10, 10))
vis.imshow(rgbd_im.depth,
vmin=self.policy_config["vis"]["vmin"],
vmax=self.policy_config["vis"]["vmax"])
vis.grasp(action.grasp,
scale=2.5,
show_center=False,
show_axis=True)
vis.title("Planned grasp at depth {0:.3f}m with Q={1:.3f}".format(
action.grasp.depth, action.q_value))
vis.show(f"tmp/grasp2_{mesh_obj.name}_{self.iter_id}.png")
vis.figure(size=(10, 10))
vis.imshow(im_tf,
vmin=self.policy_config["vis"]["vmin"],
vmax=self.policy_config["vis"]["vmax"])
vis.show(f"tmp/cropd_{mesh_obj.name}_{self.iter_id}.png")
import ipdb
ipdb.set_trace()
return XYZ_origin[:, 0], gripper_quat
import ipdb
ipdb.set_trace()
for ind in range(num_image_tensors):
print("File: " + str(ind))
image_arr = np.load(image_tensors[ind])['arr_0']
binary_arr = np.load(binary_image_tensors[ind])['arr_0']
start_ind_arr = np.load(start_ind_tensors[ind])['arr_0']
end_ind_arr = np.load(end_ind_tensors[ind])['arr_0']
datapoints_in_file = min(len(start_ind_arr), DATAPOINTS_PER_FILE)
for i in range(datapoints_in_file):
start_ind = start_ind_arr[i]
end_ind = end_ind_arr[i]
depth_im = DepthImage(image_arr[i, ...])
depth_resized = depth_im.resize((200, 200))
tensor_datapoint['depth_images'] = depth_resized.raw_data
tensor_dataset.add(tensor_datapoint)
# logging.info('Image creation took %.3f sec' %(time.time() - fc_im_start))
# plt.figure(figsize=(10,6))
# plt.subplot(1,2,1)
# quality_image = quality_image[50:150,50:150,:]
# plt.imshow(quality_image[:,:,0])
# plt.axis('off')
# plt.title('QUALITY', fontsize=FONT_SIZE)
# plt.subplot(1,2,2)