Spaces:
Running
Running
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| import PIL.Image | |
| from tqdm import tqdm | |
| import pycolmap | |
| from ...utils.read_write_model import write_model, read_model | |
| def scene_coordinates(p2D, R_w2c, t_w2c, depth, camera): | |
| assert len(depth) == len(p2D) | |
| ret = pycolmap.image_to_world(p2D, camera._asdict()) | |
| p2D_norm = np.asarray(ret["world_points"]) | |
| p2D_h = np.concatenate([p2D_norm, np.ones_like(p2D_norm[:, :1])], 1) | |
| p3D_c = p2D_h * depth[:, None] | |
| p3D_w = (p3D_c - t_w2c) @ R_w2c | |
| return p3D_w | |
| def interpolate_depth(depth, kp): | |
| h, w = depth.shape | |
| kp = kp / np.array([[w - 1, h - 1]]) * 2 - 1 | |
| assert np.all(kp > -1) and np.all(kp < 1) | |
| depth = torch.from_numpy(depth)[None, None] | |
| kp = torch.from_numpy(kp)[None, None] | |
| grid_sample = torch.nn.functional.grid_sample | |
| # To maximize the number of points that have depth: | |
| # do bilinear interpolation first and then nearest for the remaining points | |
| interp_lin = grid_sample(depth, kp, align_corners=True, mode="bilinear")[0, :, 0] | |
| interp_nn = torch.nn.functional.grid_sample( | |
| depth, kp, align_corners=True, mode="nearest" | |
| )[0, :, 0] | |
| interp = torch.where(torch.isnan(interp_lin), interp_nn, interp_lin) | |
| valid = ~torch.any(torch.isnan(interp), 0) | |
| interp_depth = interp.T.numpy().flatten() | |
| valid = valid.numpy() | |
| return interp_depth, valid | |
| def image_path_to_rendered_depth_path(image_name): | |
| parts = image_name.split("/") | |
| name = "_".join(["".join(parts[0].split("-")), parts[1]]) | |
| name = name.replace("color", "pose") | |
| name = name.replace("png", "depth.tiff") | |
| return name | |
| def project_to_image(p3D, R, t, camera, eps: float = 1e-4, pad: int = 1): | |
| p3D = (p3D @ R.T) + t | |
| visible = p3D[:, -1] >= eps # keep points in front of the camera | |
| p2D_norm = p3D[:, :-1] / p3D[:, -1:].clip(min=eps) | |
| ret = pycolmap.world_to_image(p2D_norm, camera._asdict()) | |
| p2D = np.asarray(ret["image_points"]) | |
| size = np.array([camera.width - pad - 1, camera.height - pad - 1]) | |
| valid = np.all((p2D >= pad) & (p2D <= size), -1) | |
| valid &= visible | |
| return p2D[valid], valid | |
| def correct_sfm_with_gt_depth(sfm_path, depth_folder_path, output_path): | |
| cameras, images, points3D = read_model(sfm_path) | |
| for imgid, img in tqdm(images.items()): | |
| image_name = img.name | |
| depth_name = image_path_to_rendered_depth_path(image_name) | |
| depth = PIL.Image.open(Path(depth_folder_path) / depth_name) | |
| depth = np.array(depth).astype("float64") | |
| depth = depth / 1000.0 # mm to meter | |
| depth[(depth == 0.0) | (depth > 1000.0)] = np.nan | |
| R_w2c, t_w2c = img.qvec2rotmat(), img.tvec | |
| camera = cameras[img.camera_id] | |
| p3D_ids = img.point3D_ids | |
| p3Ds = np.stack([points3D[i].xyz for i in p3D_ids[p3D_ids != -1]], 0) | |
| p2Ds, valids_projected = project_to_image(p3Ds, R_w2c, t_w2c, camera) | |
| invalid_p3D_ids = p3D_ids[p3D_ids != -1][~valids_projected] | |
| interp_depth, valids_backprojected = interpolate_depth(depth, p2Ds) | |
| scs = scene_coordinates( | |
| p2Ds[valids_backprojected], | |
| R_w2c, | |
| t_w2c, | |
| interp_depth[valids_backprojected], | |
| camera, | |
| ) | |
| invalid_p3D_ids = np.append( | |
| invalid_p3D_ids, | |
| p3D_ids[p3D_ids != -1][valids_projected][~valids_backprojected], | |
| ) | |
| for p3did in invalid_p3D_ids: | |
| if p3did == -1: | |
| continue | |
| else: | |
| obs_imgids = points3D[p3did].image_ids | |
| invalid_imgids = list(np.where(obs_imgids == img.id)[0]) | |
| points3D[p3did] = points3D[p3did]._replace( | |
| image_ids=np.delete(obs_imgids, invalid_imgids), | |
| point2D_idxs=np.delete( | |
| points3D[p3did].point2D_idxs, invalid_imgids | |
| ), | |
| ) | |
| new_p3D_ids = p3D_ids.copy() | |
| sub_p3D_ids = new_p3D_ids[new_p3D_ids != -1] | |
| valids = np.ones(np.count_nonzero(new_p3D_ids != -1), dtype=bool) | |
| valids[~valids_projected] = False | |
| valids[valids_projected] = valids_backprojected | |
| sub_p3D_ids[~valids] = -1 | |
| new_p3D_ids[new_p3D_ids != -1] = sub_p3D_ids | |
| img = img._replace(point3D_ids=new_p3D_ids) | |
| assert len(img.point3D_ids[img.point3D_ids != -1]) == len( | |
| scs | |
| ), f"{len(scs)}, {len(img.point3D_ids[img.point3D_ids != -1])}" | |
| for i, p3did in enumerate(img.point3D_ids[img.point3D_ids != -1]): | |
| points3D[p3did] = points3D[p3did]._replace(xyz=scs[i]) | |
| images[imgid] = img | |
| output_path.mkdir(parents=True, exist_ok=True) | |
| write_model(cameras, images, points3D, output_path) | |
| if __name__ == "__main__": | |
| dataset = Path("datasets/7scenes") | |
| outputs = Path("outputs/7Scenes") | |
| SCENES = ["chess", "fire", "heads", "office", "pumpkin", "redkitchen", "stairs"] | |
| for scene in SCENES: | |
| sfm_path = outputs / scene / "sfm_superpoint+superglue" | |
| depth_path = dataset / f"depth/7scenes_{scene}/train/depth" | |
| output_path = outputs / scene / "sfm_superpoint+superglue+depth" | |
| correct_sfm_with_gt_depth(sfm_path, depth_path, output_path) | |