def start_rendering(self): self._load_file_ids() for object_id in self.all_objects: self._load_data(object_id) for i, stable_pose in enumerate(self.stable_poses): try: candidate_grasp_info = self.candidate_grasps_dict[ stable_pose.id] except KeyError: continue if not candidate_grasp_info: Warning("Candidate grasp info of object id %s empty" % object_id) Warning("Continue.") continue T_obj_stp = stable_pose.T_obj_table.as_frames('obj', 'stp') T_obj_stp = self.object_mesh.get_T_surface_obj(T_obj_stp) T_table_obj = RigidTransform(from_frame='table', to_frame='obj') scene_objs = { 'table': SceneObject(self.table_mesh, T_table_obj) } urv = UniformPlanarWorksurfaceImageRandomVariable( self.object_mesh, [RenderMode.DEPTH_SCENE, RenderMode.SEGMASK], 'camera', self.config['env_rv_params'], scene_objs=scene_objs, stable_pose=stable_pose) render_sample = urv.rvs(size=self.random_positions) # for render_sample in render_samples: binary_im = render_sample.renders[RenderMode.SEGMASK].image depth_im = render_sample.renders[ RenderMode.DEPTH_SCENE].image.crop(300, 300) orig_im = Image.fromarray(self._scale_image(depth_im.data)) if self.show_images: orig_im.show() orig_im.convert('RGB').save(self.output_dir + '/images/' + object_id + '_elev_' + str(self.elev) + '_original.png') print("Saved original") T_stp_camera = render_sample.camera.object_to_camera_pose shifted_camera_intr = render_sample.camera.camera_intr.crop( 300, 300, 240, 320) depth_points = self._reproject_to_3D(depth_im, shifted_camera_intr) transformed_points, T_camera = self._transformation( depth_points) camera_dir = np.dot(T_camera.rotation, np.array([0.0, 0.0, -1.0])) pcd = o3d.geometry.PointCloud() # print(camera_dir) pcd.points = o3d.utility.Vector3dVector(transformed_points.T) # TODO check normals!! # pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30)) # pcd.normals = o3d.utility.Vector3dVector(-np.asarray(pcd.normals)) normals = np.repeat([camera_dir], len(transformed_points.T), axis=0) pcd.normals = o3d.utility.Vector3dVector(normals) # cs_points = [[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]] # cs_lines = [[0, 1], [0, 2], [0, 3]] # colors = [[1, 0, 0], [0, 1, 0], [0, 0, 1]] # cs = o3d.geometry.LineSet(points=o3d.utility.Vector3dVector(cs_points), # lines=o3d.utility.Vector2iVector(cs_lines)) # cs.colors = o3d.utility.Vector3dVector(colors) # o3d.visualization.draw_geometries([pcd]) depth = self._o3d_meshing(pcd) # projected_depth_im,new_camera_intr,table_height = self._projection(new_points,shifted_camera_intr) new_camera_intr = shifted_camera_intr new_camera_intr.cx = 150 new_camera_intr.cy = 150 projected_depth_im = np.asarray(depth) projected_depth_im[projected_depth_im == 0.0] = -1.0 table_height = np.median( projected_depth_im[projected_depth_im != -1.0].flatten()) print("Minimum depth:", min(projected_depth_im.flatten())) print("Maximum depth:", max(projected_depth_im.flatten())) im = Image.fromarray(self._scale_image(projected_depth_im)) projected_depth_im = DepthImage(projected_depth_im, frame='new_camera') cx = projected_depth_im.center[1] cy = projected_depth_im.center[0] # Grasp conversion T_obj_old_camera = T_stp_camera * T_obj_stp.as_frames( 'obj', T_stp_camera.from_frame) T_obj_camera = T_camera.dot(T_obj_old_camera) for grasp_info in candidate_grasp_info: grasp = grasp_info.grasp collision_free = grasp_info.collision_free grasp_2d = grasp.project_camera(T_obj_camera, new_camera_intr) dx = cx - grasp_2d.center.x dy = cy - grasp_2d.center.y translation = np.array([dy, dx]) # Project 3D old_camera_cs contact points into new camera cs contact_points = np.append(grasp_info.contact_point1, 1).T new_cam = np.dot(T_obj_camera.matrix, contact_points) c1 = new_camera_intr.project( Point(new_cam[0:3], frame=new_camera_intr.frame)) contact_points = np.append(grasp_info.contact_point2, 1).T new_cam = np.dot(T_obj_camera.matrix, contact_points) c2 = new_camera_intr.project( Point(new_cam[0:3], frame=new_camera_intr.frame)) # Check if there are occlusions at contact points if projected_depth_im.data[ c1.x, c1.y] == -1.0 or projected_depth_im.data[ c2.x, c2.y] == -1.0: print("Contact point at occlusion") contact_occlusion = True else: contact_occlusion = False # Mark contact points in image im = im.convert('RGB') if False: im_draw = ImageDraw.Draw(im) im_draw.line([(c1[0], c1[1] - 10), (c1[0], c1[1] + 10)], fill=(255, 0, 0, 255)) im_draw.line([(c1[0] - 10, c1[1]), (c1[0] + 10, c1[1])], fill=(255, 0, 0, 255)) im_draw.line([(c2[0], c2[1] - 10), (c2[0], c2[1] + 10)], fill=(255, 0, 0, 255)) im_draw.line([(c2[0] - 10, c2[1]), (c2[0] + 10, c2[1])], fill=(255, 0, 0, 255)) if self.show_images: im.show() im.save(self.output_dir + '/images/' + object_id + '_elev_' + str(self.elev) + '_reprojected.png') # Transform and crop image depth_im_tf = projected_depth_im.transform( translation, grasp_2d.angle) depth_im_tf = depth_im_tf.crop(96, 96) # Apply transformation to contact points trans_map = np.array([[1, 0, dx], [0, 1, dy]]) rot_map = cv2.getRotationMatrix2D( (cx, cy), np.rad2deg(grasp_2d.angle), 1) trans_map_aff = np.r_[trans_map, [[0, 0, 1]]] rot_map_aff = np.r_[rot_map, [[0, 0, 1]]] full_map = rot_map_aff.dot(trans_map_aff) # print("Full map",full_map) c1_rotated = (np.dot(full_map, np.r_[c1.vector, [1]]) - np.array([150 - 48, 150 - 48, 0])) / 3 c2_rotated = (np.dot(full_map, np.r_[c2.vector, [1]]) - np.array([150 - 48, 150 - 48, 0])) / 3 grasp_line = depth_im_tf.data[48] occlusions = len(np.where(np.squeeze(grasp_line) == -1)[0]) # Set occlusions to table height for resizing image depth_im_tf.data[depth_im_tf.data == -1.0] = table_height depth_image = Image.fromarray(np.asarray(depth_im_tf.data))\ .resize((32, 32), resample=Image.BILINEAR) depth_im_tf_table = np.asarray(depth_image).reshape( 32, 32, 1) # depth_im_tf_table = depth_im_tf.resize((32, 32,), interp='bilinear') im = Image.fromarray( self._scale_image(depth_im_tf_table.reshape( 32, 32))).convert('RGB') draw = ImageDraw.Draw(im) draw.line([(c1_rotated[0], c1_rotated[1] - 3), (c1_rotated[0], c1_rotated[1] + 3)], fill=(255, 0, 0, 255)) draw.line([(c1_rotated[0] - 3, c1_rotated[1]), (c1_rotated[0] + 3, c1_rotated[1])], fill=(255, 0, 0, 255)) draw.line([(c2_rotated[0], c2_rotated[1] - 3), (c2_rotated[0], c2_rotated[1] + 3)], fill=(255, 0, 0, 255)) draw.line([(c2_rotated[0] - 3, c2_rotated[1]), (c2_rotated[0] + 3, c2_rotated[1])], fill=(255, 0, 0, 255)) if self.show_images: im.show() im.save(self.output_dir + '/images/' + object_id + '_elev_' + str(self.elev) + '_transformed.png') hand_pose = np.r_[grasp_2d.center.y, grasp_2d.center.x, grasp_2d.depth, grasp_2d.angle, grasp_2d.center.y - new_camera_intr.cy, grasp_2d.center.x - new_camera_intr.cx, grasp_2d.width_px / 3] self.tensor_datapoint[ 'depth_ims_tf_table'] = depth_im_tf_table self.tensor_datapoint['hand_poses'] = hand_pose self.tensor_datapoint['obj_labels'] = self.cur_obj_label self.tensor_datapoint['collision_free'] = collision_free self.tensor_datapoint['pose_labels'] = self.cur_pose_label self.tensor_datapoint[ 'image_labels'] = self.cur_image_label self.tensor_datapoint['files'] = [self.tensor, self.array] self.tensor_datapoint['occlusions'] = occlusions self.tensor_datapoint[ 'contact_occlusion'] = contact_occlusion for metric_name, metric_val in self.grasp_metrics[str( grasp.id)].iteritems(): coll_free_metric = (1 * collision_free) * metric_val self.tensor_datapoint[metric_name] = coll_free_metric self.tensor_dataset.add(self.tensor_datapoint) print("Saved dataset point") self.cur_image_label += 1 self.cur_pose_label += 1 gc.collect() self.cur_obj_label += 1 self.tensor_dataset.flush()
cy = depth_im.center[0] for g, grasp in enumerate(grasps): current_index = start_ind + g if (current_index < 3057000): continue grasp_x = grasp[1] grasp_y = grasp[0] grasp_angle = grasp[3] # center images on the grasp, rotate to image x axis dx = cx - grasp_x dy = cy - grasp_y translation = np.array([dy, dx]) depth_im_tf_table = depth_im.transform(translation, grasp_angle) # crop to image size depth_im_tf_table = depth_im_tf_table.crop( im_crop_height, im_crop_width) # resize to image size depth_im_tf_table = depth_im_tf_table.resize( (im_final_height, im_final_width)) tensor_datapoint[ 'depth_ims_tf_table_96'] = depth_im_tf_table.raw_data tensor_dataset.add(tensor_datapoint) gc.collect()