#[cfg(feature = "ml-refiner")]
use crate::ml_refiner;
use crate::{ChessConfig, ChessParams, DetectorMode};
use box_image_pyramid::{build_pyramid, PyramidBuffers, PyramidParams};
use chess_corners_core::descriptor::{corners_to_descriptors_with_method, Corner};
use chess_corners_core::detect::{detect_corners_from_response_with_refiner, merge_corners_simple};
use chess_corners_core::response::{chess_response_u8, chess_response_u8_patch, Roi};
use chess_corners_core::{
detect_corners_from_radon, radon_response_u8, CornerDescriptor, CornerRefiner, RadonBuffers,
};
use chess_corners_core::{ImageView, Refiner, RefinerKind, ResponseMap};
fn to_pyramid_view(v: ImageView<'_>) -> box_image_pyramid::ImageView<'_> {
box_image_pyramid::ImageView::new(v.width, v.height, v.data).unwrap()
}
#[cfg(feature = "rayon")]
use rayon::prelude::*;
#[cfg(feature = "tracing")]
use tracing::{info_span, instrument};
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct CoarseToFineParams {
pub pyramid: PyramidParams,
pub refinement_radius: u32,
pub merge_radius: f32,
}
impl Default for CoarseToFineParams {
fn default() -> Self {
Self {
pyramid: PyramidParams::default(),
refinement_radius: 3,
merge_radius: 3.0,
}
}
}
impl CoarseToFineParams {
pub fn new() -> Self {
Self::default()
}
}
#[cfg(feature = "ml-refiner")]
fn detect_with_ml_refiner(
resp: &ResponseMap,
params: &ChessParams,
image: Option<ImageView<'_>>,
ml_state: &mut ml_refiner::MlRefinerState,
) -> Vec<Corner> {
ml_refiner::detect_corners_with_ml(resp, params, image, ml_state)
}
fn detect_with_refiner_kind(
resp: &ResponseMap,
params: &ChessParams,
image: Option<ImageView<'_>>,
refiner_kind: &RefinerKind,
) -> Vec<Corner> {
let mut refiner = Refiner::from_kind(refiner_kind.clone());
detect_corners_from_response_with_refiner(resp, params, image, &mut refiner)
}
fn refiner_radius(refiner_kind: &RefinerKind) -> i32 {
Refiner::from_kind(refiner_kind.clone()).radius()
}
struct RoiContext {
inv_scale: f32,
border: i32,
safe_margin: i32,
roi_r: i32,
base_w_i: i32,
base_h_i: i32,
}
impl RoiContext {
fn compute_roi(&self, c: &Corner) -> Option<(i32, i32, i32, i32)> {
let cx = (c.x * self.inv_scale).round() as i32;
let cy = (c.y * self.inv_scale).round() as i32;
if cx < self.safe_margin
|| cy < self.safe_margin
|| cx >= self.base_w_i - self.safe_margin
|| cy >= self.base_h_i - self.safe_margin
{
return None;
}
let mut x0 = cx - self.roi_r;
let mut y0 = cy - self.roi_r;
let mut x1 = cx + self.roi_r + 1;
let mut y1 = cy + self.roi_r + 1;
let min_xy = self.border;
let max_x = self.base_w_i - self.border;
let max_y = self.base_h_i - self.border;
if x0 < min_xy {
x0 = min_xy;
}
if y0 < min_xy {
y0 = min_xy;
}
if x1 > max_x {
x1 = max_x;
}
if y1 > max_y {
y1 = max_y;
}
if x1 - x0 <= 2 * self.border || y1 - y0 <= 2 * self.border {
return None;
}
Some((x0, y0, x1, y1))
}
}
fn refine_seed_in_roi(
base: ImageView<'_>,
params: &ChessParams,
roi_bounds: (i32, i32, i32, i32),
mut detect: impl FnMut(&ResponseMap, &ChessParams, Option<ImageView<'_>>) -> Vec<Corner>,
) -> Option<Vec<Corner>> {
let (x0, y0, x1, y1) = roi_bounds;
let base_w = base.width;
let base_h = base.height;
let roi = Roi::new(x0 as usize, y0 as usize, x1 as usize, y1 as usize)?;
let patch_resp = chess_response_u8_patch(base.data, base_w, base_h, params, roi);
if patch_resp.width() == 0 || patch_resp.height() == 0 {
return None;
}
let refine_view = ImageView::with_origin(base_w, base_h, base.data, [x0, y0])
.expect("base image dimensions must match buffer length");
let mut patch_corners = detect(&patch_resp, params, Some(refine_view));
for pc in &mut patch_corners {
pc.x += x0 as f32;
pc.y += y0 as f32;
}
if patch_corners.is_empty() {
None
} else {
Some(patch_corners)
}
}
fn merge_and_describe(
base: ImageView<'_>,
params: &ChessParams,
merge_radius: f32,
refined: &mut Vec<Corner>,
) -> Vec<CornerDescriptor> {
#[cfg(feature = "tracing")]
let merge_span = info_span!(
"merge",
merge_radius = merge_radius,
candidates = refined.len()
)
.entered();
let merged = merge_corners_simple(refined, merge_radius);
#[cfg(feature = "tracing")]
drop(merge_span);
corners_to_descriptors_with_method(
base.data,
base.width,
base.height,
params.descriptor_ring_radius(),
merged,
params.orientation_method,
)
}
#[cfg(feature = "ml-refiner")]
fn coarse_to_fine_with<R>(
base: ImageView<'_>,
cfg: &ChessConfig,
buffers: &mut PyramidBuffers,
coarse_detect: &RefinerKind,
refine_border: i32,
detect_fn: &mut R,
) -> Vec<CornerDescriptor>
where
R: FnMut(&ResponseMap, &ChessParams, Option<ImageView<'_>>) -> Vec<Corner>,
{
let params = cfg.to_chess_params();
let cf = cfg.to_coarse_to_fine_params();
let pyramid = build_pyramid(to_pyramid_view(base), &cf.pyramid, buffers);
if pyramid.levels.is_empty() {
return Vec::new();
}
if pyramid.levels.len() == 1 {
let lvl = &pyramid.levels[0];
let resp = chess_response_u8(lvl.img.data, lvl.img.width, lvl.img.height, ¶ms);
let view = ImageView::from_u8_slice(lvl.img.width, lvl.img.height, lvl.img.data)
.expect("image dimensions must match buffer length");
let mut raw = detect_fn(&resp, ¶ms, Some(view));
let merged = merge_corners_simple(&mut raw, cf.merge_radius);
return corners_to_descriptors_with_method(
lvl.img.data,
lvl.img.width,
lvl.img.height,
params.descriptor_ring_radius(),
merged,
params.orientation_method,
);
}
let coarse_lvl = pyramid.levels.last().unwrap();
let coarse_w = coarse_lvl.img.width;
let coarse_h = coarse_lvl.img.height;
#[cfg(feature = "tracing")]
let coarse_span = info_span!("coarse_detect", w = coarse_w, h = coarse_h).entered();
let coarse_resp = chess_response_u8(coarse_lvl.img.data, coarse_w, coarse_h, ¶ms);
let coarse_view = ImageView::from_u8_slice(coarse_w, coarse_h, coarse_lvl.img.data).unwrap();
let coarse_corners =
detect_with_refiner_kind(&coarse_resp, ¶ms, Some(coarse_view), coarse_detect);
#[cfg(feature = "tracing")]
drop(coarse_span);
if coarse_corners.is_empty() {
return Vec::new();
}
let roi_ctx = make_roi_context(base, coarse_lvl.scale, ¶ms, refine_border, &cf);
#[cfg(feature = "tracing")]
let refine_span = info_span!(
"refine",
seeds = coarse_corners.len(),
roi_r = roi_ctx.roi_r
)
.entered();
let mut refined: Vec<Corner> = coarse_corners
.into_iter()
.filter_map(|c| {
let roi_bounds = roi_ctx.compute_roi(&c)?;
refine_seed_in_roi(base, ¶ms, roi_bounds, &mut *detect_fn)
})
.flatten()
.collect();
#[cfg(feature = "tracing")]
drop(refine_span);
merge_and_describe(base, ¶ms, cf.merge_radius, &mut refined)
}
pub fn find_chess_corners_buff(
base: ImageView<'_>,
cfg: &ChessConfig,
buffers: &mut PyramidBuffers,
) -> Vec<CornerDescriptor> {
let refiner = cfg.refiner.to_refiner_kind();
find_chess_corners_buff_with_refiner(base, cfg, buffers, &refiner)
}
pub fn find_chess_corners_buff_with_refiner(
base: ImageView<'_>,
cfg: &ChessConfig,
buffers: &mut PyramidBuffers,
refiner: &RefinerKind,
) -> Vec<CornerDescriptor> {
if matches!(cfg.detector_mode, DetectorMode::Radon) {
return detect_with_radon(base, cfg);
}
let params = cfg.to_chess_params();
let cf = cfg.to_coarse_to_fine_params();
let border = refiner_radius(refiner);
let pyramid = build_pyramid(to_pyramid_view(base), &cf.pyramid, buffers);
if pyramid.levels.is_empty() {
return Vec::new();
}
if pyramid.levels.len() == 1 {
let lvl = &pyramid.levels[0];
let resp = chess_response_u8(lvl.img.data, lvl.img.width, lvl.img.height, ¶ms);
let refine_view = ImageView::from_u8_slice(lvl.img.width, lvl.img.height, lvl.img.data)
.expect("image dimensions must match buffer length");
let mut raw = detect_with_refiner_kind(&resp, ¶ms, Some(refine_view), refiner);
let merged = merge_corners_simple(&mut raw, cf.merge_radius);
return corners_to_descriptors_with_method(
lvl.img.data,
lvl.img.width,
lvl.img.height,
params.descriptor_ring_radius(),
merged,
params.orientation_method,
);
}
let coarse_lvl = pyramid.levels.last().unwrap();
let coarse_w = coarse_lvl.img.width;
let coarse_h = coarse_lvl.img.height;
#[cfg(feature = "tracing")]
let coarse_span = info_span!("coarse_detect", w = coarse_w, h = coarse_h).entered();
let coarse_resp = chess_response_u8(coarse_lvl.img.data, coarse_w, coarse_h, ¶ms);
let coarse_view = ImageView::from_u8_slice(coarse_w, coarse_h, coarse_lvl.img.data).unwrap();
let coarse_corners =
detect_with_refiner_kind(&coarse_resp, ¶ms, Some(coarse_view), refiner);
#[cfg(feature = "tracing")]
drop(coarse_span);
if coarse_corners.is_empty() {
return Vec::new();
}
let roi_ctx = make_roi_context(base, coarse_lvl.scale, ¶ms, border, &cf);
#[cfg(feature = "tracing")]
let refine_span = info_span!(
"refine",
seeds = coarse_corners.len(),
roi_r = roi_ctx.roi_r
)
.entered();
let refine_one = |c: Corner| -> Option<Vec<Corner>> {
let roi_bounds = roi_ctx.compute_roi(&c)?;
refine_seed_in_roi(base, ¶ms, roi_bounds, |resp, p, image| {
detect_with_refiner_kind(resp, p, image, refiner)
})
};
#[cfg(feature = "rayon")]
let mut refined: Vec<Corner> = coarse_corners
.into_par_iter()
.filter_map(refine_one)
.flatten()
.collect();
#[cfg(not(feature = "rayon"))]
let mut refined: Vec<Corner> = coarse_corners
.into_iter()
.filter_map(refine_one)
.flatten()
.collect();
#[cfg(feature = "tracing")]
drop(refine_span);
merge_and_describe(base, ¶ms, cf.merge_radius, &mut refined)
}
fn detect_with_radon(base: ImageView<'_>, cfg: &ChessConfig) -> Vec<CornerDescriptor> {
#[cfg(feature = "tracing")]
let span = info_span!(
"radon_detect",
w = base.width,
h = base.height,
upsample = cfg.radon_detector.image_upsample,
)
.entered();
let mut rb = RadonBuffers::new();
let resp = radon_response_u8(
base.data,
base.width,
base.height,
&cfg.radon_detector,
&mut rb,
);
let corners = detect_corners_from_radon(&resp, &cfg.radon_detector);
if corners.is_empty() {
#[cfg(feature = "tracing")]
drop(span);
return Vec::new();
}
let params = cfg.to_chess_params();
let mut merged = corners;
let merged = merge_corners_simple(&mut merged, cfg.merge_radius);
let out = corners_to_descriptors_with_method(
base.data,
base.width,
base.height,
params.descriptor_ring_radius(),
merged,
params.orientation_method,
);
#[cfg(feature = "tracing")]
drop(span);
out
}
#[cfg(feature = "ml-refiner")]
pub fn find_chess_corners_buff_with_ml(
base: ImageView<'_>,
cfg: &ChessConfig,
buffers: &mut PyramidBuffers,
) -> Vec<CornerDescriptor> {
let ml_params = ml_refiner::MlRefinerParams::default();
let fallback_refiner = cfg.refiner.to_refiner_kind();
let mut ml_state = ml_refiner::MlRefinerState::new(&ml_params, &fallback_refiner);
find_chess_corners_buff_with_ml_state(base, cfg, buffers, &ml_params, &mut ml_state)
}
#[cfg(feature = "ml-refiner")]
fn find_chess_corners_buff_with_ml_state(
base: ImageView<'_>,
cfg: &ChessConfig,
buffers: &mut PyramidBuffers,
ml: &ml_refiner::MlRefinerParams,
ml_state: &mut ml_refiner::MlRefinerState,
) -> Vec<CornerDescriptor> {
if matches!(cfg.detector_mode, DetectorMode::Radon) {
return detect_with_radon(base, cfg);
}
let params = cfg.to_chess_params();
let ml_border = ml_refiner::patch_radius(ml);
coarse_to_fine_with(
base,
cfg,
buffers,
¶ms.refiner.clone(),
ml_border,
&mut |resp, p, image| detect_with_ml_refiner(resp, p, image, ml_state),
)
}
fn make_roi_context(
base: ImageView<'_>,
coarse_scale: f32,
params: &ChessParams,
refine_border: i32,
cf: &CoarseToFineParams,
) -> RoiContext {
let ring_r = params.ring_radius() as i32;
let nms_r = params.nms_radius as i32;
let border = (ring_r + nms_r + refine_border).max(0);
let safe_margin = border + 1;
let roi_r_base = (cf.refinement_radius as f32 / coarse_scale).ceil() as i32;
let min_roi_r = border + 2;
RoiContext {
inv_scale: 1.0 / coarse_scale,
border,
safe_margin,
roi_r: roi_r_base.max(min_roi_r),
base_w_i: base.width as i32,
base_h_i: base.height as i32,
}
}
#[must_use]
#[cfg_attr(
feature = "tracing",
instrument(
level = "info",
skip(base, cfg),
fields(levels = cfg.pyramid_levels, min_size = cfg.pyramid_min_size)
)
)]
pub fn find_chess_corners(base: ImageView<'_>, cfg: &ChessConfig) -> Vec<CornerDescriptor> {
let refiner = cfg.refiner.to_refiner_kind();
find_chess_corners_with_refiner(base, cfg, &refiner)
}
#[must_use]
pub fn find_chess_corners_with_refiner(
base: ImageView<'_>,
cfg: &ChessConfig,
refiner: &RefinerKind,
) -> Vec<CornerDescriptor> {
let mut buffers = PyramidBuffers::with_capacity(cfg.pyramid_levels);
find_chess_corners_buff_with_refiner(base, cfg, &mut buffers, refiner)
}
#[cfg(feature = "ml-refiner")]
#[must_use]
pub fn find_chess_corners_with_ml(base: ImageView<'_>, cfg: &ChessConfig) -> Vec<CornerDescriptor> {
let mut buffers = PyramidBuffers::with_capacity(cfg.pyramid_levels);
find_chess_corners_buff_with_ml(base, cfg, &mut buffers)
}
#[cfg(test)]
mod tests {
use super::*;
use box_image_pyramid::ImageBuffer;
#[test]
fn default_coarse_to_fine_config_is_single_scale() {
let cfg = CoarseToFineParams::default();
assert_eq!(cfg.pyramid.num_levels, 1);
assert_eq!(cfg.pyramid.min_size, 128);
assert_eq!(cfg.refinement_radius, 3);
assert_eq!(cfg.merge_radius, 3.0);
}
#[test]
fn chess_config_multiscale_preset_has_expected_pyramid() {
let cfg = ChessConfig::multiscale();
assert_eq!(cfg.pyramid_levels, 3);
assert_eq!(cfg.pyramid_min_size, 128);
assert_eq!(cfg.refinement_radius, 3);
assert_eq!(cfg.merge_radius, 3.0);
}
#[test]
fn coarse_to_fine_trace_reports_timings() {
let buf = ImageBuffer::new(32, 32);
let view = ImageView::from_u8_slice(buf.width, buf.height, &buf.data)
.expect("dimensions must match");
let cfg = ChessConfig::default();
let corners = find_chess_corners(view, &cfg);
assert!(corners.is_empty());
}
}