def get_pts_3d(vote_type,class_type): linemod_db=LineModModelDB() if vote_type==VotingType.BB8C: points_3d = linemod_db.get_corners_3d(class_type) points_3d = np.concatenate([points_3d,linemod_db.get_centers_3d(class_type)[None,:]],0) elif vote_type==VotingType.BB8S: points_3d = linemod_db.get_small_bbox(class_type) points_3d = np.concatenate([points_3d,linemod_db.get_centers_3d(class_type)[None,:]],0) elif vote_type==VotingType.Farthest: points_3d = linemod_db.get_farthest_3d(class_type) points_3d = np.concatenate([points_3d, linemod_db.get_centers_3d(class_type)[None, :]], 0) elif vote_type==VotingType.Farthest4: points_3d = linemod_db.get_farthest_3d(class_type,4) points_3d = np.concatenate([points_3d, linemod_db.get_centers_3d(class_type)[None, :]], 0) elif vote_type==VotingType.Farthest12: points_3d = linemod_db.get_farthest_3d(class_type,12) points_3d = np.concatenate([points_3d, linemod_db.get_centers_3d(class_type)[None, :]], 0) elif vote_type==VotingType.Farthest16: points_3d = linemod_db.get_farthest_3d(class_type,16) points_3d = np.concatenate([points_3d, linemod_db.get_centers_3d(class_type)[None, :]], 0) elif vote_type==VotingType.Farthest20: points_3d = linemod_db.get_farthest_3d(class_type,20) points_3d = np.concatenate([points_3d, linemod_db.get_centers_3d(class_type)[None, :]], 0) else: # BB8 points_3d = linemod_db.get_corners_3d(class_type) return points_3d
image_db = LineModImageDB('duck', has_ro_set=False, has_ra_set=False, has_plane_set=False, has_render_set=False, has_ms_set=False,has_fuse_set=False) random.shuffle(image_db.real_set) dataset = LineModDatasetRealAug(image_db.real_set[:5], data_prefix=image_db.linemod_dir, vote_type=VotingType.Extreme, augment=False) sampler = RandomSampler(dataset) batch_sampler = ImageSizeBatchSampler(sampler, 5, False) loader = DataLoader(dataset, batch_sampler=batch_sampler, num_workers=8) modeldb=LineModModelDB() camera_matrix=Projector().intrinsic_matrix['linemod'].astype(np.float32) for i, data in enumerate(loader): rgb, mask, vertex, vertex_weight, pose, gt_corners = data pts2d=gt_corners[0].numpy()[:,:2].astype(np.float32) pts3d=modeldb.get_extreme_3d('duck') pts3d=np.concatenate([pts3d,modeldb.get_centers_3d('duck')[None,:]],0).astype(np.float32) wgt2d=np.zeros([pts2d.shape[0],3]).astype(np.float32) wgt2d[:,(0,2)]=1.0 for k in range(pts2d.shape[0]): if np.random.random()<0.5: scale = np.random.uniform(1, 8) else: scale = np.random.uniform(32, 48) pts2d[k]+=np.random.normal(0,scale,2) wgt2d[k,(0,2)]=1/scale wgt2d/=wgt2d.max() pose_pred=uncertainty_pnp(pts2d,wgt2d,pts3d,camera_matrix) pose_pred2=pnp(pts3d,pts2d,camera_matrix)
dataset = LineModDatasetRealAug(image_db.real_set[:5], data_prefix=image_db.linemod_dir, vote_type=VotingType.Extreme, augment=False) sampler = RandomSampler(dataset) batch_sampler = ImageSizeBatchSampler(sampler, 5, False) loader = DataLoader(dataset, batch_sampler=batch_sampler, num_workers=8) modeldb = LineModModelDB() camera_matrix = Projector().intrinsic_matrix['linemod'].astype(np.float32) for i, data in enumerate(loader): rgb, mask, vertex, vertex_weight, pose, gt_corners = data pts2d = gt_corners[0].numpy()[:, :2].astype(np.float32) pts3d = modeldb.get_extreme_3d('duck') pts3d = np.concatenate( [pts3d, modeldb.get_centers_3d('duck')[None, :]], 0).astype(np.float32) wgt2d = np.zeros([pts2d.shape[0], 3]).astype(np.float32) wgt2d[:, (0, 2)] = 1.0 for k in range(pts2d.shape[0]): if np.random.random() < 0.5: scale = np.random.uniform(1, 8) else: scale = np.random.uniform(32, 48) pts2d[k] += np.random.normal(0, scale, 2) wgt2d[k, (0, 2)] = 1 / scale wgt2d /= wgt2d.max() pose_pred = uncertainty_pnp(pts2d, wgt2d, pts3d, camera_matrix) pose_pred2 = pnp(pts3d, pts2d, camera_matrix)