Spaces:
Running
on
Zero
Running
on
Zero
| import numpy as np | |
| import argparse | |
| import cv2 | |
| from joblib import delayed, Parallel | |
| import json | |
| from read_write_model import * | |
| def get_scales(key, cameras, images, points3d_ordered, args): | |
| image_meta = images[key] | |
| cam_intrinsic = cameras[image_meta.camera_id] | |
| pts_idx = images_metas[key].point3D_ids | |
| mask = pts_idx >= 0 | |
| mask *= pts_idx < len(points3d_ordered) | |
| pts_idx = pts_idx[mask] | |
| valid_xys = image_meta.xys[mask] | |
| if len(pts_idx) > 0: | |
| pts = points3d_ordered[pts_idx] | |
| else: | |
| pts = np.array([0, 0, 0]) | |
| R = qvec2rotmat(image_meta.qvec) | |
| pts = np.dot(pts, R.T) + image_meta.tvec | |
| invcolmapdepth = 1. / pts[..., 2] | |
| n_remove = len(image_meta.name.split('.')[-1]) + 1 | |
| invmonodepthmap = cv2.imread(f"{args.depths_dir}/{image_meta.name[:-n_remove]}.png", cv2.IMREAD_UNCHANGED) | |
| if invmonodepthmap is None: | |
| return None | |
| if invmonodepthmap.ndim != 2: | |
| invmonodepthmap = invmonodepthmap[..., 0] | |
| invmonodepthmap = invmonodepthmap.astype(np.float32) / (2**16) | |
| s = invmonodepthmap.shape[0] / cam_intrinsic.height | |
| maps = (valid_xys * s).astype(np.float32) | |
| valid = ( | |
| (maps[..., 0] >= 0) * | |
| (maps[..., 1] >= 0) * | |
| (maps[..., 0] < cam_intrinsic.width * s) * | |
| (maps[..., 1] < cam_intrinsic.height * s) * (invcolmapdepth > 0)) | |
| if valid.sum() > 10 and (invcolmapdepth.max() - invcolmapdepth.min()) > 1e-3: | |
| maps = maps[valid, :] | |
| invcolmapdepth = invcolmapdepth[valid] | |
| invmonodepth = cv2.remap(invmonodepthmap, maps[..., 0], maps[..., 1], interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)[..., 0] | |
| ## Median / dev | |
| t_colmap = np.median(invcolmapdepth) | |
| s_colmap = np.mean(np.abs(invcolmapdepth - t_colmap)) | |
| t_mono = np.median(invmonodepth) | |
| s_mono = np.mean(np.abs(invmonodepth - t_mono)) | |
| scale = s_colmap / s_mono | |
| offset = t_colmap - t_mono * scale | |
| else: | |
| scale = 0 | |
| offset = 0 | |
| return {"image_name": image_meta.name[:-n_remove], "scale": scale, "offset": offset} | |
| if __name__ == '__main__': | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--base_dir', default="../data/big_gaussians/standalone_chunks/campus") | |
| parser.add_argument('--depths_dir', default="../data/big_gaussians/standalone_chunks/campus/depths_any") | |
| parser.add_argument('--model_type', default="bin") | |
| args = parser.parse_args() | |
| cam_intrinsics, images_metas, points3d = read_model(os.path.join(args.base_dir, "sparse", "0"), ext=f".{args.model_type}") | |
| pts_indices = np.array([points3d[key].id for key in points3d]) | |
| pts_xyzs = np.array([points3d[key].xyz for key in points3d]) | |
| points3d_ordered = np.zeros([pts_indices.max()+1, 3]) | |
| points3d_ordered[pts_indices] = pts_xyzs | |
| # depth_param_list = [get_scales(key, cam_intrinsics, images_metas, points3d_ordered, args) for key in images_metas] | |
| depth_param_list = Parallel(n_jobs=-1, backend="threading")( | |
| delayed(get_scales)(key, cam_intrinsics, images_metas, points3d_ordered, args) for key in images_metas | |
| ) | |
| depth_params = { | |
| depth_param["image_name"]: {"scale": depth_param["scale"], "offset": depth_param["offset"]} | |
| for depth_param in depth_param_list if depth_param != None | |
| } | |
| with open(f"{args.base_dir}/sparse/0/depth_params.json", "w") as f: | |
| json.dump(depth_params, f, indent=2) | |
| print(0) | |