53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
|
import pycolmap
|
||
|
|
from omegaconf import OmegaConf
|
||
|
|
import torch
|
||
|
|
from ...geometry.wrappers import Pose
|
||
|
|
|
||
|
|
from ..base_estimator import BaseEstimator
|
||
|
|
|
||
|
|
|
||
|
|
class PycolmapTwoViewEstimator(BaseEstimator):
|
||
|
|
default_conf = {
|
||
|
|
"ransac_th": 4.0,
|
||
|
|
"options": {**pycolmap.TwoViewGeometryOptions().todict()},
|
||
|
|
}
|
||
|
|
|
||
|
|
required_data_keys = ["m_kpts0", "m_kpts1", "camera0", "camera1"]
|
||
|
|
|
||
|
|
def _init(self, conf):
|
||
|
|
opts = OmegaConf.to_container(conf.options)
|
||
|
|
self.options = pycolmap.TwoViewGeometryOptions(opts)
|
||
|
|
self.options.ransac.max_error = conf.ransac_th
|
||
|
|
|
||
|
|
def _forward(self, data):
|
||
|
|
pts0, pts1 = data["m_kpts0"], data["m_kpts1"]
|
||
|
|
camera0 = data["camera0"]
|
||
|
|
camera1 = data["camera1"]
|
||
|
|
info = pycolmap.two_view_geometry_estimation(
|
||
|
|
pts0.numpy(),
|
||
|
|
pts1.numpy(),
|
||
|
|
camera0.to_cameradict(),
|
||
|
|
camera1.to_cameradict(),
|
||
|
|
self.options,
|
||
|
|
)
|
||
|
|
success = info["success"]
|
||
|
|
if success:
|
||
|
|
R = pycolmap.qvec_to_rotmat(info["qvec"])
|
||
|
|
t = info["tvec"]
|
||
|
|
M = Pose.from_Rt(torch.tensor(R), torch.tensor(t)).to(pts0)
|
||
|
|
inl = torch.tensor(info.pop("inliers")).to(pts0)
|
||
|
|
else:
|
||
|
|
M = Pose.from_4x4mat(torch.eye(4)).to(pts0)
|
||
|
|
inl = torch.zeros_like(pts0[:, 0]).bool()
|
||
|
|
|
||
|
|
estimation = {
|
||
|
|
"success": success,
|
||
|
|
"M_0to1": M,
|
||
|
|
"inliers": inl,
|
||
|
|
"type": str(
|
||
|
|
info.get("configuration_type", pycolmap.TwoViewGeometry.UNDEFINED)
|
||
|
|
),
|
||
|
|
}
|
||
|
|
|
||
|
|
return estimation
|