calib_targets_chessboard/
mesh_warp.rs

1use calib_targets_core::{
2    estimate_homography_rect_to_img, sample_bilinear_u8, GrayImage, GrayImageView, GridCoords,
3    Homography, LabeledCorner,
4};
5use std::collections::HashMap;
6
7use nalgebra::Point2;
8
9#[derive(thiserror::Error, Debug)]
10pub enum MeshWarpError {
11    #[error("not enough labeled corners with grid coords")]
12    NotEnoughLabeledCorners,
13    #[error("no valid grid cells found (need 2x2 corners at least)")]
14    NoValidCells,
15    #[error("homography estimation failed for at least one cell")]
16    HomographyFailed,
17}
18
19// ---- Mesh warp rectification ----
20#[derive(Clone, Debug)]
21pub struct RectifiedMeshView {
22    pub rect: GrayImage,
23
24    // Rectified grid cell layout (cell indices, not corner indices):
25    // cell (ci, cj) corresponds to corner indices:
26    // (min_i + ci, min_j + cj) .. (min_i + ci + 1, min_j + cj + 1)
27    pub min_i: i32,
28    pub min_j: i32,
29    pub cells_x: usize,
30    pub cells_y: usize,
31
32    pub px_per_square: f32,
33
34    // How many cells were valid (had all 4 corners)
35    pub valid_cells: usize,
36
37    // Per-cell homographies (cell-local rect -> image), in row-major (cj * cells_x + ci).
38    // This is intentionally kept private; use `cell_rect_to_img`/`rect_to_img` accessors.
39    cells: Vec<Cell>,
40}
41
42// Internal cell storage
43#[derive(Clone, Copy, Debug)]
44struct Cell {
45    h_img_from_cellrect: Homography,
46    valid: bool,
47}
48
49impl RectifiedMeshView {
50    /// Map a point in **global rectified pixel coordinates** into the original image,
51    /// using the homography of the cell that contains it.
52    ///
53    /// Returns `None` if the point lies outside the rectified image or the cell is invalid.
54    pub fn rect_to_img(&self, p_rect: Point2<f32>) -> Option<Point2<f32>> {
55        let s = self.px_per_square;
56        if s <= 0.0 {
57            return None;
58        }
59
60        let ci = (p_rect.x / s).floor() as i32;
61        let cj = (p_rect.y / s).floor() as i32;
62        if ci < 0 || cj < 0 || ci >= self.cells_x as i32 || cj >= self.cells_y as i32 {
63            return None;
64        }
65
66        let x_local = p_rect.x - (ci as f32) * s;
67        let y_local = p_rect.y - (cj as f32) * s;
68        self.cell_rect_to_img(ci as usize, cj as usize, Point2::new(x_local, y_local))
69    }
70
71    /// Map a point in **cell-local rectified pixel coordinates** into the original image.
72    ///
73    /// - `ci`, `cj`: cell indices in `0..cells_x × 0..cells_y`
74    /// - `p_cell`: point in `[0..px_per_square]²` (cell-local)
75    pub fn cell_rect_to_img(
76        &self,
77        ci: usize,
78        cj: usize,
79        p_cell: Point2<f32>,
80    ) -> Option<Point2<f32>> {
81        let idx = cj.checked_mul(self.cells_x)?.checked_add(ci)?;
82        let cell = *self.cells.get(idx)?;
83        if !cell.valid {
84            return None;
85        }
86        Some(cell.h_img_from_cellrect.apply(p_cell))
87    }
88
89    /// Convenience: map the four corners of a cell into image coordinates.
90    pub fn cell_corners_img(&self, ci: usize, cj: usize) -> Option<[Point2<f32>; 4]> {
91        let s = self.px_per_square;
92        let pts = [
93            Point2::new(0.0, 0.0),
94            Point2::new(s, 0.0),
95            Point2::new(s, s),
96            Point2::new(0.0, s),
97        ];
98        Some([
99            self.cell_rect_to_img(ci, cj, pts[0])?,
100            self.cell_rect_to_img(ci, cj, pts[1])?,
101            self.cell_rect_to_img(ci, cj, pts[2])?,
102            self.cell_rect_to_img(ci, cj, pts[3])?,
103        ])
104    }
105}
106
107fn build_corners_to_pix_map(
108    corners: &[LabeledCorner],
109    inliers: &[usize],
110) -> HashMap<GridCoords, Point2<f32>> {
111    let mut map: HashMap<GridCoords, Point2<f32>> = HashMap::new();
112    for &idx in inliers {
113        if let Some(c) = corners.get(idx) {
114            if let Some(g) = c.grid {
115                map.insert(g, c.position);
116            }
117        }
118    }
119    map
120}
121
122/// Build a rectified “board view” by piecewise homographies per grid cell.
123/// This is robust to lens distortion because it does not assume a single global H.
124///
125/// - `corners`: your detection.corners
126/// - `inliers`: indices into `corners` that you trust
127/// - `px_per_square`: rectified pixels per chess square (recommend 60..120, preferably an integer)
128pub fn rectify_mesh_from_grid(
129    src: &GrayImageView<'_>,
130    corners: &[LabeledCorner],
131    inliers: &[usize],
132    px_per_square: f32,
133) -> Result<RectifiedMeshView, MeshWarpError> {
134    // 1) Build map: (i,j) -> image point
135    let map = build_corners_to_pix_map(corners, inliers);
136    if map.len() < 4 {
137        return Err(MeshWarpError::NotEnoughLabeledCorners);
138    }
139
140    // 2) Determine bounding box in corner-index space
141    let (mut min_i, mut min_j) = (i32::MAX, i32::MAX);
142    let (mut max_i, mut max_j) = (i32::MIN, i32::MIN);
143    for g in map.keys() {
144        min_i = min_i.min(g.i);
145        min_j = min_j.min(g.j);
146        max_i = max_i.max(g.i);
147        max_j = max_j.max(g.j);
148    }
149
150    // Need at least 2x2 corners => at least 1x1 cell
151    if max_i - min_i < 1 || max_j - min_j < 1 {
152        return Err(MeshWarpError::NoValidCells);
153    }
154
155    let cells_x = (max_i - min_i) as usize; // number of squares horizontally
156    let cells_y = (max_j - min_j) as usize; // number of squares vertically
157
158    // Output size: exactly cells * px_per_square (floor to ensure stable indexing)
159    let out_w = ((cells_x as f32) * px_per_square).floor().max(1.0) as usize;
160    let out_h = ((cells_y as f32) * px_per_square).floor().max(1.0) as usize;
161
162    // 3) Precompute per-cell homographies (cell-rect -> image)
163    let mut cells = vec![
164        Cell {
165            h_img_from_cellrect: Homography::zero(),
166            valid: false
167        };
168        cells_x * cells_y
169    ];
170
171    let s = px_per_square;
172
173    // Rectified cell corner coordinates (in pixels within the cell)
174    let cell_rect = [
175        Point2::new(0.0, 0.0),
176        Point2::new(s, 0.0),
177        Point2::new(0.0, s),
178        Point2::new(s, s),
179    ];
180
181    let mut valid_cells = 0usize;
182
183    for cj in 0..cells_y {
184        for ci in 0..cells_x {
185            let i0 = min_i + ci as i32;
186            let j0 = min_j + cj as i32;
187
188            let g00 = GridCoords { i: i0, j: j0 };
189            let g10 = GridCoords { i: i0 + 1, j: j0 };
190            let g01 = GridCoords { i: i0, j: j0 + 1 };
191            let g11 = GridCoords {
192                i: i0 + 1,
193                j: j0 + 1,
194            };
195
196            let Some(p00) = map.get(&g00).copied() else {
197                continue;
198            };
199            let Some(p10) = map.get(&g10).copied() else {
200                continue;
201            };
202            let Some(p01) = map.get(&g01).copied() else {
203                continue;
204            };
205            let Some(p11) = map.get(&g11).copied() else {
206                continue;
207            };
208
209            let img_quad = [p00, p10, p01, p11];
210
211            let Some(h) = estimate_homography_rect_to_img(&cell_rect, &img_quad) else {
212                return Err(MeshWarpError::HomographyFailed);
213            };
214
215            let idx = cj * cells_x + ci;
216            cells[idx] = Cell {
217                h_img_from_cellrect: h,
218                valid: true,
219            };
220            valid_cells += 1;
221        }
222    }
223
224    if valid_cells == 0 {
225        return Err(MeshWarpError::NoValidCells);
226    }
227
228    // 4) Warp: each output pixel chooses its cell and uses that cell homography
229    let mut out = vec![0u8; out_w * out_h];
230
231    for y in 0..out_h {
232        let yf = y as f32 + 0.5;
233        let cj = (yf / s).floor() as i32;
234        if cj < 0 || cj >= cells_y as i32 {
235            continue;
236        }
237        let cj_u = cj as usize;
238        let y_local = yf - (cj as f32) * s;
239
240        for x in 0..out_w {
241            let xf = x as f32 + 0.5;
242            let ci = (xf / s).floor() as i32;
243            if ci < 0 || ci >= cells_x as i32 {
244                continue;
245            }
246            let ci_u = ci as usize;
247            let x_local = xf - (ci as f32) * s;
248
249            let cell = cells[cj_u * cells_x + ci_u];
250            if cell.valid {
251                let p_cell = Point2::new(x_local, y_local);
252                let p_img = cell.h_img_from_cellrect.apply(p_cell);
253                out[y * out_w + x] = sample_bilinear_u8(src, p_img.x, p_img.y);
254            }
255        }
256    }
257
258    Ok(RectifiedMeshView {
259        rect: GrayImage {
260            width: out_w,
261            height: out_h,
262            data: out,
263        },
264        min_i,
265        min_j,
266        cells_x,
267        cells_y,
268        px_per_square,
269        valid_cells,
270        cells,
271    })
272}