(info_file, save_path, total_scenes_index, total_scenes_count)
| 72 | |
| 73 | |
| 74 | def process(info_file, save_path, total_scenes_index, total_scenes_count): |
| 75 | |
| 76 | # gt depth data loader |
| 77 | width, height = 640, 480 |
| 78 | transform = transforms.Compose([ |
| 79 | transforms.ResizeImage((width,height)), |
| 80 | transforms.ToTensor(), |
| 81 | ]) |
| 82 | dataset = SceneDataset(info_file, transform, frame_types=['depth']) |
| 83 | dataloader = torch.utils.data.DataLoader(dataset, batch_size=None, |
| 84 | batch_sampler=None, num_workers=2) |
| 85 | scene = dataset.info['scene'] |
| 86 | |
| 87 | # get info about tsdf |
| 88 | file_tsdf_pred = os.path.join(save_path, '%s.npz'%scene) |
| 89 | temp = TSDF.load(file_tsdf_pred) |
| 90 | voxel_size = int(temp.voxel_size*100) |
| 91 | |
| 92 | # re-fuse to remove hole filling since filled holes are penalized in |
| 93 | # mesh metrics |
| 94 | vol_dim = list(temp.tsdf_vol.shape) |
| 95 | origin = temp.origin |
| 96 | tsdf_fusion = TSDFFusion(vol_dim, float(voxel_size)/100, origin, color=False) |
| 97 | device = tsdf_fusion.device |
| 98 | |
| 99 | # mesh renderer |
| 100 | renderer = Renderer() |
| 101 | mesh_file = os.path.join(save_path, '%s.ply'%scene) |
| 102 | mesh = trimesh.load(mesh_file, process=False) |
| 103 | mesh_opengl = renderer.mesh_opengl(mesh) |
| 104 | |
| 105 | for i, d in enumerate(dataloader): |
| 106 | if i%25==0: |
| 107 | print(total_scenes_index, total_scenes_count,scene, i, len(dataloader)) |
| 108 | |
| 109 | depth_trgt = d['depth'].numpy() |
| 110 | _, depth_pred = renderer(height, width, d['intrinsics'], d['pose'], mesh_opengl) |
| 111 | |
| 112 | temp = eval_depth(depth_pred, depth_trgt) |
| 113 | if i==0: |
| 114 | metrics_depth = temp |
| 115 | else: |
| 116 | metrics_depth = {key:value+temp[key] |
| 117 | for key, value in metrics_depth.items()} |
| 118 | |
| 119 | # # play video visualizations of depth |
| 120 | # viz1 = (np.clip((depth_trgt-.5)/5,0,1)*255).astype(np.uint8) |
| 121 | # viz2 = (np.clip((depth_pred-.5)/5,0,1)*255).astype(np.uint8) |
| 122 | # viz1 = cv2.applyColorMap(viz1, cv2.COLORMAP_JET) |
| 123 | # viz2 = cv2.applyColorMap(viz2, cv2.COLORMAP_JET) |
| 124 | # viz1[depth_trgt==0]=0 |
| 125 | # viz2[depth_pred==0]=0 |
| 126 | # viz = np.hstack((viz1,viz2)) |
| 127 | # cv2.imshow('test', viz) |
| 128 | # cv2.waitKey(1) |
| 129 | |
| 130 | tsdf_fusion.integrate((d['intrinsics'] @ d['pose'].inverse()[:3,:]).to(device), |
| 131 | torch.as_tensor(depth_pred).to(device)) |
no test coverage detected