File size: 2,114 Bytes
4f6b78d |
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 |
import numpy as np
import torch
import torch.nn as nn
import cv2
from dust3r.utils.geometry import inv, geotrf, depthmap_to_absolute_camera_coordinates
from dust3r.post_process import estimate_focal_knowing_depth
def camera_parameter_estimation(view_n, view_m, pred_n, pred_m, im_conf):
# for frame n and m.
# pair_n2m: n->m
conf = float(pred_n['conf'].mean() * pred_m['conf'].mean())
K, focal = estimate_intrinsic(view_n['true_shape'], pred_n['pts3d'], pred_n['conf'])
try:
rel_pose_n2m = estimate_extrinsic(view_m['true_shape'], pred_m['pts3d_in_other_view'], im_conf.cpu().numpy(), K)
R_mn, T_mn = rel_pose_n2m[:3, :3], rel_pose_n2m[:3, 3]
except Exception as e:
print(f"Error estimating extrinsic parameters: {e}")
rel_pose_n2m = torch.eye(4).cuda()
R_mn, T_mn = rel_pose_n2m[:3, :3], rel_pose_n2m[:3, 3]
# ptcloud is expressed in camera n
depth_n, depth_m = pred_n['pts3d'][..., 2], geotrf(inv(rel_pose_n2m), pred_m['pts3d_in_other_view'])[..., 2]
return conf, K, focal, R_mn, T_mn, depth_n, depth_m
def estimate_intrinsic(img_shape, pts3d, conf):
H, W = img_shape
pts3d = pts3d.cpu()
pp = torch.tensor((W/2, H/2))
focal = float(estimate_focal_knowing_depth(pts3d[None], pp, focal_mode='weiszfeld'))
K = np.float32([(focal, 0, pp[0]), (0, focal, pp[1]), (0, 0, 1)])
return K, focal
def estimate_extrinsic(img_shape, pts3d, conf, K):
min_conf_thr = 3
H, W = img_shape
H, W = H.item(), W.item()
# estimate the pose of pts1 in image 2
pts3d = pts3d.cpu().numpy()
pixels = np.mgrid[:W, :H].T.astype(np.float32)
msk = (conf > min_conf_thr)
res = cv2.solvePnPRansac(pts3d[msk], pixels[msk], K, None,
iterationsCount=100, reprojectionError=5, flags=cv2.SOLVEPNP_SQPNP)
success, R, T, inliers = res
assert success
R = cv2.Rodrigues(R)[0] # world to cam
pose = inv(np.r_[np.c_[R, T], [(0, 0, 0, 1)]]) # cam to world
return torch.from_numpy(pose.astype(np.float32)).cuda()
|