def get_dataset_poses(self): if os.path.exists(self.dataset_poses_path): poses = np.load(self.dataset_poses_path) return poses[:, :3], poses[:, 3:] eulers = [] translations = [] train_set = np.loadtxt( os.path.join(cfg.LINEMOD, '{}/training_range.txt'.format(self.class_type)), np.int32) for idx in train_set: rot_path = os.path.join(self.dir_path, 'rot{}.rot'.format(idx)) tra_path = os.path.join(self.dir_path, 'tra{}.tra'.format(idx)) pose = read_pose(rot_path, tra_path) euler = self.pose_transformer.orig_pose_to_blender_euler(pose) eulers.append(euler) translations.append(pose[:, 3]) eulers = np.array(eulers) translations = np.array(translations) np.save(self.dataset_poses_path, np.concatenate([eulers, translations], axis=-1)) return eulers, translations
def collect_real_set_info(self): # linemod standard database=[] projector=Projector() modeldb=LineModModelDB() img_num=len(os.listdir(os.path.join(self.linemod_dir,self.rgb_dir))) for k in range(img_num): data={} data['rgb_pth']=os.path.join(self.rgb_dir, '{:06}.jpg'.format(k)) data['dpt_pth']=os.path.join(self.mask_dir, '{:04}.png'.format(k)) pose=read_pose(os.path.join(self.rt_dir, 'rot{}.rot'.format(k)), os.path.join(self.rt_dir, 'tra{}.tra'.format(k))) pose_transformer = PoseTransformer(class_type=self.cls_name) data['RT'] = pose_transformer.orig_pose_to_blender_pose(pose).astype(np.float32) data['cls_typ']=self.cls_name data['rnd_typ']='real' data['corners']=projector.project(modeldb.get_corners_3d(self.cls_name),data['RT'],'linemod') data['farthest']=projector.project(modeldb.get_farthest_3d(self.cls_name),data['RT'],'linemod') for num in [4,12,16,20]: data['farthest{}'.format(num)]=projector.project(modeldb.get_farthest_3d(self.cls_name,num),data['RT'],'linemod') data['center']=projector.project(modeldb.get_centers_3d(self.cls_name)[None, :],data['RT'],'linemod') data['small_bbox'] = projector.project(modeldb.get_small_bbox(self.cls_name), data['RT'], 'linemod') axis_direct=np.concatenate([np.identity(3), np.zeros([3, 1])], 1).astype(np.float32) data['van_pts']=projector.project_h(axis_direct, data['RT'], 'linemod') database.append(data) save_pickle(database,self.set_pkl) return database
def read_benchvise_pose(index): orig_pose_dir_path = os.path.join(cfg.LINEMOD_ORIG, 'benchvise/data') pose = read_pose( os.path.join(orig_pose_dir_path, 'rot{}.rot'.format(index)), os.path.join(orig_pose_dir_path, 'tra{}.tra'.format(index))) pose_transformer = PoseTransformer(class_type='benchvise') return pose_transformer.orig_pose_to_blender_pose(pose).astype( np.float32)
def make_truncated_linemod_dataset(): for cls_name in cfg.linemod_cls_names: print(cls_name) linemod_dir = cfg.LINEMOD rgb_dir = '{}/JPEGImages'.format(cls_name) mask_dir = '{}/mask'.format(cls_name) rt_dir = os.path.join(cfg.DATA_DIR, 'LINEMOD_ORIG', cls_name, 'data') if not os.path.exists( os.path.join(linemod_dir, 'truncated', cls_name)): os.mkdir(os.path.join(linemod_dir, 'truncated', cls_name)) projector = Projector() img_num = len(os.listdir(os.path.join(linemod_dir, rgb_dir))) print(img_num) for k in range(img_num): rgb = imread( os.path.join(linemod_dir, rgb_dir, '{:06}.jpg'.format(k))) msk = imread( os.path.join(linemod_dir, mask_dir, '{:04}.png'.format(k))) msk = (np.sum(msk, 2) > 0).astype(np.uint8) before = np.sum(msk) count = 0 while True: rgb_new, msk_new, hbeg, wbeg = LineModImageDB.crop_instance( rgb, msk, 256) after = np.sum(msk_new) count += 1 if after / before >= 0.2 or count > 50: rgb, msk = rgb_new, msk_new break imsave( os.path.join(linemod_dir, 'truncated', cls_name, '{:06}_rgb.jpg'.format(k)), rgb) imsave( os.path.join(linemod_dir, 'truncated', cls_name, '{:04}_msk.png'.format(k)), msk) pose = read_pose(os.path.join(rt_dir, 'rot{}.rot'.format(k)), os.path.join(rt_dir, 'tra{}.tra'.format(k))) pose_transformer = PoseTransformer(class_type=cls_name) pose = pose_transformer.orig_pose_to_blender_pose(pose).astype( np.float32) K = projector.intrinsic_matrix['linemod'].copy() K[0, 2] += wbeg K[1, 2] += hbeg save_pickle([pose, K], os.path.join(linemod_dir, 'truncated', cls_name, '{:06}_info.pkl'.format(k))) if k % 500 == 0: print(k)
def collect_real_set_info(self): ''' pvnet 的数据集linemod有做更改, cat.ply 与原始的linemod中的cat.ply,有模型偏移和旋转, 所以 原始数据集中的pose真值,需要 简单变换,就是这边数据集的 pose 会这样问了,既然图片数据集 这是没有改变的,怎么RT改变呢 因为 这边提的3d特征是 新的model上取的,所以计算RT的时候, RT要变的 pose_real.pkl ''' database = [] projector = Projector() modeldb = LineModModelDB() img_num = len(os.listdir(os.path.join(self.linemod_dir, self.rgb_dir))) for k in range(img_num): data = {} data['rgb_pth'] = os.path.join(self.rgb_dir, '{:06}.jpg'.format(k)) data['dpt_pth'] = os.path.join(self.mask_dir, '{:04}.png'.format(k)) pose = read_pose(os.path.join(self.rt_dir, 'rot{}.rot'.format(k)), os.path.join(self.rt_dir, 'tra{}.tra'.format(k))) pose_transformer = PoseTransformer(class_type=self.cls_name) data['RT'] = pose_transformer.orig_pose_to_blender_pose( pose).astype(np.float32) data['cls_typ'] = self.cls_name data['rnd_typ'] = 'real' data['corners'] = projector.project( modeldb.get_corners_3d(self.cls_name), data['RT'], 'linemod') data['farthest'] = projector.project( modeldb.get_farthest_3d(self.cls_name), data['RT'], 'linemod') for num in [4, 12, 16, 20]: data['farthest{}'.format(num)] = projector.project( modeldb.get_farthest_3d(self.cls_name, num), data['RT'], 'linemod') data['center'] = projector.project( modeldb.get_centers_3d(self.cls_name)[None, :], data['RT'], 'linemod') data['small_bbox'] = projector.project( modeldb.get_small_bbox(self.cls_name), data['RT'], 'linemod') axis_direct = np.concatenate( [np.identity(3), np.zeros([3, 1])], 1).astype(np.float32) data['van_pts'] = projector.project_h(axis_direct, data['RT'], 'linemod') database.append(data) save_pickle(database, self.real_pkl) return database