use crate::bundle_adjustment::refine_absolute_pose;
use crate::core::Estimator;
use crate::models::AbsolutePose;
use crate::types::DataMatrix;
use nalgebra::{Vector2, Vector3};
pub struct AbsolutePoseEstimator;
impl Default for AbsolutePoseEstimator {
fn default() -> Self {
Self::new()
}
}
impl AbsolutePoseEstimator {
pub fn new() -> Self {
Self
}
}
impl Estimator for AbsolutePoseEstimator {
type Model = AbsolutePose;
fn sample_size(&self) -> usize {
3 }
fn is_valid_sample(&self, data: &DataMatrix, sample: &[usize]) -> bool {
if sample.len() < self.sample_size() {
return false;
}
for i in 0..sample.len() {
for j in (i + 1)..sample.len() {
if sample[i] == sample[j] {
return false;
}
}
}
if data.ncols() < 5 {
return false;
}
true
}
fn estimate_model(&self, data: &DataMatrix, sample: &[usize]) -> Vec<Self::Model> {
let n = sample.len();
if n < self.sample_size() || data.ncols() < 5 {
return Vec::new();
}
use nalgebra::{DMatrix, Matrix3, SVD, Vector3};
let mut a = DMatrix::<f64>::zeros(2 * n, 12);
for (i, &idx) in sample.iter().enumerate() {
let u = data[(idx, 0)];
let v = data[(idx, 1)];
let x = data[(idx, 2)];
let y = data[(idx, 3)];
let z = data[(idx, 4)];
a[(2 * i, 0)] = x;
a[(2 * i, 1)] = y;
a[(2 * i, 2)] = z;
a[(2 * i, 3)] = 1.0;
a[(2 * i, 8)] = -u * x;
a[(2 * i, 9)] = -u * y;
a[(2 * i, 10)] = -u * z;
a[(2 * i, 11)] = -u;
a[(2 * i + 1, 4)] = x;
a[(2 * i + 1, 5)] = y;
a[(2 * i + 1, 6)] = z;
a[(2 * i + 1, 7)] = 1.0;
a[(2 * i + 1, 8)] = -v * x;
a[(2 * i + 1, 9)] = -v * y;
a[(2 * i + 1, 10)] = -v * z;
a[(2 * i + 1, 11)] = -v;
}
let svd = SVD::new(a, false, true);
let vt = match svd.v_t {
Some(vt) => vt,
None => return Vec::new(),
};
let v = vt.transpose();
let last_col = v.column(v.ncols() - 1);
if last_col.iter().any(|&x| x.is_nan() || x.is_infinite()) {
return Vec::new();
}
let mut r = Matrix3::<f64>::zeros();
r[(0, 0)] = last_col[0];
r[(0, 1)] = last_col[1];
r[(0, 2)] = last_col[2];
r[(1, 0)] = last_col[4];
r[(1, 1)] = last_col[5];
r[(1, 2)] = last_col[6];
r[(2, 0)] = last_col[8];
r[(2, 1)] = last_col[9];
r[(2, 2)] = last_col[10];
let r_svd = SVD::new(r, true, true);
let u_r = r_svd.u.unwrap();
let vt_r = r_svd.v_t.unwrap();
let v_r = vt_r.transpose();
let r_ortho = u_r * v_r.transpose();
let mut r_final = r_ortho;
if r_final.determinant() < 0.0 {
let mut u_neg = u_r;
u_neg.column_mut(2).neg_mut();
r_final = u_neg * v_r.transpose();
}
let t = Vector3::<f64>::new(last_col[3], last_col[7], last_col[11]);
vec![AbsolutePose::from_rt(r_final, t)]
}
fn estimate_model_nonminimal(
&self,
data: &DataMatrix,
sample: &[usize],
weights: Option<&[f64]>,
) -> Vec<Self::Model> {
if sample.len() < 4 {
return self.estimate_model(data, sample);
}
let mut points_2d = Vec::new();
let mut points_3d = Vec::new();
for &idx in sample {
if idx >= data.nrows() || data.ncols() < 5 {
continue;
}
points_2d.push([data[(idx, 0)], data[(idx, 1)]]);
points_3d.push([data[(idx, 2)], data[(idx, 3)], data[(idx, 4)]]);
}
if points_2d.len() < 4 {
return Vec::new();
}
#[cfg(feature = "kornia-pnp")]
{
let _ = (&points_2d, &points_3d, weights);
}
let initial_models = self.estimate_model(data, sample);
if initial_models.is_empty() {
return Vec::new();
}
let rot = initial_models[0].rotation.to_rotation_matrix();
let mut r = *rot.matrix();
let mut t_vec = initial_models[0].translation.vector;
if sample.len() >= 4 {
let mut points_2d_vec = Vec::new();
let mut points_3d_vec = Vec::new();
let mut weights_vec = Vec::new();
for &idx in sample {
points_2d_vec.push(Vector2::new(data[(idx, 0)], data[(idx, 1)]));
points_3d_vec.push(Vector3::new(data[(idx, 2)], data[(idx, 3)], data[(idx, 4)]));
if let Some(w) = weights {
weights_vec.push(w[idx]);
}
}
let weights_slice = if weights_vec.is_empty() {
None
} else {
Some(weights_vec.as_slice())
};
refine_absolute_pose(
&points_2d_vec,
&points_3d_vec,
&mut r,
&mut t_vec,
weights_slice,
100,
);
}
vec![AbsolutePose::from_rt(r, t_vec)]
}
fn is_valid_model(
&self,
model: &Self::Model,
_data: &DataMatrix,
_sample: &[usize],
_threshold: f64,
) -> bool {
let r = model.rotation.to_rotation_matrix();
let det = r.matrix().determinant();
(det - 1.0).abs() < 1e-2 * _threshold.max(1.0)
}
}