sizes = [0.5 for _ in range(len(pcds))]
fig = plt.figure(figsize=(len(pcds) * 3, 9))
for i in range(3):
elev = 30
azim = -45 + 90 * i
for j, (pcd, size) in enumerate(zip(pcds, sizes)):
color = pcd[:, 0]
ax = fig.add_subplot(3, len(pcds), i * len(pcds) + j + 1, projection='3d')
ax.view_init(elev, azim)
ax.scatter(pcd[:, 0], pcd[:, 1], pcd[:, 2], zdir=zdir, c=color, s=size, cmap=cmap, vmin=-1, vmax=0.5)
ax.set_title(titles[j])
ax.set_axis_off()
ax.set_xlim(xlim)
ax.set_ylim(ylim)
ax.set_zlim(zlim)
plt.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.9, wspace=0.1, hspace=0.1)
plt.suptitle(suptitle)
fig.savefig(filename)
plt.close(fig)
if __name__ == "__main__":
filenames = ['airplane.pcd', 'car.pcd', 'chair.pcd', 'lamp.pcd'] # '../demo_data'
for file in filenames:
filename = file.replace('.pcd', '')
pcds = [np.asarray(read_point_cloud('../demo_data/' + file).points)]
titles = ['viewpoint 1', 'viewpoint 2', 'viewpoint 3']
plot_pcd_three_views(
filename, pcds, titles, suptitle=filename, sizes=None, cmap='viridis', zdir='y',
xlim=(-0.3, 0.3), ylim=(-0.3, 0.3), zlim=(-0.3, 0.3))
def preprocess_data():
# Create directory
Path(f'{SAVE_DIR}/image').mkdir(parents=True, exist_ok=True)
Path(f'{SAVE_DIR}/model').mkdir(parents=True, exist_ok=True)
# Parse metadata and load information of images with the target object
img_info = []
if not IS_TARGET:
for sc_id in range(*SCENE_ID_RANGE):
assert Path(f'{DATASET_DIR}/{sc_id:06d}/scene_gt.json').is_file()
scene_gt = io.load_scene_gt(
f'{DATASET_DIR}/{sc_id:06d}/scene_gt.json')
assert Path(
f'{DATASET_DIR}/{sc_id:06d}/scene_camera.json').is_file()
scene_cam = io.load_scene_camera(
f'{DATASET_DIR}/{sc_id:06d}/scene_camera.json')
for (im_id, im_gt), im_cam in zip(scene_gt.items(),
scene_cam.values()):
for gt_id, gt in enumerate(im_gt):
if int(gt['obj_id']) == OBJ_ID:
assert Path(
f'{DATASET_DIR}/{sc_id:06d}/depth/{im_id:06d}.png'
).is_file()
depth_path = f'{DATASET_DIR}/{sc_id:06d}/depth/{im_id:06d}.png'
assert Path(
f'{DATASET_DIR}/{sc_id:06d}/mask_visib/{im_id:06d}_{gt_id:06d}.png'
).is_file()
mask_path = f'{DATASET_DIR}/{sc_id:06d}/mask_visib/{im_id:06d}_{gt_id:06d}.png'
save_name = f'{sc_id:06d}_{im_id:06d}_{gt_id:06d}'
img_info.append(
[depth_path, mask_path, im_cam, gt, save_name])
else:
targets = io.load_json(
f'{Path(DATASET_DIR).parent}/test_targets_bop19.json')
for target in targets:
if int(target['obj_id']) == OBJ_ID:
sc_id = target['scene_id']
assert Path(
f'{DATASET_DIR}/{sc_id:06d}/scene_gt.json').is_file()
scene_gt = io.load_scene_gt(
f'{DATASET_DIR}/{sc_id:06d}/scene_gt.json')
assert Path(
f'{DATASET_DIR}/{sc_id:06d}/scene_camera.json').is_file()
scene_cam = io.load_scene_camera(
f'{DATASET_DIR}/{sc_id:06d}/scene_camera.json')
im_id = int(target['im_id'])
im_gt, im_cam = scene_gt[im_id], scene_cam[im_id]
for gt_id, gt in enumerate(im_gt):
if int(gt['obj_id']) == OBJ_ID:
assert Path(
f'{DATASET_DIR}/{sc_id:06d}/depth/{im_id:06d}.png'
).is_file()
depth_path = f'{DATASET_DIR}/{sc_id:06d}/depth/{im_id:06d}.png'
assert Path(
f'{DATASET_DIR}/{sc_id:06d}/mask_visib/{im_id:06d}_{gt_id:06d}.png'
).is_file()
mask_path = f'{DATASET_DIR}/{sc_id:06d}/mask_visib/{im_id:06d}_{gt_id:06d}.png'
save_name = f'{sc_id:06d}_{im_id:06d}_{gt_id:06d}'
img_info.append(
[depth_path, mask_path, im_cam, gt, save_name])
# Read model point cloud
model_file = f'{Path(DATASET_DIR).parent}/models/obj_{OBJ_ID:06d}.ply'
assert Path(model_file).is_file()
model_pcd = read_point_cloud(model_file).voxel_down_sample(VOXEL_SIZE)
md_pcd_pts = np.asarray(model_pcd.points)
# Create point cloud from image information
t1 = time()
with get_context('spawn').Pool(8) as pool:
jobs = [
pool.apply_async(create_patch_pair, (*info, md_pcd_pts))
for info in img_info
]
df = pd.DataFrame([j.get() for j in jobs]).dropna(axis=0)
if Path(f'{SAVE_DIR}/labels.csv').is_file():
old_df = pd.read_csv(f'{SAVE_DIR}/labels.csv', header=None)
df = pd.concat([old_df, df])
df.to_csv(f'{SAVE_DIR}/labels.csv',
header=['filename', 'num_patches'],
index=False)
t2 = time()
print(
f'Created patch_pairs from {len(df)} images. Time elapsed: {t2 - t1 :.3f}'
)
def read_pcd(filename):
pcd = read_point_cloud(filename)
return np.array(pcd.points)