use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use crate::core::image::OcrImage;
use crate::core::text::BoundingBox;
use crate::utils::Result;
use ndarray::Array2;
#[allow(async_fn_in_trait)]
pub trait TextRecognizer: Send + Sync {
async fn recognize(&self, image: &OcrImage) -> Result<RecognitionResult>;
}
pub trait TrainableModel: Send + Sync {
fn forward_train(&self, input: &Array2<f32>) -> Result<Array2<f32>>;
fn backward_train(&mut self, input: &Array2<f32>, output_grad: &Array2<f32>) -> Result<()>;
fn get_params_and_grads(&mut self) -> Vec<(&mut Array2<f32>, &Array2<f32>)>;
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ModelType {
LSTM,
Transformer,
VisionTransformer,
CNN,
Hybrid,
EndToEnd,
Custom(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum RecognitionEngine {
LSTM,
PatternMatching,
Hybrid,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecognitionConfig {
pub engine: RecognitionEngine,
pub language: String,
pub confidence_threshold: f32,
pub character_whitelist: Option<Vec<char>>,
pub character_blacklist: Option<Vec<char>>,
pub enable_dictionary_correction: bool,
pub enable_language_model: bool,
pub parameters: HashMap<String, String>,
}
impl Default for RecognitionConfig {
fn default() -> Self {
Self {
engine: RecognitionEngine::LSTM,
language: "en".to_string(),
confidence_threshold: 0.5,
character_whitelist: None,
character_blacklist: None,
enable_dictionary_correction: true,
enable_language_model: true,
parameters: HashMap::new(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecognitionResult {
pub text: String,
pub confidence: f32,
pub characters: Vec<CharacterRecognition>,
pub words: Vec<WordRecognition>,
pub lines: Vec<LineRecognition>,
pub metadata: RecognitionMetadata,
pub model_type: Option<ModelType>,
pub processing_time_ms: Option<u64>,
pub language: Option<String>,
pub character_results: Vec<CharacterRecognition>,
pub word_results: Vec<WordRecognition>,
pub line_results: Vec<LineRecognition>,
}
impl RecognitionResult {
pub fn new(text: String, confidence: f32) -> Self {
Self {
text,
confidence,
characters: Vec::new(),
words: Vec::new(),
lines: Vec::new(),
metadata: RecognitionMetadata::default(),
model_type: None,
processing_time_ms: None,
language: None,
character_results: Vec::new(),
word_results: Vec::new(),
line_results: Vec::new(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CharacterRecognition {
pub character: char,
pub confidence: f32,
pub code_point: u32,
pub properties: CharacterProperties,
pub alternatives: Vec<CharacterAlternative>,
pub bounding_box: Option<BoundingBox>,
pub features: Option<(Vec<u8>, Vec<u8>)>,
}
impl CharacterRecognition {
pub fn new(character: char, confidence: f32) -> Self {
Self {
character,
confidence,
code_point: character as u32,
properties: CharacterProperties::default(),
alternatives: Vec::new(),
bounding_box: None,
features: None,
}
}
pub fn with_bounding_box(character: char, confidence: f32, bounding_box: BoundingBox) -> Self {
Self {
character,
confidence,
code_point: character as u32,
properties: CharacterProperties::default(),
alternatives: Vec::new(),
bounding_box: Some(bounding_box),
features: None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CharacterAlternative {
pub character: char,
pub confidence: f32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CharacterProperties {
pub is_digit: bool,
pub is_letter: bool,
pub is_whitespace: bool,
pub is_punctuation: bool,
pub width: f32,
pub height: f32,
pub baseline: f32,
pub x_height: f32,
}
impl Default for CharacterProperties {
fn default() -> Self {
Self {
is_digit: false,
is_letter: false,
is_whitespace: false,
is_punctuation: false,
width: 0.0,
height: 0.0,
baseline: 0.0,
x_height: 0.0,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WordRecognition {
pub word: String,
pub confidence: f32,
pub characters: Vec<CharacterRecognition>,
pub properties: WordProperties,
pub alternatives: Vec<WordAlternative>,
pub bounding_box: Option<BoundingBox>,
}
impl WordRecognition {
pub fn new(word: String, confidence: f32) -> Self {
Self {
word,
confidence,
characters: Vec::new(),
properties: WordProperties::default(),
alternatives: Vec::new(),
bounding_box: None,
}
}
pub fn with_bounding_box(word: String, confidence: f32, bounding_box: BoundingBox) -> Self {
Self {
word,
confidence,
characters: Vec::new(),
properties: WordProperties::default(),
alternatives: Vec::new(),
bounding_box: Some(bounding_box),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WordAlternative {
pub word: String,
pub confidence: f32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WordProperties {
pub is_dictionary_word: bool,
pub length: usize,
pub average_character_width: f32,
pub average_character_height: f32,
pub language: Option<String>,
}
impl Default for WordProperties {
fn default() -> Self {
Self {
is_dictionary_word: false,
length: 0,
average_character_width: 0.0,
average_character_height: 0.0,
language: None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LineRecognition {
pub line: String,
pub confidence: f32,
pub words: Vec<WordRecognition>,
pub properties: LineProperties,
pub bounding_box: Option<BoundingBox>,
}
impl LineRecognition {
pub fn new(line: String, confidence: f32) -> Self {
Self {
line,
confidence,
words: Vec::new(),
properties: LineProperties::default(),
bounding_box: None,
}
}
pub fn with_bounding_box(line: String, confidence: f32, bounding_box: BoundingBox) -> Self {
Self {
line,
confidence,
words: Vec::new(),
properties: LineProperties::default(),
bounding_box: Some(bounding_box),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LineProperties {
pub height: f32,
pub baseline: f32,
pub line_spacing: f32,
pub alignment: TextAlignment,
pub reading_order: ReadingOrder,
}
impl Default for LineProperties {
fn default() -> Self {
Self {
height: 0.0,
baseline: 0.0,
line_spacing: 0.0,
alignment: TextAlignment::Left,
reading_order: ReadingOrder::LeftToRight,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecognitionMetadata {
pub engine: RecognitionEngine,
pub language: Option<String>,
pub processing_time_ms: u64,
pub model_version: Option<String>,
pub additional: HashMap<String, String>,
}
impl Default for RecognitionMetadata {
fn default() -> Self {
Self {
engine: RecognitionEngine::LSTM,
language: None,
processing_time_ms: 0,
model_version: None,
additional: HashMap::new(),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum TextAlignment {
Left,
Right,
Center,
Justified,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ReadingOrder {
LeftToRight,
RightToLeft,
TopToBottom,
BottomToTop,
}
pub trait RecognitionEngineTrait {
fn initialize(&mut self, config: &RecognitionConfig) -> anyhow::Result<()>;
fn recognize(
&self,
image_data: &[u8],
width: u32,
height: u32,
) -> anyhow::Result<RecognitionResult>;
fn recognize_region(
&self,
image_data: &[u8],
width: u32,
height: u32,
region: &ImageRegion,
) -> anyhow::Result<RecognitionResult>;
fn capabilities(&self) -> EngineCapabilities;
fn info(&self) -> EngineInfo;
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EngineCapabilities {
pub supported_languages: Vec<String>,
pub max_image_size: (u32, u32),
pub min_image_size: (u32, u32),
pub supported_formats: Vec<String>,
pub supports_character_level: bool,
pub supports_word_level: bool,
pub supports_line_level: bool,
pub supports_confidence: bool,
pub supports_alternatives: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EngineInfo {
pub name: String,
pub version: String,
pub description: String,
pub author: String,
pub license: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct ImageRegion {
pub left: u32,
pub top: u32,
pub right: u32,
pub bottom: u32,
}
impl ImageRegion {
pub fn new(left: u32, top: u32, right: u32, bottom: u32) -> Self {
Self {
left,
top,
right,
bottom,
}
}
pub fn width(&self) -> u32 {
self.right - self.left
}
pub fn height(&self) -> u32 {
self.bottom - self.top
}
pub fn area(&self) -> u32 {
self.width() * self.height()
}
}