Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,486 Bytes
5f9d349 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
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)
|