use box_image_pyramid::PyramidBuffers;
use chess_corners_core::{ChessBuffers, RadonBuffers};
#[cfg(feature = "ml-refiner")]
use crate::ml_refiner;
use crate::multiscale;
use crate::upscale::{self, UpscaleBuffers};
use crate::{low_level::ImageView, ChessError, CornerDescriptor, DetectorConfig};
pub struct Detector {
cfg: DetectorConfig,
pyramid: PyramidBuffers,
chess_buffers: ChessBuffers,
radon_buffers: RadonBuffers,
upscale: UpscaleBuffers,
#[cfg(feature = "ml-refiner")]
ml_state: Option<ml_refiner::MlRefinerState>,
#[cfg(feature = "ml-refiner")]
ml_params: ml_refiner::MlRefinerParams,
}
impl Detector {
pub fn new(cfg: DetectorConfig) -> Result<Self, ChessError> {
cfg.upscale.validate()?;
Ok(Self {
cfg,
pyramid: PyramidBuffers::default(),
chess_buffers: ChessBuffers::default(),
radon_buffers: RadonBuffers::default(),
upscale: UpscaleBuffers::new(),
#[cfg(feature = "ml-refiner")]
ml_state: None,
#[cfg(feature = "ml-refiner")]
ml_params: ml_refiner::MlRefinerParams::default(),
})
}
pub fn with_default() -> Self {
Self::new(DetectorConfig::default()).expect("default DetectorConfig is always valid")
}
pub fn config(&self) -> &DetectorConfig {
&self.cfg
}
pub fn set_config(&mut self, cfg: DetectorConfig) -> Result<(), ChessError> {
cfg.upscale.validate()?;
self.cfg = cfg;
#[cfg(feature = "ml-refiner")]
{
self.ml_state = None;
}
Ok(())
}
pub fn config_mut(&mut self) -> &mut DetectorConfig {
#[cfg(feature = "ml-refiner")]
{
self.ml_state = None;
}
&mut self.cfg
}
pub fn detect_u8(
&mut self,
img: &[u8],
width: u32,
height: u32,
) -> Result<Vec<CornerDescriptor>, ChessError> {
let src_w = width as usize;
let src_h = height as usize;
let expected = src_w * src_h;
if img.len() != expected {
return Err(ChessError::DimensionMismatch {
expected,
actual: img.len(),
});
}
let factor = self.cfg.upscale.effective_factor();
if factor <= 1 {
let view =
ImageView::from_u8_slice(src_w, src_h, img).expect("dimensions were checked above");
return Ok(Self::detect_view_inner(
&self.cfg,
&mut self.pyramid,
&mut self.chess_buffers,
&mut self.radon_buffers,
#[cfg(feature = "ml-refiner")]
&mut self.ml_state,
#[cfg(feature = "ml-refiner")]
&self.ml_params,
view,
));
}
let upscaled = upscale::upscale_bilinear_u8(img, src_w, src_h, factor, &mut self.upscale)?;
let mut corners = Self::detect_view_inner(
&self.cfg,
&mut self.pyramid,
&mut self.chess_buffers,
&mut self.radon_buffers,
#[cfg(feature = "ml-refiner")]
&mut self.ml_state,
#[cfg(feature = "ml-refiner")]
&self.ml_params,
upscaled,
);
upscale::rescale_descriptors_to_input(&mut corners, factor);
Ok(corners)
}
#[cfg(feature = "image")]
pub fn detect(&mut self, img: &image::GrayImage) -> Result<Vec<CornerDescriptor>, ChessError> {
self.detect_u8(img.as_raw(), img.width(), img.height())
}
pub fn diagnostics(&self) -> crate::diagnostics::DetectorDiagnostics<'_> {
crate::diagnostics::DetectorDiagnostics::new(self)
}
#[allow(clippy::too_many_arguments)]
fn detect_view_inner(
cfg: &DetectorConfig,
pyramid: &mut PyramidBuffers,
chess_buffers: &mut ChessBuffers,
radon_buffers: &mut RadonBuffers,
#[cfg(feature = "ml-refiner")] ml_state: &mut Option<ml_refiner::MlRefinerState>,
#[cfg(feature = "ml-refiner")] ml_params: &ml_refiner::MlRefinerParams,
view: ImageView<'_>,
) -> Vec<CornerDescriptor> {
#[cfg(feature = "ml-refiner")]
if Self::is_ml_refiner(cfg) {
if ml_state.is_none() {
let fallback = chess_corners_core::RefinerKind::CenterOfMass(
chess_corners_core::CenterOfMassConfig::default(),
);
*ml_state = Some(ml_refiner::MlRefinerState::new(ml_params, &fallback));
}
let state = ml_state.as_mut().expect("ml_state initialised above");
return multiscale::detect_with_ml(
view,
cfg,
pyramid,
chess_buffers,
radon_buffers,
ml_params,
state,
);
}
multiscale::detect_with_buffers(view, cfg, pyramid, chess_buffers, radon_buffers)
}
#[cfg(feature = "ml-refiner")]
#[inline]
fn is_ml_refiner(cfg: &DetectorConfig) -> bool {
matches!(
&cfg.strategy,
crate::DetectionStrategy::Chess(c) if matches!(c.refiner, crate::ChessRefiner::Ml)
)
}
}