use serde::{Deserialize, Serialize};
use crate::error::{OcrError, Result};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Decoder {
#[default]
Greedy,
BeamSearch,
WordBeamSearch,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct RecognitionConfig {
pub decoder: Decoder,
pub beam_width: usize,
pub batch_size: usize,
pub allowlist: String,
pub blocklist: String,
pub contrast_ths: f32,
pub adjust_contrast: f32,
pub filter_ths: f32,
}
impl RecognitionConfig {
pub(crate) fn validate(&self) -> Result<()> {
if !(0.0..=1.0).contains(&self.filter_ths) {
return Err(OcrError::config(format!(
"recognition.filter_ths must be finite and within [0, 1], got {}",
self.filter_ths
)));
}
if self.beam_width == 0 {
return Err(OcrError::config("recognition.beam_width must be at least 1, got 0"));
}
Ok(())
}
}
impl Default for RecognitionConfig {
fn default() -> Self {
Self {
decoder: Decoder::Greedy,
beam_width: 5,
batch_size: 1,
allowlist: String::new(),
blocklist: String::new(),
contrast_ths: 0.1,
adjust_contrast: 0.5,
filter_ths: 0.1,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn should_accept_filter_threshold_bounds() {
for filter_ths in [0.0, 1.0] {
let config = RecognitionConfig {
filter_ths,
..RecognitionConfig::default()
};
config
.validate()
.expect("inclusive filter threshold bound should be valid");
}
}
#[test]
fn should_default_to_quality_corpus_filter_threshold() {
assert_eq!(RecognitionConfig::default().filter_ths, 0.1);
}
#[test]
fn should_reject_invalid_filter_thresholds() {
for filter_ths in [-0.1, 1.1, f32::NAN, f32::INFINITY, f32::NEG_INFINITY] {
let config = RecognitionConfig {
filter_ths,
..RecognitionConfig::default()
};
let error = config.validate().expect_err("invalid filter threshold should fail");
assert!(matches!(error, OcrError::Config { .. }));
}
}
#[test]
fn should_default_to_greedy_decoding_with_easyocr_beam_width() {
let config = RecognitionConfig::default();
assert_eq!(config.decoder, Decoder::Greedy);
assert_eq!(config.beam_width, 5);
}
#[test]
fn should_reject_a_zero_beam_width() {
let config = RecognitionConfig {
beam_width: 0,
..RecognitionConfig::default()
};
let error = config.validate().expect_err("a zero beam width should fail");
assert!(matches!(error, OcrError::Config { .. }));
}
}