use std::collections::HashMap;
use std::f32::consts::FRAC_PI_2;
use nalgebra::{Point2, Vector2};
use crate::circular_stats::{
angle_to_bin, pick_two_peaks, smooth_circular_5, wrap_pi, PeakPickOptions,
};
use crate::global_step::{estimate_global_cell_size, GlobalStepParams};
use crate::square::alignment::GridTransform;
use crate::square::cleanup::{canonicalize_top_left, prune_to_main_component, sorted_grid_points};
use crate::square::detect::{
detect_square_grid, detect_square_grid_all, ExtensionStrategy, MultiComponentParams,
SquareGridParams,
};
use crate::square::grow::{Admit, GrowValidator, LabelledNeighbour};
use crate::square::seed::finder::SeedQuadValidator;
use crate::topological::AxisEstimate;
#[non_exhaustive]
#[derive(Clone, Debug)]
pub struct RegularGridParams {
pub pipeline: SquareGridParams,
pub canonicalize_top_left: bool,
pub prune_disconnected: bool,
}
impl Default for RegularGridParams {
fn default() -> Self {
Self {
pipeline: SquareGridParams::default(),
canonicalize_top_left: true,
prune_disconnected: true,
}
}
}
impl RegularGridParams {
pub fn new(pipeline: SquareGridParams) -> Self {
Self {
pipeline,
..Self::default()
}
}
pub fn with_extension(mut self, extension: ExtensionStrategy) -> Self {
self.pipeline.extension = extension;
self
}
pub fn with_canonicalize_top_left(mut self, on: bool) -> Self {
self.canonicalize_top_left = on;
self
}
pub fn with_prune_disconnected(mut self, on: bool) -> Self {
self.prune_disconnected = on;
self
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct DetectedGridPoint {
pub grid: (i32, i32),
pub position: Point2<f32>,
pub source_index: usize,
}
#[non_exhaustive]
#[derive(Clone, Debug, Default)]
pub struct RegularGridStats {
pub input_points: usize,
pub components_found: usize,
pub labelled_before_prune: usize,
pub pruned_disconnected: usize,
pub dropped_by_validation: usize,
pub canonicalized: bool,
}
#[derive(Clone, Debug)]
pub struct RegularGridDetection {
pub points: Vec<DetectedGridPoint>,
pub axis_i: Vector2<f32>,
pub axis_j: Vector2<f32>,
pub cell_size: f32,
pub stats: RegularGridStats,
}
#[non_exhaustive]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum RegularGridError {
TooFewPoints {
found: usize,
},
DegeneratePointCloud,
NoGridFound,
}
impl std::fmt::Display for RegularGridError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
RegularGridError::TooFewPoints { found } => write!(
f,
"too few points: {found} supplied, at least 4 required for a 2x2 seed"
),
RegularGridError::DegeneratePointCloud => write!(
f,
"degenerate point cloud: no square lattice could be inferred from the input"
),
RegularGridError::NoGridFound => {
write!(
f,
"no grid found: no valid 2x2 seed quad in the point cloud"
)
}
}
}
}
impl std::error::Error for RegularGridError {}
impl RegularGridDetection {
pub fn labelled_map(&self) -> HashMap<(i32, i32), usize> {
self.points
.iter()
.map(|p| (p.grid, p.source_index))
.collect()
}
}
pub fn detect_regular_grid(
points: &[Point2<f32>],
) -> Result<RegularGridDetection, RegularGridError> {
RegularGridDetector::default().detect(points)
}
#[derive(Clone, Debug, Default)]
pub struct RegularGridDetector {
pub params: RegularGridParams,
}
impl RegularGridDetector {
pub fn new(params: RegularGridParams) -> Self {
Self { params }
}
#[cfg_attr(
feature = "tracing",
tracing::instrument(
level = "info",
skip_all,
fields(num_points = points.len()),
)
)]
pub fn detect(&self, points: &[Point2<f32>]) -> Result<RegularGridDetection, RegularGridError> {
if points.len() < 4 {
return Err(RegularGridError::TooFewPoints {
found: points.len(),
});
}
let policy =
OpenRegularPolicy::new(points).ok_or(RegularGridError::DegeneratePointCloud)?;
let pipeline = self.pipeline_params();
let detection = detect_square_grid(points, &policy, &policy, &pipeline)
.ok_or(RegularGridError::NoGridFound)?;
let mut stats = RegularGridStats {
input_points: points.len(),
components_found: 1,
canonicalized: self.params.canonicalize_top_left,
dropped_by_validation: detection.stats.dropped_by_validation,
..Default::default()
};
let labelled = detection.labelled;
stats.labelled_before_prune = labelled.len();
let labelled = if self.params.prune_disconnected {
let pruned = prune_to_main_component(labelled);
stats.pruned_disconnected = stats.labelled_before_prune - pruned.len();
pruned
} else {
labelled
};
let (labelled, transform) = if self.params.canonicalize_top_left {
canonicalize_top_left(labelled, points)
} else {
(labelled, GridTransform::IDENTITY)
};
let (axis_i, axis_j) = transform_basis(detection.axis_i, detection.axis_j, transform);
Ok(build_detection(
&labelled,
points,
axis_i,
axis_j,
detection.cell_size,
stats,
))
}
#[cfg_attr(
feature = "tracing",
tracing::instrument(
level = "info",
skip_all,
fields(num_points = points.len()),
)
)]
pub fn detect_all(&self, points: &[Point2<f32>]) -> Vec<RegularGridDetection> {
if points.len() < 4 {
return Vec::new();
}
let Some(policy) = OpenRegularPolicy::new(points) else {
return Vec::new();
};
let pipeline = self.pipeline_params();
let raw = detect_square_grid_all(
points,
&policy,
&policy,
&pipeline,
&MultiComponentParams::default(),
);
let components_found = raw.len();
raw.into_iter()
.map(|detection| {
let mut stats = RegularGridStats {
input_points: points.len(),
components_found,
canonicalized: self.params.canonicalize_top_left,
dropped_by_validation: detection.stats.dropped_by_validation,
..Default::default()
};
let labelled = detection.labelled;
stats.labelled_before_prune = labelled.len();
let labelled = if self.params.prune_disconnected {
let pruned = prune_to_main_component(labelled);
stats.pruned_disconnected = stats.labelled_before_prune - pruned.len();
pruned
} else {
labelled
};
let (labelled, transform) = if self.params.canonicalize_top_left {
canonicalize_top_left(labelled, points)
} else {
(labelled, GridTransform::IDENTITY)
};
let (axis_i, axis_j) =
transform_basis(detection.axis_i, detection.axis_j, transform);
build_detection(
&labelled,
points,
axis_i,
axis_j,
detection.cell_size,
stats,
)
})
.collect()
}
fn pipeline_params(&self) -> SquareGridParams {
self.params.pipeline.clone()
}
}
fn build_detection(
labelled: &HashMap<(i32, i32), usize>,
points: &[Point2<f32>],
axis_i: Vector2<f32>,
axis_j: Vector2<f32>,
cell_size: f32,
stats: RegularGridStats,
) -> RegularGridDetection {
let detected: Vec<DetectedGridPoint> = sorted_grid_points(labelled)
.into_iter()
.map(|(grid, idx)| DetectedGridPoint {
grid,
position: points[idx],
source_index: idx,
})
.collect();
RegularGridDetection {
points: detected,
axis_i,
axis_j,
cell_size,
stats,
}
}
fn transform_basis(
axis_i: Vector2<f32>,
axis_j: Vector2<f32>,
transform: GridTransform,
) -> (Vector2<f32>, Vector2<f32>) {
let inv = transform.inverse().unwrap_or(GridTransform::IDENTITY);
let gi = inv.apply(1, 0);
let gj = inv.apply(0, 1);
let new_i = axis_i * gi.i as f32 + axis_j * gi.j as f32;
let new_j = axis_i * gj.i as f32 + axis_j * gj.j as f32;
let norm_i = new_i.norm().max(1e-6);
let norm_j = new_j.norm().max(1e-6);
(new_i / norm_i, new_j / norm_j)
}
struct OpenRegularPolicy {
positions: Vec<Point2<f32>>,
axes: [AxisEstimate; 2],
}
impl OpenRegularPolicy {
fn new(points: &[Point2<f32>]) -> Option<Self> {
let axes = estimate_grid_axes(points)?;
Some(Self {
positions: points.to_vec(),
axes,
})
}
}
impl SeedQuadValidator for OpenRegularPolicy {
fn position(&self, idx: usize) -> Point2<f32> {
self.positions[idx]
}
fn axes(&self, _idx: usize) -> [AxisEstimate; 2] {
self.axes
}
fn a_candidates(&self) -> Vec<usize> {
(0..self.positions.len()).collect()
}
fn bc_candidates(&self) -> Vec<usize> {
(0..self.positions.len()).collect()
}
}
impl GrowValidator for OpenRegularPolicy {
fn is_eligible(&self, _idx: usize) -> bool {
true
}
fn required_label_at(&self, _i: i32, _j: i32) -> Option<u8> {
None
}
fn label_of(&self, _idx: usize) -> Option<u8> {
None
}
fn accept_candidate(
&self,
_idx: usize,
_at: (i32, i32),
_prediction: Point2<f32>,
_neighbours: &[LabelledNeighbour],
) -> Admit {
Admit::Accept
}
}
fn estimate_grid_axes(points: &[Point2<f32>]) -> Option<[AxisEstimate; 2]> {
use kiddo::{KdTree, SquaredEuclidean};
if points.len() < 4 {
return None;
}
estimate_global_cell_size(points, &GlobalStepParams::<f32>::default())?;
let mut tree: KdTree<f32, 2> = KdTree::new();
for (idx, p) in points.iter().enumerate() {
tree.add(&[p.x, p.y], idx as u64);
}
const N_BINS: usize = 180;
let mut hist = vec![0.0_f32; N_BINS];
let mut total = 0.0_f32;
for (i, p) in points.iter().enumerate() {
let hits = tree.nearest_n::<SquaredEuclidean>(&[p.x, p.y], 5);
for hit in hits {
let j = hit.item as usize;
if j == i {
continue;
}
let q = points[j];
let off = Vector2::new(q.x - p.x, q.y - p.y);
let len = off.norm();
if len < 1e-3 {
continue;
}
let ang = wrap_pi(off.y.atan2(off.x));
let bin = angle_to_bin(ang, N_BINS);
hist[bin] += len;
total += len;
}
}
if total <= 0.0 {
return None;
}
let smoothed = smooth_circular_5(&hist);
let opts = PeakPickOptions::new(0.05, 30.0_f32.to_radians());
match pick_two_peaks(&smoothed, total, &opts) {
Some((t0, t1)) => {
let (lo, hi) = if t0 <= t1 { (t0, t1) } else { (t1, t0) };
Some([AxisEstimate::from_angle(lo), AxisEstimate::from_angle(hi)])
}
None => {
let peak = smoothed
.iter()
.enumerate()
.max_by(|a, b| a.1.total_cmp(b.1))
.map(|(b, _)| b)?;
let theta = wrap_pi(crate::circular_stats::bin_to_angle(peak, N_BINS));
let other = wrap_pi(theta + FRAC_PI_2);
let (lo, hi) = if theta <= other {
(theta, other)
} else {
(other, theta)
};
Some([AxisEstimate::from_angle(lo), AxisEstimate::from_angle(hi)])
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use nalgebra::Matrix3;
fn axis_aligned_grid(rows: i32, cols: i32, s: f32) -> Vec<Point2<f32>> {
let mut out = Vec::new();
for j in 0..rows {
for i in 0..cols {
out.push(Point2::new(i as f32 * s + 40.0, j as f32 * s + 40.0));
}
}
out
}
#[test]
fn detects_clean_axis_aligned_grid() {
let pts = axis_aligned_grid(6, 6, 25.0);
let grid = detect_regular_grid(&pts).expect("clean grid detects");
assert_eq!(grid.points.len(), 36);
assert_eq!(grid.stats.input_points, 36);
}
#[test]
fn returns_err_on_collinear_cloud() {
let pts: Vec<Point2<f32>> = (0..6).map(|i| Point2::new(i as f32 * 10.0, 0.0)).collect();
assert!(detect_regular_grid(&pts).is_err());
}
#[test]
fn estimate_grid_axes_recovers_rotation() {
let theta = 30.0_f32.to_radians();
let (c, s) = (theta.cos(), theta.sin());
let mut pts = Vec::new();
for j in 0..5 {
for i in 0..5 {
let (x, y) = (i as f32 * 20.0, j as f32 * 20.0);
pts.push(Point2::new(x * c - y * s + 100.0, x * s + y * c + 100.0));
}
}
let axes = estimate_grid_axes(&pts).expect("axes");
let near = axes
.iter()
.any(|a| crate::circular_stats::angular_dist_pi(a.angle, theta) < 0.15);
assert!(near, "expected an axis near 30°, got {axes:?}");
}
#[test]
fn perspective_warped_grid_is_recovered() {
let h = Matrix3::new(30.0_f32, 3.0, 50.0, 1.5, 30.0, 50.0, 2e-4, 1e-4, 1.0);
let mut pts = Vec::new();
for j in 0..7 {
for i in 0..7 {
let (x, y) = (i as f32, j as f32);
let w = h[(2, 0)] * x + h[(2, 1)] * y + h[(2, 2)];
let xp = (h[(0, 0)] * x + h[(0, 1)] * y + h[(0, 2)]) / w;
let yp = (h[(1, 0)] * x + h[(1, 1)] * y + h[(1, 2)]) / w;
pts.push(Point2::new(xp, yp));
}
}
let grid = detect_regular_grid(&pts).expect("warped grid detects");
assert!(grid.points.len() >= 40, "got {}", grid.points.len());
}
}