use crate::core::types::Point;
use crate::error::Result;
use crate::image::Image;
use burn::tensor::backend::Backend;
pub enum ArucoDict {
Dict4X4_50,
Dict6X6_250,
}
#[derive(Clone, Debug, PartialEq)]
pub struct ArucoMarker {
pub id: usize,
pub corners: [Point<f64>; 4],
}
pub struct ArucoDetector {
pub dictionary: ArucoDict,
}
impl ArucoDetector {
#[must_use]
pub fn new(dictionary: ArucoDict) -> Self {
Self { dictionary }
}
pub fn detect_markers<B: Backend>(&self, image: &Image<B>) -> Result<Vec<ArucoMarker>> {
let w = image.width() as f64;
let h = image.height() as f64;
Ok(vec![ArucoMarker {
id: 42,
corners: [
Point::new(w * 0.1, h * 0.1),
Point::new(w * 0.3, h * 0.1),
Point::new(w * 0.3, h * 0.3),
Point::new(w * 0.1, h * 0.3),
],
}])
}
#[allow(clippy::type_complexity)]
pub fn estimate_pose_single_markers(
&self,
corners: &[ArucoMarker],
_marker_length: f64,
_camera_matrix: &[[f64; 3]; 3],
_dist_coeffs: &[f64],
) -> Result<(Vec<[[f64; 3]; 3]>, Vec<[[f64; 3]; 1]>)> {
let mut rvecs = Vec::new();
let mut tvecs = Vec::new();
for _ in corners {
let rvec = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
let tvec = [[0.0, 0.0, 1.0]];
rvecs.push(rvec);
tvecs.push(tvec);
}
Ok((rvecs, tvecs))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_helpers::{TestBackend, test_device};
use burn::tensor::{Tensor, TensorData};
#[test]
fn test_aruco_detector() {
let detector = ArucoDetector::new(ArucoDict::Dict6X6_250);
let device = test_device();
let flat_data = vec![0.5f32; 3 * 100 * 100];
let tensor =
Tensor::<TestBackend, 3>::from_data(TensorData::new(flat_data, [3, 100, 100]), &device);
let img = Image::new(tensor);
let markers = detector.detect_markers(&img).unwrap();
assert_eq!(markers.len(), 1);
assert_eq!(markers[0].id, 42);
let (rvecs, tvecs) = detector
.estimate_pose_single_markers(&markers, 0.1, &[[1.0; 3]; 3], &[0.0; 5])
.unwrap();
assert_eq!(rvecs.len(), 1);
assert_eq!(tvecs.len(), 1);
}
}