45 lines
1.2 KiB
Python
45 lines
1.2 KiB
Python
|
import poselib
|
||
|
|
from omegaconf import OmegaConf
|
||
|
|
import torch
|
||
|
|
from ...geometry.wrappers import Pose
|
||
|
|
|
||
|
|
from ..base_estimator import BaseEstimator
|
||
|
|
|
||
|
|
|
||
|
|
class PoseLibRelativePoseEstimator(BaseEstimator):
|
||
|
|
default_conf = {"ransac_th": 2.0, "options": {}}
|
||
|
|
|
||
|
|
required_data_keys = ["m_kpts0", "m_kpts1", "camera0", "camera1"]
|
||
|
|
|
||
|
|
def _init(self, conf):
|
||
|
|
pass
|
||
|
|
|
||
|
|
def _forward(self, data):
|
||
|
|
pts0, pts1 = data["m_kpts0"], data["m_kpts1"]
|
||
|
|
camera0 = data["camera0"]
|
||
|
|
camera1 = data["camera1"]
|
||
|
|
M, info = poselib.estimate_relative_pose(
|
||
|
|
pts0.numpy(),
|
||
|
|
pts1.numpy(),
|
||
|
|
camera0.to_cameradict(),
|
||
|
|
camera1.to_cameradict(),
|
||
|
|
{
|
||
|
|
"max_epipolar_error": self.conf.ransac_th,
|
||
|
|
**OmegaConf.to_container(self.conf.options),
|
||
|
|
},
|
||
|
|
)
|
||
|
|
success = M is not None
|
||
|
|
if success:
|
||
|
|
M = Pose.from_Rt(torch.tensor(M.R), torch.tensor(M.t)).to(pts0)
|
||
|
|
else:
|
||
|
|
M = Pose.from_4x4mat(torch.eye(4)).to(pts0)
|
||
|
|
|
||
|
|
estimation = {
|
||
|
|
"success": success,
|
||
|
|
"M_0to1": M,
|
||
|
|
"inliers": torch.tensor(info.pop("inliers")).to(pts0),
|
||
|
|
**info,
|
||
|
|
}
|
||
|
|
|
||
|
|
return estimation
|