def get_object_pose_in_world(object_pose, camera_pose, fat_world_pose=None, type='quat'): # Returns in cm object_pose_matrix = np.zeros((4,4)) object_pose_matrix[:3,:3] = RT_transform.quat2mat(get_wxyz_quaternion(object_pose['quaternion_xyzw'])) object_pose_matrix[:,3] = object_pose['location'] + [1] # camera_pose_matrix = np.zeros((4,4)) # camera_pose_matrix[:3, :3] = RT_transform.quat2mat(get_wxyz_quaternion(camera_pose['quaternion_xyzw_worldframe'])) # camera_pose_matrix[:, 3] = camera_pose['location_worldframe'] + [1] camera_pose_matrix = get_camera_pose_in_world(camera_pose, fat_world_pose, type='rot', cam_to_body=None) object_pose_world = np.matmul(camera_pose_matrix, object_pose_matrix) # object_pose_world = np.matmul(np.linalg.inv(camera_pose_matrix), object_pose_matrix) # scale = np.array([[0.01,0,0,0],[0,0.01,0,0],[0,0,0.01,0],[0,0,0,1]]) # object_pose_world = np.matmul(scale, object_pose_world) if fat_world_pose is not None: fat_world_matrix = np.zeros((4,4)) fat_world_matrix[:3,:3] = RT_transform.quat2mat(get_wxyz_quaternion(fat_world_pose['quaternion_xyzw'])) fat_world_matrix[:,3] = fat_world_pose['location'] + [1] object_pose_world = np.matmul(fat_world_matrix, object_pose_world) # print(object_pose_world) if type == 'quat': quat = RT_transform.mat2quat(object_pose_world[:3, :3]).tolist() return object_pose_world[:3,3].tolist(), get_xyzw_quaternion(quat)
def get_camera_pose_in_world(camera_pose, fat_world_pose=None, type='quat', cam_to_body=None): # print(camera_pose) # this matrix gives world to camera transform camera_pose_matrix = np.zeros((4,4)) # camera_rotation = np.linalg.inv(RT_transform.quat2mat(get_wxyz_quaternion(camera_pose['quaternion_xyzw_worldframe']))) camera_rotation = RT_transform.quat2mat(get_wxyz_quaternion(camera_pose['quaternion_xyzw_worldframe'])) camera_pose_matrix[:3, :3] = camera_rotation camera_location = [i for i in camera_pose['location_worldframe']] # camera_pose_matrix[:, 3] = np.matmul(-1 * camera_rotation, camera_location).tolist() + [1] camera_pose_matrix[:, 3] = camera_location + [1] # # make it camera to world # camera_pose_matrix = np.linalg.inv(camera_pose_matrix) if cam_to_body is not None: camera_pose_matrix = np.matmul(camera_pose_matrix, np.linalg.inv(cam_to_body)) if fat_world_pose is not None: fat_world_matrix = np.zeros((4,4)) fat_world_matrix[:3,:3] = RT_transform.quat2mat(get_wxyz_quaternion(fat_world_pose['quaternion_xyzw'])) fat_world_matrix[:,3] = fat_world_pose['location'] + [1] camera_pose_world = np.matmul(fat_world_matrix, camera_pose_matrix) else: camera_pose_world = camera_pose_matrix # make it camera to world # camera_pose_world = np.linalg.inv(camera_pose_world) if type == 'quat': quat = RT_transform.mat2quat(camera_pose_world[:3, :3]).tolist() return camera_pose_world[:3,3], get_xyzw_quaternion(quat) elif type == 'rot': return camera_pose_world
def render(self, r, t, light_position, light_intensity, class_name, brightness_k=0, r_type='quat'): ''' :param r: :param t: :param light_position: :param light_intensity: :param brightness_k: choose which brightness in __init__ :param r_type: :return: ''' if r_type == 'quat': R = RT_transform.quat2mat(r) elif r_type == 'mat': R = r self.class_name = class_name self.brightness_k = brightness_k self.render_kernels[ self.class_name][brightness_k]['u_view'] = self._get_view_mtx( R, t) self.render_kernels[ self.class_name][brightness_k]['u_light_position'] = light_position self.render_kernels[self.class_name][brightness_k]['u_normal'] = \ np.array(np.matrix(np.dot( self.render_kernels[self.class_name][brightness_k]['u_view'].reshape(4,4), self.render_kernels[self.class_name][brightness_k]['u_model'].reshape(4,4))).I.T) self.render_kernels[self.class_name][brightness_k][ 'u_light_intensity'] = light_intensity app.run(framecount=0, framerate=0) rgb_buffer = np.zeros((self.height, self.width, 4), dtype=np.float32) gl.glReadPixels(0, 0, self.width, self.height, gl.GL_RGBA, gl.GL_FLOAT, rgb_buffer) rgb_gl = np.copy(rgb_buffer) rgb_gl.shape = 480, 640, 4 rgb_gl = rgb_gl[::-1, :] rgb_gl = np.round(rgb_gl[:, :, :3] * 255).astype( np.uint8) # Convert to [0, 255] rgb_gl = rgb_gl[:, :, [2, 1, 0]] depth_buffer = np.zeros((self.height, self.width), dtype=np.float32) gl.glReadPixels(0, 0, self.width, self.height, gl.GL_DEPTH_COMPONENT, gl.GL_FLOAT, depth_buffer) depth_gl = np.copy(depth_buffer) depth_gl.shape = 480, 640 depth_gl = depth_gl[::-1, :] depth_bg = depth_gl == 1 depth_gl = 2 * self.zFar * self.zNear / (self.zFar + self.zNear - (self.zFar - self.zNear) * (2 * depth_gl - 1)) depth_gl[depth_bg] = 0 return rgb_gl, depth_gl
def gen_poses(): pz = np.array([0, 0, 1]) pose_dict, quat_stat, trans_stat, new_points = stat_poses() observed_prefix_list = ['{:06d}'.format(i + 1) for i in range(NUM_IMAGES)] sel_classes = classes observed_pose_dict = { cls_name: np.zeros((NUM_IMAGES, 7)) for cls_name in sel_classes } for cls_i, cls_name in enumerate(sel_classes): # uncomment here to only generate data for ape # if cls_name not in ['ape']: # continue src_quat_mean = quat_stat[cls_name]['quat_mean'] src_quat_std = quat_stat[cls_name]['quat_std'] src_trans_mean = trans_stat[cls_name]['trans_mean'] src_trans_std = trans_stat[cls_name]['trans_std'] deg_max = new_points[cls_name]['angle_max'] for i in tqdm(range(NUM_IMAGES)): observed_prefix = observed_prefix_list[i] # randomly generate a pose tgt_quat = np.random.normal(0, 1, 4) tgt_quat = tgt_quat / np.linalg.norm(tgt_quat) if tgt_quat[0] < 0: tgt_quat *= -1 tgt_trans = np.random.normal(src_trans_mean, src_trans_std) tgt_rot_m = se3.quat2mat(tgt_quat) new_pz = np.dot(tgt_rot_m, pz.reshape((-1, 1))).reshape((3, )) pz_mean = new_points[cls_name]['pz_mean'] deg = angle(new_pz, pz_mean) # r_dist, t_dist = calc_rt_dist_q(tgt_quat, src_quat, tgt_trans, src_trans) transform = np.matmul(K, tgt_trans.reshape(3, 1)) center_x = float(transform[0] / transform[2]) center_y = float(transform[1] / transform[2]) count = 0 while (deg > deg_max or not (48 < center_x < (640 - 48) and 48 < center_y < (480 - 48))): # randomly generate a pose tgt_quat = np.random.normal(0, 1, 4) tgt_quat = tgt_quat / np.linalg.norm(tgt_quat) if tgt_quat[0] < 0: tgt_quat *= -1 tgt_trans = np.random.normal(src_trans_mean, src_trans_std) tgt_rot_m = se3.quat2mat(tgt_quat) new_pz = np.dot(tgt_rot_m, pz.reshape((-1, 1))).reshape((3, )) pz_mean = new_points[cls_name]['pz_mean'] deg = angle(new_pz, pz_mean) transform = np.matmul(K, tgt_trans.reshape(3, 1)) center_x = float(transform[0] / transform[2]) center_y = float(transform[1] / transform[2]) count += 1 if count % 100 == 0: print( observed_prefix, cls_name, count, "deg < deg_max={}: {}, 48 < center_x < (640-48): {}, 48 < center_y < (480-48): {}" .format(deg_max, deg <= deg_max, 48 < center_x < (640 - 48), 48 < center_y < (480 - 48))) tgt_pose_q = np.zeros((7, ), dtype='float32') tgt_pose_q[:4] = tgt_quat tgt_pose_q[4:] = tgt_trans observed_pose_dict[cls_name][i, :] = tgt_pose_q # write pose poses_file = os.path.join(pose_dir, 'LM6d_ds_train_observed_pose_all.pkl') with open(poses_file, 'wb') as f: cPickle.dump(observed_pose_dict, f, 2)
def main(): gen_images = True for class_idx, class_name in idx2class.items(): train_pair = [] val_pair = [] print("start ", class_idx, class_name) if class_name in ["__back_ground__"]: continue # uncomment here to only generate data for ape # if class_name not in ['ape']: # continue if gen_images: # init renderer model_dir = os.path.join(LM6d_root, "models", class_name) render_machine = Render_Py(model_dir, K, width, height, ZNEAR, ZFAR) for set_type in ["all"]: with open( os.path.join(observed_set_root, "{}_{}.txt".format(class_name, "all")), "r", ) as f: all_observed_list = [x.strip() for x in f.readlines()] # with open( # os.path.join(observed_set_root, '{}_{}.txt'.format( # class_name, 'train')), 'r') as f: # train_observed_list = [x.strip() for x in f.readlines()] with open( os.path.join(observed_set_root, "{}_{}.txt".format(class_name, "test")), "r", ) as f: test_observed_list = [x.strip() for x in f.readlines()] with open(rendered_pose_path.format(set_type, class_name)) as f: str_rendered_pose_list = [x.strip().split(" ") for x in f.readlines()] rendered_pose_list = np.array( [[float(x) for x in each_pose] for each_pose in str_rendered_pose_list] ) rendered_per_observed = 10 assert len(rendered_pose_list) == 10 * len( all_observed_list ), "{} vs {}".format(len(rendered_pose_list), len(all_observed_list)) for idx, observed_index in enumerate(tqdm(all_observed_list)): video_name, observed_prefix = observed_index.split("/") rendered_dir = os.path.join(rendered_root_dir, class_name) mkdir_if_missing(rendered_dir) for inner_idx in range(rendered_per_observed): if gen_images: # if gen_images and observed_index in test_observed_list and inner_idx == 0: # only generate my_val_v{} # noqa:E501 image_file = os.path.join( rendered_dir, "{}_{}-color.png".format(observed_prefix, inner_idx), ) depth_file = os.path.join( rendered_dir, "{}_{}-depth.png".format(observed_prefix, inner_idx), ) rendered_idx = idx * rendered_per_observed + inner_idx pose_rendered_q = rendered_pose_list[rendered_idx] rgb_gl, depth_gl = render_machine.render( pose_rendered_q[:4], pose_rendered_q[4:] ) rgb_gl = rgb_gl.astype("uint8") depth_gl = (depth_gl * depth_factor).astype(np.uint16) cv2.imwrite(image_file, rgb_gl) cv2.imwrite(depth_file, depth_gl) pose_rendered_file = os.path.join( rendered_dir, "{}_{}-pose.txt".format(observed_prefix, inner_idx), ) text_file = open(pose_rendered_file, "w") text_file.write("{}\n".format(class_idx)) pose_rendered_m = np.zeros((3, 4)) pose_rendered_m[:, :3] = RT_transform.quat2mat( pose_rendered_q[:4] ) pose_rendered_m[:, 3] = pose_rendered_q[4:] pose_ori_m = pose_rendered_m pose_str = "{} {} {} {}\n{} {} {} {}\n{} {} {} {}".format( pose_ori_m[0, 0], pose_ori_m[0, 1], pose_ori_m[0, 2], pose_ori_m[0, 3], pose_ori_m[1, 0], pose_ori_m[1, 1], pose_ori_m[1, 2], pose_ori_m[1, 3], pose_ori_m[2, 0], pose_ori_m[2, 1], pose_ori_m[2, 2], pose_ori_m[2, 3], ) text_file.write(pose_str) if observed_index in test_observed_list: if inner_idx == 0: val_pair.append( "{} {}/{}_{}".format( observed_index, class_name, observed_prefix, inner_idx, ) ) else: train_pair.append( "{} {}/{}_{}".format( observed_index, class_name, observed_prefix, inner_idx ) ) train_pair_set_file = os.path.join( pair_set_dir, "train_{}.txt".format(class_name) ) train_pair = sorted(train_pair) with open(train_pair_set_file, "w") as text_file: for x in train_pair: text_file.write("{}\n".format(x)) test_pair_set_file = os.path.join( pair_set_dir, "my_val_{}.txt".format(class_name) ) val_pair = sorted(val_pair) with open(test_pair_set_file, "w") as text_file: for x in val_pair: text_file.write("{}\n".format(x)) print(class_name, " done")
def main(): gen_images = ( True ) ##################################################################### gen_images for class_idx, class_name in enumerate(class_list): # train_pair = [] val_pair = [] if class_name in ["__back_ground__"]: continue if not class_name in sel_classes: continue print("start ", class_name) model_path = os.path.join( TLESS_root, "models/{}/obj_{}_scaled.ply".format(class_name, class_name) ) model = inout.load_ply(model_path) for video_name in sel_videos: # 02 for set_type in ["test"]: with open( os.path.join( real_set_root, "{}_{}_{}.txt".format(class_name, video_name, "test"), ), "r", ) as f: test_real_list = [x.strip() for x in f.readlines()] with open( rendered_pose_path.format(set_type, class_name, video_name) ) as f: str_rendered_pose_list = [ x.strip().split(" ") for x in f.readlines() ] rendered_pose_list = np.array( [ [float(x) for x in each_pose] for each_pose in str_rendered_pose_list ] ) rendered_per_real = 1 assert len(rendered_pose_list) == 1 * len( test_real_list ), "{} vs {}".format(len(rendered_pose_list), len(test_real_list)) for idx, real_index in enumerate(tqdm(test_real_list)): _, cls_name, vid_name, real_prefix = real_index.split("/") rendered_dir = os.path.join( rendered_root_dir, class_name, video_name ) mkdir_if_missing(rendered_dir) for inner_idx in range(rendered_per_real): if gen_images: # if gen_images and real_index in test_real_list and inner_idx == 0: # only generate my_val_v{} image_file = os.path.join( rendered_dir, "{}-color.png".format(real_prefix) ) depth_file = os.path.join( rendered_dir, "{}-depth.png".format(real_prefix) ) rendered_idx = idx * rendered_per_real + inner_idx pose_rendered_q = rendered_pose_list[rendered_idx] R = se3.quat2mat(pose_rendered_q[:4]) t = pose_rendered_q[4:] surf_color = None # (1, 0, 0) # ????? im_size = (width, height) # (w, h) ren_rgb, ren_depth = renderer.render( model, im_size, K_0, R, t, clip_near=ZNEAR, clip_far=ZFAR, surf_color=surf_color, mode="rgb+depth", ) ren_depth = (ren_depth * depth_factor).astype(np.uint16) cv2.imwrite(image_file, ren_rgb) cv2.imwrite(depth_file, ren_depth) pose_rendered_file = os.path.join( rendered_dir, "{}-pose.txt".format(real_prefix) ) text_file = open(pose_rendered_file, "w") text_file.write("{}\n".format(class_idx)) pose_rendered_m = np.zeros((3, 4)) pose_rendered_m[:, :3] = se3.quat2mat(pose_rendered_q[:4]) pose_rendered_m[:, 3] = pose_rendered_q[4:] pose_ori_m = pose_rendered_m pose_str = "{} {} {} {}\n{} {} {} {}\n{} {} {} {}".format( pose_ori_m[0, 0], pose_ori_m[0, 1], pose_ori_m[0, 2], pose_ori_m[0, 3], pose_ori_m[1, 0], pose_ori_m[1, 1], pose_ori_m[1, 2], pose_ori_m[1, 3], pose_ori_m[2, 0], pose_ori_m[2, 1], pose_ori_m[2, 2], pose_ori_m[2, 3], ) text_file.write(pose_str) val_pair.append( "{} test/{}/{}/{}".format( real_index, class_name, video_name, real_prefix ) ) test_pair_set_file = os.path.join( pair_set_dir, "my_val_{}_video{}_{}.txt".format(version, video_name, class_name), ) val_pair = sorted(val_pair) with open(test_pair_set_file, "w") as text_file: for x in val_pair: text_file.write("{}\n".format(x)) print(class_name, video_name, " done")
pose_rendered_q[:4], pose_rendered_q[4:]) rgb_gl = rgb_gl.astype('uint8') depth_gl = (depth_gl * depth_factor).astype(np.uint16) cv2.imwrite(image_file, rgb_gl) cv2.imwrite(depth_file, depth_gl) pose_file = os.path.join( rendered_dir, '{}_{}_{}-pose.txt'.format(class_name, real_prefix, inner_idx)) text_file = open(pose_file, 'w') text_file.write("{}\n".format(class_idx)) pose_rendered_m = np.zeros((3, 4)) pose_rendered_m[:, :3] = se3.quat2mat( pose_rendered_q[:4]) pose_rendered_m[:, 3] = pose_rendered_q[4:] pose_ori_m = pose_rendered_m pose_str = "{} {} {} {}\n{} {} {} {}\n{} {} {} {}"\ .format(pose_ori_m[0, 0], pose_ori_m[0, 1], pose_ori_m[0, 2], pose_ori_m[0, 3], pose_ori_m[1, 0], pose_ori_m[1, 1], pose_ori_m[1, 2], pose_ori_m[1, 3], pose_ori_m[2, 0], pose_ori_m[2, 1], pose_ori_m[2, 2], pose_ori_m[2, 3]) text_file.write(pose_str) train_pair.append("{} {}/{}_{}_{}".format( real_index, video_name, class_name, real_prefix, inner_idx)) pair_set_file = os.path.join(pair_set_dir, "train_{}.txt".format(class_name)) train_pair = sorted(train_pair)
def gen_poses(): pz = np.array([0,0,1]) pose_dict, quat_stat, trans_stat, new_points = stat_poses() real_prefix_list = ['{:06d}'.format(i + 1) for i in range(NUM_IMAGES)] sel_classes = classes num_class = len(sel_classes) # syn_pose_dict = {prefix: np.zeros((num_class, 7), dtype='float32') # for prefix in real_prefix_list} # store poses syn_pose_dict = {cls_name: np.zeros((NUM_IMAGES, 7)) for cls_name in sel_classes} syn_pose_dir = os.path.join(cur_path, '..', 'data/LINEMOD_6D/LM6d_converted/LM6d_render_v1/syn_poses_single') mkdir_if_missing(syn_pose_dir) for cls_i, cls_name in enumerate(sel_classes): if cls_name != 'ape': continue src_quat_mean = quat_stat[cls_name]['quat_mean'] src_quat_std = quat_stat[cls_name]['quat_std'] src_trans_mean = trans_stat[cls_name]['trans_mean'] src_trans_std = trans_stat[cls_name]['trans_std'] deg_max = new_points[cls_name]['angle_max'] for i in tqdm(range(NUM_IMAGES)): real_prefix = real_prefix_list[i] # if cls_name == 'ape': # continue # randomly generate a pose tgt_quat = np.random.normal(0, 1, 4) tgt_quat = tgt_quat/np.linalg.norm(tgt_quat) if tgt_quat[0] < 0: tgt_quat *= -1 tgt_trans = np.random.normal(src_trans_mean, src_trans_std) tgt_rot_m = se3.quat2mat(tgt_quat) new_pz = np.dot(tgt_rot_m, pz.reshape((-1, 1))).reshape((3,)) pz_mean = new_points[cls_name]['pz_mean'] deg = angle(new_pz, pz_mean) # r_dist, t_dist = calc_rt_dist_q(tgt_quat, src_quat, tgt_trans, src_trans) transform = np.matmul(K, tgt_trans.reshape(3, 1)) center_x = float(transform[0] / transform[2]) center_y = float(transform[1] / transform[2]) count = 0 while (deg > deg_max or not (48 < center_x < (640 - 48) and 48 < center_y < (480 - 48))): # randomly generate a pose tgt_quat = np.random.normal(0, 1, 4) tgt_quat = tgt_quat / np.linalg.norm(tgt_quat) if tgt_quat[0] < 0: tgt_quat *= -1 tgt_trans = np.random.normal(src_trans_mean, src_trans_std) tgt_rot_m = se3.quat2mat(tgt_quat) new_pz = np.dot(tgt_rot_m, pz.reshape((-1, 1))).reshape((3,)) pz_mean = new_points[cls_name]['pz_mean'] deg = angle(new_pz, pz_mean) transform = np.matmul(K, tgt_trans.reshape(3, 1)) center_x = float(transform[0] / transform[2]) center_y = float(transform[1] / transform[2]) count += 1 if count % 100 == 0: print(real_prefix, cls_name, count, "deg < deg_max={}: {}, 48 < center_x < (640-48): {}, 48 < center_y < (480-48): {}".format( deg_max, deg <= deg_max, 48 < center_x < (640 - 48), 48 < center_y < (480 - 48))) tgt_pose_q = np.zeros((7,), dtype='float32') tgt_pose_q[:4] = tgt_quat tgt_pose_q[4:] = tgt_trans syn_pose_dict[cls_name][i, :] = tgt_pose_q # write pose poses_file = os.path.join(syn_pose_dir, 'LM6d_ds_v1_all_syn_pose.pkl') with open(poses_file, 'wb') as f: cPickle.dump(syn_pose_dict, f, 2)
def main(): gen_images = True ############ for class_idx, class_name in idx2class.items(): val_pair = [] print("start ", class_name) if class_name in ['__back_ground__']: continue if not class_name in ['bowl', 'cup']: continue if gen_images: # init render model_dir = os.path.join(LM6d_root, 'models', class_name) render_machine = Render_Py(model_dir, K, width, height, ZNEAR, ZFAR) for set_type in ['all']: with open( os.path.join(real_set_root, '{}_{}.txt'.format(class_name, 'all')), 'r') as f: all_real_list = [x.strip() for x in f.readlines()] with open( os.path.join(real_set_root, '{}_{}.txt'.format(class_name, 'test')), 'r') as f: test_real_list = [x.strip() for x in f.readlines()] with open(rendered_pose_path.format(set_type, class_name)) as f: str_rendered_pose_list = [ x.strip().split(' ') for x in f.readlines() ] rendered_pose_list = np.array( [[float(x) for x in each_pose] for each_pose in str_rendered_pose_list]) rendered_per_real = 1 assert (len(rendered_pose_list) == 1*len(all_real_list)), \ '{} vs {}'.format(len(rendered_pose_list), len(all_real_list)) for idx, real_index in enumerate(tqdm(all_real_list)): video_name, real_prefix = real_index.split('/') rendered_dir = os.path.join(rendered_root_dir, class_name) mkdir_if_missing(rendered_dir) for inner_idx in range(rendered_per_real): if gen_images: image_file = os.path.join( rendered_dir, '{}_{}-color.png'.format(real_prefix, inner_idx)) depth_file = os.path.join( rendered_dir, '{}_{}-depth.png'.format(real_prefix, inner_idx)) rendered_idx = idx * rendered_per_real + inner_idx pose_rendered_q = rendered_pose_list[rendered_idx] rgb_gl, depth_gl = render_machine.render( pose_rendered_q[:4], pose_rendered_q[4:]) rgb_gl = rgb_gl.astype('uint8') depth_gl = (depth_gl * depth_factor).astype(np.uint16) cv2.imwrite(image_file, rgb_gl) cv2.imwrite(depth_file, depth_gl) pose_rendered_file = os.path.join( rendered_dir, '{}_{}-pose.txt'.format(real_prefix, inner_idx)) text_file = open(pose_rendered_file, 'w') text_file.write("{}\n".format(class_idx)) pose_rendered_m = np.zeros((3, 4)) pose_rendered_m[:, :3] = RT_transform.quat2mat( pose_rendered_q[:4]) pose_rendered_m[:, 3] = pose_rendered_q[4:] pose_ori_m = pose_rendered_m pose_str = "{} {} {} {}\n{} {} {} {}\n{} {} {} {}"\ .format(pose_ori_m[0, 0], pose_ori_m[0, 1], pose_ori_m[0, 2], pose_ori_m[0, 3], pose_ori_m[1, 0], pose_ori_m[1, 1], pose_ori_m[1, 2], pose_ori_m[1, 3], pose_ori_m[2, 0], pose_ori_m[2, 1], pose_ori_m[2, 2], pose_ori_m[2, 3]) text_file.write(pose_str) if real_index in test_real_list: if inner_idx == 0: val_pair.append("{} {}/{}_{}".format( real_index, class_name, real_prefix, inner_idx)) test_pair_set_file = os.path.join( pair_set_dir, "my_val_{}.txt".format(class_name)) val_pair.pose_est with open(test_pair_set_file, "w") as text_file: for x in val_pair: text_file.write("{}\n".format(x)) print(class_name, " done")
def main(): gen_images = True for class_idx, class_name in enumerate(tqdm(classes)): train_pair = [] print("start ", class_name) if class_name in ["__back_ground__"]: continue if gen_images: # init render machine model_dir = os.path.join( cur_path, "../data/LINEMOD_6D/LM6d_converted/models/{}".format( class_name), ) render_machine = Render_Py(model_dir, K, width, height, ZNEAR, ZFAR) for set_type in ["NDtrain"]: # observed index list observed_list_path = os.path.join( observed_set_dir, "NDtrain_observed_{}.txt".format(class_name)) with open(observed_list_path, "r") as f: observed_list = [x.strip() for x in f.readlines()] # rendered poses rendered_pose_path = os.path.join( rendered_pose_dir, "LM6d_occ_dsm_{}_NDtrain_rendered_pose_{}.txt".format( version, class_name), ) with open(rendered_pose_path, "r") as f: str_rendered_pose_list = [ x.strip().split(" ") for x in f.readlines() ] rendered_pose_list = np.array( [[float(x) for x in each_pose] for each_pose in str_rendered_pose_list]) rendered_per_observed = 1 assert len(rendered_pose_list ) == 1 * len(observed_list), "{} vs {}".format( len(rendered_pose_list), len(observed_list)) for idx, observed_index in enumerate(tqdm(observed_list)): video_name, observed_prefix = observed_index.split( "/") # ./prefix rendered_dir = os.path.join(rendered_root_dir, video_name) mkdir_if_missing(rendered_dir) rendered_dir = os.path.join(rendered_dir, class_name) mkdir_if_missing(rendered_dir) for inner_idx in range(rendered_per_observed): if gen_images: image_file = os.path.join( rendered_dir, "{}_{}-color.png".format(observed_prefix, inner_idx), ) depth_file = os.path.join( rendered_dir, "{}_{}-depth.png".format(observed_prefix, inner_idx), ) # if os.path.exists(image_file) and os.path.exists(depth_file): # continue rendered_idx = idx * rendered_per_observed + inner_idx pose_rendered_q = rendered_pose_list[rendered_idx] rgb_gl, depth_gl = render_machine.render( pose_rendered_q[:4], pose_rendered_q[4:]) rgb_gl = rgb_gl.astype("uint8") depth_gl = (depth_gl * depth_factor).astype(np.uint16) cv2.imwrite(image_file, rgb_gl) cv2.imwrite(depth_file, depth_gl) pose_rendered_file = os.path.join( rendered_dir, "{}_{}-pose.txt".format(observed_prefix, inner_idx), ) text_file = open(pose_rendered_file, "w") text_file.write("{}\n".format(class2idx(class_name))) pose_rendered_m = np.zeros((3, 4)) pose_rendered_m[:, :3] = se3.quat2mat( pose_rendered_q[:4]) pose_rendered_m[:, 3] = pose_rendered_q[4:] pose_ori_m = pose_rendered_m pose_str = "{} {} {} {}\n{} {} {} {}\n{} {} {} {}".format( pose_ori_m[0, 0], pose_ori_m[0, 1], pose_ori_m[0, 2], pose_ori_m[0, 3], pose_ori_m[1, 0], pose_ori_m[1, 1], pose_ori_m[1, 2], pose_ori_m[1, 3], pose_ori_m[2, 0], pose_ori_m[2, 1], pose_ori_m[2, 2], pose_ori_m[2, 3], ) text_file.write(pose_str) train_pair.append("{} {}/{}_{}".format( observed_index, class_name, observed_prefix, inner_idx)) pair_set_file = os.path.join(image_set_dir, "train_{}.txt".format(class_name)) train_pair = sorted(train_pair) with open(pair_set_file, "w") as text_file: for x in train_pair: text_file.write("{}\n".format(x)) print(class_name, " done")
def main(): gen_images = ( True ) ##################################################################### gen_images for class_idx, class_name in enumerate(class_list): train_pair = [] # val_pair = [] if class_name in ["__back_ground__"]: continue if not class_name in sel_classes: continue print("start ", class_name) if gen_images: # init render model_dir = os.path.join(TLESS_root, "models", class_name) render_machine = Render_Py(model_dir, K, width, height, ZNEAR, ZFAR) for set_type in ["train"]: with open( os.path.join(real_set_root, "{}_{}.txt".format(class_name, "train")), "r", ) as f: train_real_list = [x.strip() for x in f.readlines()] with open(rendered_pose_path.format(set_type, class_name)) as f: str_rendered_pose_list = [ x.strip().split(" ") for x in f.readlines() ] rendered_pose_list = np.array( [[float(x) for x in each_pose] for each_pose in str_rendered_pose_list]) rendered_per_real = 10 assert len(rendered_pose_list) == rendered_per_real * len( train_real_list), "{} vs {}".format(len(rendered_pose_list), len(train_real_list)) for idx, real_index in enumerate(tqdm(train_real_list)): video_name, real_prefix = real_index.split("/") rendered_dir = os.path.join(rendered_root_dir, class_name) mkdir_if_missing(rendered_dir) for inner_idx in range(rendered_per_real): if gen_images: # if gen_images and real_index in test_real_list and inner_idx == 0: # only generate my_val_v{} image_file = os.path.join( rendered_dir, "{}_{}-color.png".format(real_prefix, inner_idx), ) depth_file = os.path.join( rendered_dir, "{}_{}-depth.png".format(real_prefix, inner_idx), ) rendered_idx = idx * rendered_per_real + inner_idx pose_rendered_q = rendered_pose_list[rendered_idx] rgb_gl, depth_gl = render_machine.render( pose_rendered_q[:4], pose_rendered_q[4:]) rgb_gl = rgb_gl.astype("uint8") depth_gl = (depth_gl * depth_factor).astype(np.uint16) cv2.imwrite(image_file, rgb_gl) cv2.imwrite(depth_file, depth_gl) pose_rendered_file = os.path.join( rendered_dir, "{}_{}-pose.txt".format(real_prefix, inner_idx), ) text_file = open(pose_rendered_file, "w") text_file.write("{}\n".format(class_idx)) pose_rendered_m = np.zeros((3, 4)) pose_rendered_m[:, :3] = se3.quat2mat( pose_rendered_q[:4]) pose_rendered_m[:, 3] = pose_rendered_q[4:] pose_ori_m = pose_rendered_m pose_str = "{} {} {} {}\n{} {} {} {}\n{} {} {} {}".format( pose_ori_m[0, 0], pose_ori_m[0, 1], pose_ori_m[0, 2], pose_ori_m[0, 3], pose_ori_m[1, 0], pose_ori_m[1, 1], pose_ori_m[1, 2], pose_ori_m[1, 3], pose_ori_m[2, 0], pose_ori_m[2, 1], pose_ori_m[2, 2], pose_ori_m[2, 3], ) text_file.write(pose_str) train_pair.append("{} {}/{}_{}".format( real_index, class_name, real_prefix, inner_idx)) train_pair_set_file = os.path.join( pair_set_dir, "train_{}_{}.txt".format(version, class_name)) train_pair = sorted(train_pair) with open(train_pair_set_file, "w") as text_file: for x in train_pair: text_file.write("{}\n".format(x)) print(class_name, " done")
def pose_q2m(pose_q): pose = np.zeros((3, 4), dtype="float32") pose[:3, :3] = RT_transform.quat2mat(pose_q[:4]) pose[:3, 3] = pose_q[4:] return pose