chalkydri/
calibration.rs

1use std::{collections::HashMap, path::Path, sync::Arc, time::Duration};
2
3use aprilgrid::{
4    TagFamily,
5    detector::{DetectorParams, TagDetector},
6};
7use camera_intrinsic_calibration::{
8    board::{Board, create_default_6x6_board},
9    detected_points::{FeaturePoint, FrameFeature},
10    io::write_report,
11    types::{CalibParams, RvecTvec, ToRvecTvec},
12    util::*,
13    visualization::*,
14};
15use camera_intrinsic_model::{
16    self as model, CameraModel, GenericModel, OpenCVModel5, model_to_json,
17};
18use gstreamer_app::AppSink;
19use image::{ColorType, DynamicImage, GrayImage, RgbImage};
20
21use gstreamer::{
22    Buffer, Element, State,
23    glib::{WeakRef, object::ObjectExt},
24    prelude::{ElementExt, ElementExtManual},
25};
26use model::model_from_json;
27use tokio::{sync::watch, time::Instant};
28
29pub struct CalibratedModel {
30    model: GenericModel<f64>,
31}
32impl CalibratedModel {
33    pub fn new(calib: serde_json::Value) -> Self {
34        // Load the camera model
35        let model = serde_json::from_value(calib).unwrap();
36
37        Self { model }
38    }
39
40    pub const fn inner_model(&self) -> GenericModel<f64> {
41        self.model
42    }
43}
44
45const MIN_CORNERS: usize = 24;
46
47/// A camera calibrator
48pub struct Calibrator {
49    valve: WeakRef<Element>,
50    det: TagDetector,
51    board: Board,
52    frame_feats: Vec<FrameFeature>,
53    cam_model: GenericModel<f64>,
54    start: Instant,
55    rx: watch::Receiver<Option<Buffer>>,
56}
57impl Calibrator {
58    /// Initialize a new calibrator
59    pub fn new(valve: WeakRef<Element>, rx: watch::Receiver<Option<Buffer>>) -> Self {
60        Self {
61            det: TagDetector::new(&TagFamily::T36H11, None),
62            board: create_default_6x6_board(),
63            frame_feats: Vec::new(),
64            cam_model: GenericModel::OpenCVModel5(OpenCVModel5::zeros()),
65            start: Instant::now(),
66            rx,
67            valve,
68        }
69    }
70
71    /// Process a frame
72    pub fn step(&mut self) -> usize {
73        let valve = self.valve.upgrade().unwrap();
74        valve.set_property("drop", false);
75
76        if self.frame_feats.len() < 200 {
77            let mut frame_feat = None;
78            while frame_feat.is_none() {
79                if self.rx.has_changed().is_ok() && self.rx.has_changed().unwrap() {
80                    let val = self.rx.borrow_and_update().clone();
81                    debug!("got frame");
82                    //valve.set_property("drop", true);
83                    let img = DynamicImage::ImageRgb8(
84                        RgbImage::from_vec(
85                            1280,
86                            720,
87                            val.unwrap().into_mapped_buffer_readable().unwrap().to_vec(),
88                        )
89                        .unwrap(),
90                    );
91
92                    frame_feat =
93                        camera_intrinsic_calibration::data_loader::image_to_option_feature_frame(
94                            &self.det,
95                            &img,
96                            &create_default_6x6_board(),
97                            MIN_CORNERS,
98                            self.start.elapsed().as_nanos() as i64,
99                        );
100                }
101            }
102
103            self.frame_feats.push(frame_feat.unwrap());
104        }
105
106        self.frame_feats.len()
107    }
108
109    /// Calibrate
110    pub fn calibrate(&mut self) -> Option<GenericModel<f64>> {
111        let mut calib_res = None;
112
113        let valve = self.valve.upgrade().unwrap();
114        valve.set_property("drop", true);
115
116        for i in 0..5 {
117            calib_res = init_and_calibrate_one_camera(
118                0,
119                &[self
120                    .frame_feats
121                    .clone()
122                    .into_iter()
123                    .map(|f| Some(f))
124                    .collect()],
125                &self.cam_model,
126                &CalibParams {
127                    one_focal: false,
128                    fixed_focal: None,
129                    disabled_distortion_num: 0,
130                },
131                i > 1,
132            );
133            if calib_res.is_some() {
134                break;
135            }
136        }
137
138        self.frame_feats.clear();
139
140        if calib_res.is_none() {
141            error!("failed to calibrate camera");
142            None
143        } else {
144            Some(calib_res.unwrap().0)
145        }
146    }
147}