Spaces:
Running
Running
File size: 1,251 Bytes
4d4dd90 |
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 |
import poselib
import torch
from omegaconf import OmegaConf
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
|