Skip to main content

ocr_rs/
rec.rs

1//! Text Recognition Model
2//!
3//! Provides text recognition functionality based on PaddleOCR recognition models
4
5use image::{DynamicImage, RgbImage};
6use imageproc::geometric_transformations::Projection;
7use imageproc::point::Point;
8use ndarray::{Array4, ArrayD, ArrayViewD, Axis};
9use std::{borrow::Cow, path::Path};
10
11use crate::error::{OcrError, OcrResult};
12use crate::mnn::{InferenceConfig, InferenceEngine};
13use crate::postprocess::TextBox;
14use crate::preprocess::{preprocess_for_rec, NormalizeParams};
15
16/// Recognition result
17#[derive(Debug, Clone)]
18pub struct RecognitionResult {
19    /// Recognized text
20    pub text: String,
21    /// Confidence score (0.0 - 1.0)
22    pub confidence: f32,
23    /// Confidence score for each character
24    pub char_scores: Vec<(char, f32)>,
25}
26
27impl RecognitionResult {
28    /// Create a new recognition result
29    pub fn new(text: String, confidence: f32, char_scores: Vec<(char, f32)>) -> Self {
30        Self {
31            text,
32            confidence,
33            char_scores,
34        }
35    }
36
37    /// Check if the result is valid (confidence above threshold)
38    pub fn is_valid(&self, threshold: f32) -> bool {
39        self.confidence >= threshold
40    }
41}
42
43/// Recognition options
44#[derive(Debug, Clone)]
45pub struct RecOptions {
46    /// Target height (recognition model input height)
47    pub target_height: u32,
48    /// Minimum confidence threshold (characters below this value will be filtered)
49    pub min_score: f32,
50    /// Minimum confidence threshold for punctuation
51    pub punct_min_score: f32,
52    /// Batch size
53    pub batch_size: usize,
54    /// Whether to enable batch processing
55    pub enable_batch: bool,
56}
57
58impl Default for RecOptions {
59    fn default() -> Self {
60        Self {
61            target_height: 48,
62            min_score: 0.3, // Lower threshold, model output is raw logit
63            punct_min_score: 0.1,
64            batch_size: 8,
65            enable_batch: true,
66        }
67    }
68}
69
70impl RecOptions {
71    /// Create new recognition options
72    pub fn new() -> Self {
73        Self::default()
74    }
75
76    /// Set target height
77    pub fn with_target_height(mut self, height: u32) -> Self {
78        self.target_height = height;
79        self
80    }
81
82    /// Set minimum confidence
83    pub fn with_min_score(mut self, score: f32) -> Self {
84        self.min_score = score;
85        self
86    }
87
88    /// Set punctuation minimum confidence
89    pub fn with_punct_min_score(mut self, score: f32) -> Self {
90        self.punct_min_score = score;
91        self
92    }
93
94    /// Set batch size
95    pub fn with_batch_size(mut self, size: usize) -> Self {
96        self.batch_size = size;
97        self
98    }
99
100    /// Enable/disable batch processing
101    pub fn with_batch(mut self, enable: bool) -> Self {
102        self.enable_batch = enable;
103        self
104    }
105}
106
107/// Text recognition model
108pub struct RecModel {
109    engine: InferenceEngine,
110    /// Character set (index to character mapping)
111    charset: Vec<char>,
112    options: RecOptions,
113    normalize_params: NormalizeParams,
114}
115
116/// Common punctuation marks
117const PUNCTUATIONS: [char; 49] = [
118    ',', '.', '!', '?', ';', ':', '"', '\'', '(', ')', '[', ']', '{', '}', '-', '_', '/', '\\',
119    '|', '@', '#', '$', '%', '&', '*', '+', '=', '~', ',', '。', '!', '?', ';', ':', '、',
120    '「', '」', '『', '』', '(', ')', '【', '】', '《', '》', '—', '…', '·', '~',
121];
122
123impl RecModel {
124    /// Create recognizer from model file and charset file
125    ///
126    /// # Parameters
127    /// - `model_path`: Model file path (.mnn format)
128    /// - `charset_path`: Charset file path (one character per line)
129    /// - `config`: Optional inference config
130    pub fn from_file(
131        model_path: impl AsRef<Path>,
132        charset_path: impl AsRef<Path>,
133        config: Option<InferenceConfig>,
134    ) -> OcrResult<Self> {
135        let engine = InferenceEngine::from_file(model_path, config)?;
136        let charset = Self::load_charset_from_file(charset_path)?;
137
138        Ok(Self {
139            engine,
140            charset,
141            options: RecOptions::default(),
142            normalize_params: NormalizeParams::paddle_rec(),
143        })
144    }
145
146    /// Create recognizer from model bytes and charset file
147    pub fn from_bytes(
148        model_bytes: &[u8],
149        charset_path: impl AsRef<Path>,
150        config: Option<InferenceConfig>,
151    ) -> OcrResult<Self> {
152        let engine = InferenceEngine::from_buffer(model_bytes, config)?;
153        let charset = Self::load_charset_from_file(charset_path)?;
154
155        Ok(Self {
156            engine,
157            charset,
158            options: RecOptions::default(),
159            normalize_params: NormalizeParams::paddle_rec(),
160        })
161    }
162
163    /// Create recognizer from model bytes and charset bytes
164    pub fn from_bytes_with_charset(
165        model_bytes: &[u8],
166        charset_bytes: &[u8],
167        config: Option<InferenceConfig>,
168    ) -> OcrResult<Self> {
169        let engine = InferenceEngine::from_buffer(model_bytes, config)?;
170        let charset = Self::parse_charset(charset_bytes)?;
171
172        Ok(Self {
173            engine,
174            charset,
175            options: RecOptions::default(),
176            normalize_params: NormalizeParams::paddle_rec(),
177        })
178    }
179
180    /// Load charset from file
181    fn load_charset_from_file(path: impl AsRef<Path>) -> OcrResult<Vec<char>> {
182        let content = std::fs::read_to_string(path)?;
183        Self::parse_charset(content.as_bytes())
184    }
185
186    /// Parse charset data
187    fn parse_charset(data: &[u8]) -> OcrResult<Vec<char>> {
188        let content = std::str::from_utf8(data)
189            .map_err(|e| OcrError::CharsetError(format!("UTF-8 decode error: {}", e)))?;
190
191        // Charset format: one character per line
192        // Add space at beginning and end as blank and padding
193        let mut charset: Vec<char> = vec![' ']; // blank token at start
194
195        for ch in content.chars() {
196            if ch != '\n' && ch != '\r' {
197                charset.push(ch);
198            }
199        }
200
201        charset.push(' '); // padding token at end
202
203        if charset.len() < 3 {
204            return Err(OcrError::CharsetError("Charset too small".to_string()));
205        }
206
207        Ok(charset)
208    }
209
210    /// Set recognition options
211    pub fn with_options(mut self, options: RecOptions) -> Self {
212        self.options = options;
213        self
214    }
215
216    /// Get current recognition options
217    pub fn options(&self) -> &RecOptions {
218        &self.options
219    }
220
221    /// Modify recognition options
222    pub fn options_mut(&mut self) -> &mut RecOptions {
223        &mut self.options
224    }
225
226    /// Get charset size
227    pub fn charset_size(&self) -> usize {
228        self.charset.len()
229    }
230
231    /// Recognize a single image
232    ///
233    /// # Parameters
234    /// - `image`: Input image (text line image)
235    ///
236    /// # Returns
237    /// Recognition result
238    pub fn recognize(&self, image: &DynamicImage) -> OcrResult<RecognitionResult> {
239        // Preprocess
240        let input = preprocess_for_rec(image, self.options.target_height, &self.normalize_params)?;
241
242        // Inference (using dynamic shape)
243        let output = self.engine.run_dynamic(input.view().into_dyn())?;
244
245        // Decode
246        self.decode_output_view(output.view())
247    }
248
249    /// Recognize a single image, return text only
250    pub fn recognize_text(&self, image: &DynamicImage) -> OcrResult<String> {
251        let result = self.recognize(image)?;
252        Ok(result.text)
253    }
254
255    /// Batch recognize images
256    ///
257    /// # Parameters
258    /// - `images`: List of input images
259    ///
260    /// # Returns
261    /// List of recognition results
262    pub fn recognize_batch(&self, images: &[DynamicImage]) -> OcrResult<Vec<RecognitionResult>> {
263        if images.is_empty() {
264            return Ok(Vec::new());
265        }
266
267        // For small number of images, process individually
268        if images.len() <= 2 || !self.options.enable_batch {
269            return images.iter().map(|img| self.recognize(img)).collect();
270        }
271
272        // Batch processing
273        let mut results = Vec::with_capacity(images.len());
274
275        for chunk in images.chunks(self.options.batch_size) {
276            let batch_results = self.recognize_batch_internal(chunk)?;
277            results.extend(batch_results);
278        }
279
280        Ok(results)
281    }
282
283    /// Batch recognize images (borrowed version, avoid cloning)
284    ///
285    /// # Parameters
286    /// - `images`: List of input image references
287    ///
288    /// # Returns
289    /// List of recognition results
290    pub fn recognize_batch_ref(
291        &self,
292        images: &[&DynamicImage],
293    ) -> OcrResult<Vec<RecognitionResult>> {
294        if images.is_empty() {
295            return Ok(Vec::new());
296        }
297
298        // For small number of images, process individually
299        if images.len() <= 2 || !self.options.enable_batch {
300            return images.iter().map(|img| self.recognize(img)).collect();
301        }
302
303        // Batch processing
304        let mut results = Vec::with_capacity(images.len());
305
306        for chunk in images.chunks(self.options.batch_size) {
307            // Dereference and convert to Vec<DynamicImage>
308            let chunk_owned: Vec<DynamicImage> = chunk.iter().map(|img| (*img).clone()).collect();
309            let batch_results = self.recognize_batch_internal(&chunk_owned)?;
310            results.extend(batch_results);
311        }
312
313        Ok(results)
314    }
315
316    pub(crate) fn recognize_regions(
317        &self,
318        image: &DynamicImage,
319        boxes: &[TextBox],
320    ) -> OcrResult<Vec<RecognitionResult>> {
321        if boxes.is_empty() {
322            return Ok(Vec::new());
323        }
324
325        let source = image.to_rgb8();
326        let mut results = Vec::with_capacity(boxes.len());
327        let batch_size = self.options.batch_size.max(1);
328
329        for chunk in boxes.chunks(batch_size) {
330            let batch_input = self.preprocess_regions_batch(&source, chunk)?;
331            let batch_output = self.engine.run_dynamic(batch_input.view().into_dyn())?;
332
333            let shape = batch_output.shape();
334            if shape.len() != 3 {
335                return Err(OcrError::PostprocessError(format!(
336                    "Region batch inference output shape error: {:?}",
337                    shape
338                )));
339            }
340
341            for i in 0..shape[0] {
342                let sample_output = batch_output.index_axis(Axis(0), i).into_dyn();
343                results.push(self.decode_output_view(sample_output)?);
344            }
345        }
346
347        Ok(results)
348    }
349
350    fn preprocess_regions_batch(
351        &self,
352        source: &RgbImage,
353        boxes: &[TextBox],
354    ) -> OcrResult<Array4<f32>> {
355        if boxes.is_empty() {
356            return Ok(Array4::<f32>::zeros((
357                0,
358                3,
359                self.options.target_height as usize,
360                0,
361            )));
362        }
363
364        if self.options.target_height == 0 {
365            return Err(OcrError::InvalidParameter(
366                "Recognition target height must be greater than 0".into(),
367            ));
368        }
369
370        let target_height = self.options.target_height;
371        let target_widths = boxes
372            .iter()
373            .map(|text_box| region_target_width(text_box, target_height))
374            .collect::<Vec<_>>();
375        let max_width = target_widths.iter().copied().max().unwrap_or(1) as usize;
376        let batch_size = boxes.len();
377        let target_height_usize = target_height as usize;
378        let sample_size = 3 * target_height_usize * max_width;
379        let plane_size = target_height_usize * max_width;
380
381        let mut batch = Array4::<f32>::zeros((batch_size, 3, target_height_usize, max_width));
382        let data = batch
383            .as_slice_mut()
384            .expect("Array4 created by zeros should be contiguous");
385        let scales = [
386            1.0 / (255.0 * self.normalize_params.std[0]),
387            1.0 / (255.0 * self.normalize_params.std[1]),
388            1.0 / (255.0 * self.normalize_params.std[2]),
389        ];
390        let offsets = [
391            -self.normalize_params.mean[0] / self.normalize_params.std[0],
392            -self.normalize_params.mean[1] / self.normalize_params.std[1],
393            -self.normalize_params.mean[2] / self.normalize_params.std[2],
394        ];
395
396        for (i, (text_box, &target_width)) in boxes.iter().zip(target_widths.iter()).enumerate() {
397            let projection =
398                target_to_source_projection(source, text_box, target_width, target_height)
399                    .ok_or_else(|| {
400                        OcrError::PreprocessError(format!(
401                            "Failed to render recognition region: {:?}",
402                            text_box.rect
403                        ))
404                    })?;
405            let target_width = target_width as usize;
406            let sample_offset = i * sample_size;
407
408            write_projected_region_to_tensor(
409                source,
410                projection,
411                target_width,
412                target_height_usize,
413                max_width,
414                sample_offset,
415                plane_size,
416                data,
417                &scales,
418                &offsets,
419            );
420        }
421
422        Ok(batch)
423    }
424
425    /// Internal batch recognition
426    fn recognize_batch_internal(
427        &self,
428        images: &[DynamicImage],
429    ) -> OcrResult<Vec<RecognitionResult>> {
430        if images.is_empty() {
431            return Ok(Vec::new());
432        }
433
434        // If only one image, process individually
435        if images.len() == 1 {
436            return Ok(vec![self.recognize(&images[0])?]);
437        }
438
439        // Batch preprocessing
440        let batch_input = crate::preprocess::preprocess_batch_for_rec(
441            images,
442            self.options.target_height,
443            &self.normalize_params,
444        )?;
445
446        // Batch inference
447        let batch_output = self.engine.run_dynamic(batch_input.view().into_dyn())?;
448
449        // Decode output for each sample
450        let shape = batch_output.shape();
451        if shape.len() != 3 {
452            return Err(OcrError::PostprocessError(format!(
453                "Batch inference output shape error: {:?}",
454                shape
455            )));
456        }
457
458        let batch_size = shape[0];
459        let mut results = Vec::with_capacity(batch_size);
460
461        for i in 0..batch_size {
462            let sample_output = batch_output.index_axis(Axis(0), i).into_dyn();
463            let result = self.decode_output_view(sample_output)?;
464            results.push(result);
465        }
466
467        Ok(results)
468    }
469
470    fn decode_output_view(&self, output: ArrayViewD<'_, f32>) -> OcrResult<RecognitionResult> {
471        let shape = output.shape();
472        let output_data = match output.as_slice_memory_order() {
473            Some(slice) => Cow::Borrowed(slice),
474            None => Cow::Owned(output.iter().copied().collect()),
475        };
476
477        // Output shape should be [batch, seq_len, num_classes] or [seq_len, num_classes]
478        let (seq_len, num_classes) = if shape.len() == 3 {
479            (shape[1], shape[2])
480        } else if shape.len() == 2 {
481            (shape[0], shape[1])
482        } else {
483            return Err(OcrError::PostprocessError(format!(
484                "Invalid output shape: {:?}",
485                shape
486            )));
487        };
488
489        if num_classes == 0 {
490            return Err(OcrError::PostprocessError(
491                "Invalid output shape with zero classes".into(),
492            ));
493        }
494
495        // CTC decoding
496        let mut char_scores = Vec::with_capacity(seq_len.min(32));
497        let mut text = String::new();
498        let mut score_sum = 0.0f32;
499        let mut prev_idx = 0usize;
500
501        for t in 0..seq_len {
502            // Find character with maximum probability at current time step
503            let start = t * num_classes;
504            let end = start + num_classes;
505            let probs = &output_data[start..end];
506
507            let mut max_idx = 0usize;
508            let mut max_prob = f32::NEG_INFINITY;
509            for (idx, &prob) in probs.iter().enumerate() {
510                if prob > max_prob {
511                    max_idx = idx;
512                    max_prob = prob;
513                }
514            }
515
516            // CTC decoding rule: skip blank (index 0) and duplicate characters
517            if max_idx != 0 && max_idx != prev_idx {
518                if max_idx < self.charset.len() {
519                    let ch = self.charset[max_idx];
520
521                    // Use raw logit value as confidence (model output is already softmax probability)
522                    // For large character sets, softmax scores can be very small, so use max_prob directly
523                    let score = max_prob;
524
525                    // Only filter out very low confidence characters
526                    let threshold = if Self::is_punctuation(ch) {
527                        self.options.punct_min_score
528                    } else {
529                        self.options.min_score
530                    };
531
532                    if score >= threshold {
533                        text.push(ch);
534                        score_sum += score;
535                        char_scores.push((ch, score));
536                    }
537                }
538            }
539
540            prev_idx = max_idx;
541        }
542
543        // Calculate average confidence
544        let confidence = if char_scores.is_empty() {
545            0.0
546        } else {
547            score_sum / char_scores.len() as f32
548        };
549
550        Ok(RecognitionResult::new(text, confidence, char_scores))
551    }
552
553    /// Check if character is punctuation
554    fn is_punctuation(ch: char) -> bool {
555        PUNCTUATIONS.contains(&ch)
556    }
557}
558
559fn region_target_width(text_box: &TextBox, target_height: u32) -> u32 {
560    let (width, height) = region_dimensions(text_box);
561    ((width / height.max(1.0)) * target_height as f32)
562        .round()
563        .max(2.0) as u32
564}
565
566fn target_to_source_projection(
567    source: &RgbImage,
568    text_box: &TextBox,
569    target_width: u32,
570    target_height: u32,
571) -> Option<Projection> {
572    if target_width < 2 || target_height < 2 {
573        return None;
574    }
575
576    if let Some(source_points) =
577        source_points_for_text_box(text_box, source.width(), source.height())
578    {
579        if let Some(projection) =
580            build_target_to_source_projection(source_points, target_width, target_height)
581        {
582            return Some(projection);
583        }
584    }
585
586    let source_points = rect_source_points_for_text_box(text_box, source.width(), source.height())?;
587    build_target_to_source_projection(source_points, target_width, target_height)
588}
589
590fn build_target_to_source_projection(
591    source_points: [(f32, f32); 4],
592    target_width: u32,
593    target_height: u32,
594) -> Option<Projection> {
595    let target_points = [
596        (0.0, 0.0),
597        (target_width.saturating_sub(1) as f32, 0.0),
598        (
599            target_width.saturating_sub(1) as f32,
600            target_height.saturating_sub(1) as f32,
601        ),
602        (0.0, target_height.saturating_sub(1) as f32),
603    ];
604
605    Projection::from_control_points(source_points, target_points)
606        .map(|projection| projection.invert())
607}
608
609#[allow(clippy::too_many_arguments)]
610fn write_projected_region_to_tensor(
611    source: &RgbImage,
612    target_to_source: Projection,
613    target_width: usize,
614    target_height: usize,
615    max_width: usize,
616    sample_offset: usize,
617    plane_size: usize,
618    data: &mut [f32],
619    scales: &[f32; 3],
620    offsets: &[f32; 3],
621) {
622    let source_width = source.width() as usize;
623    let source_height = source.height() as usize;
624    let source_data = source.as_raw();
625
626    for y in 0..target_height {
627        let dst_row = y * max_width;
628
629        for x in 0..target_width {
630            let (source_x, source_y) = target_to_source * (x as f32, y as f32);
631            let dst = sample_offset + dst_row + x;
632            write_normalized_sample(
633                source_data,
634                source_width,
635                source_height,
636                source_x,
637                source_y,
638                data,
639                dst,
640                sample_offset + plane_size + dst_row + x,
641                sample_offset + plane_size * 2 + dst_row + x,
642                scales,
643                offsets,
644            );
645        }
646    }
647}
648
649#[allow(clippy::too_many_arguments)]
650#[inline(always)]
651fn write_normalized_sample(
652    source_data: &[u8],
653    source_width: usize,
654    source_height: usize,
655    x: f32,
656    y: f32,
657    data: &mut [f32],
658    dst_r: usize,
659    dst_g: usize,
660    dst_b: usize,
661    scales: &[f32; 3],
662    offsets: &[f32; 3],
663) {
664    let left = x.floor();
665    let right = left + 1.0;
666    let top = y.floor();
667    let bottom = top + 1.0;
668
669    if !(left >= 0.0 && right < source_width as f32 && top >= 0.0 && bottom < source_height as f32)
670    {
671        data[dst_r] = 255.0 * scales[0] + offsets[0];
672        data[dst_g] = 255.0 * scales[1] + offsets[1];
673        data[dst_b] = 255.0 * scales[2] + offsets[2];
674        return;
675    }
676
677    let right_weight = x - left;
678    let bottom_weight = y - top;
679    let left = left as usize;
680    let right = right as usize;
681    let top = top as usize;
682    let bottom = bottom as usize;
683    let top_left = (top * source_width + left) * 3;
684    let top_right = (top * source_width + right) * 3;
685    let bottom_left = (bottom * source_width + left) * 3;
686    let bottom_right = (bottom * source_width + right) * 3;
687
688    let r = bilinear_channel(
689        source_data[top_left],
690        source_data[top_right],
691        source_data[bottom_left],
692        source_data[bottom_right],
693        right_weight,
694        bottom_weight,
695    );
696    let g = bilinear_channel(
697        source_data[top_left + 1],
698        source_data[top_right + 1],
699        source_data[bottom_left + 1],
700        source_data[bottom_right + 1],
701        right_weight,
702        bottom_weight,
703    );
704    let b = bilinear_channel(
705        source_data[top_left + 2],
706        source_data[top_right + 2],
707        source_data[bottom_left + 2],
708        source_data[bottom_right + 2],
709        right_weight,
710        bottom_weight,
711    );
712
713    data[dst_r] = r as f32 * scales[0] + offsets[0];
714    data[dst_g] = g as f32 * scales[1] + offsets[1];
715    data[dst_b] = b as f32 * scales[2] + offsets[2];
716}
717
718#[inline(always)]
719fn bilinear_channel(
720    top_left: u8,
721    top_right: u8,
722    bottom_left: u8,
723    bottom_right: u8,
724    right_weight: f32,
725    bottom_weight: f32,
726) -> u8 {
727    let top = lerp(top_left as f32, top_right as f32, right_weight);
728    let bottom = lerp(bottom_left as f32, bottom_right as f32, right_weight);
729    clamp_to_u8(lerp(top, bottom, bottom_weight))
730}
731
732#[inline]
733fn lerp(left: f32, right: f32, weight: f32) -> f32 {
734    (1.0 - weight) * left + weight * right
735}
736
737#[inline]
738fn clamp_to_u8(value: f32) -> u8 {
739    if value < u8::MAX as f32 {
740        if value > u8::MIN as f32 {
741            value as u8
742        } else {
743            u8::MIN
744        }
745    } else {
746        u8::MAX
747    }
748}
749
750fn source_points_for_text_box(
751    text_box: &TextBox,
752    image_width: u32,
753    image_height: u32,
754) -> Option<[(f32, f32); 4]> {
755    if let Some(points) = text_box.points {
756        let max_x = image_width.saturating_sub(1) as f32;
757        let max_y = image_height.saturating_sub(1) as f32;
758        return Some(points.map(|point| (point.x.clamp(0.0, max_x), point.y.clamp(0.0, max_y))));
759    }
760
761    rect_source_points_for_text_box(text_box, image_width, image_height)
762}
763
764fn rect_source_points_for_text_box(
765    text_box: &TextBox,
766    image_width: u32,
767    image_height: u32,
768) -> Option<[(f32, f32); 4]> {
769    let left = text_box.rect.left().max(0) as u32;
770    let top = text_box.rect.top().max(0) as u32;
771    let right = left
772        .saturating_add(text_box.rect.width())
773        .min(image_width)
774        .saturating_sub(1);
775    let bottom = top
776        .saturating_add(text_box.rect.height())
777        .min(image_height)
778        .saturating_sub(1);
779
780    if right <= left || bottom <= top {
781        return None;
782    }
783
784    Some([
785        (left as f32, top as f32),
786        (right as f32, top as f32),
787        (right as f32, bottom as f32),
788        (left as f32, bottom as f32),
789    ])
790}
791
792fn region_dimensions(text_box: &TextBox) -> (f32, f32) {
793    if let Some(points) = text_box.points {
794        let width = distance(points[0], points[1]).max(distance(points[3], points[2]));
795        let height = distance(points[0], points[3]).max(distance(points[1], points[2]));
796        (width.max(1.0), height.max(1.0))
797    } else {
798        (
799            text_box.rect.width().max(1) as f32,
800            text_box.rect.height().max(1) as f32,
801        )
802    }
803}
804
805fn distance(a: Point<f32>, b: Point<f32>) -> f32 {
806    let dx = a.x - b.x;
807    let dy = a.y - b.y;
808    (dx * dx + dy * dy).sqrt()
809}
810
811/// Low-level recognition API
812impl RecModel {
813    /// Raw inference interface
814    ///
815    /// Execute model inference directly without preprocessing and postprocessing
816    ///
817    /// # Parameters
818    /// - `input`: Preprocessed input tensor [1, 3, H, W]
819    ///
820    /// # Returns
821    /// Model raw output
822    pub fn run_raw(&self, input: ndarray::ArrayViewD<f32>) -> OcrResult<ArrayD<f32>> {
823        Ok(self.engine.run_dynamic(input)?)
824    }
825
826    /// Get model input shape
827    pub fn input_shape(&self) -> &[usize] {
828        self.engine.input_shape()
829    }
830
831    /// Get model output shape
832    pub fn output_shape(&self) -> &[usize] {
833        self.engine.output_shape()
834    }
835
836    /// Get charset
837    pub fn charset(&self) -> &[char] {
838        &self.charset
839    }
840
841    /// Get character by index
842    pub fn get_char(&self, index: usize) -> Option<char> {
843        self.charset.get(index).copied()
844    }
845}
846
847#[cfg(test)]
848mod tests {
849    use super::*;
850
851    #[test]
852    fn test_rec_options_default() {
853        let opts = RecOptions::default();
854        assert_eq!(opts.target_height, 48);
855        assert_eq!(opts.min_score, 0.3);
856        assert_eq!(opts.punct_min_score, 0.1);
857        assert_eq!(opts.batch_size, 8);
858        assert!(opts.enable_batch);
859    }
860
861    #[test]
862    fn test_rec_options_builder() {
863        let opts = RecOptions::new()
864            .with_target_height(32)
865            .with_min_score(0.6)
866            .with_punct_min_score(0.2)
867            .with_batch_size(16)
868            .with_batch(false);
869
870        assert_eq!(opts.target_height, 32);
871        assert_eq!(opts.min_score, 0.6);
872        assert_eq!(opts.punct_min_score, 0.2);
873        assert_eq!(opts.batch_size, 16);
874        assert!(!opts.enable_batch);
875    }
876
877    #[test]
878    fn test_recognition_result_new() {
879        let char_scores = vec![
880            ('H', 0.99),
881            ('e', 0.94),
882            ('l', 0.93),
883            ('l', 0.95),
884            ('o', 0.94),
885        ];
886        let result = RecognitionResult::new("Hello".to_string(), 0.95, char_scores.clone());
887
888        assert_eq!(result.text, "Hello");
889        assert_eq!(result.confidence, 0.95);
890        assert_eq!(result.char_scores.len(), 5);
891        assert_eq!(result.char_scores[0].0, 'H');
892        assert_eq!(result.char_scores[0].1, 0.99);
893    }
894
895    #[test]
896    fn test_recognition_result_is_valid() {
897        let result = RecognitionResult::new(
898            "Hello".to_string(),
899            0.95,
900            vec![
901                ('H', 0.99),
902                ('e', 0.94),
903                ('l', 0.93),
904                ('l', 0.95),
905                ('o', 0.94),
906            ],
907        );
908
909        assert!(result.is_valid(0.9));
910        assert!(result.is_valid(0.95));
911        assert!(!result.is_valid(0.96));
912        assert!(!result.is_valid(0.99));
913    }
914
915    #[test]
916    fn test_recognition_result_empty() {
917        let result = RecognitionResult::new(String::new(), 0.0, vec![]);
918
919        assert!(result.text.is_empty());
920        assert_eq!(result.confidence, 0.0);
921        assert!(!result.is_valid(0.1));
922    }
923
924    #[test]
925    fn test_region_target_width_avoids_projection_degenerate_width() {
926        let text_box = TextBox::with_points(
927            imageproc::rect::Rect::at(747, 14).of_size(61, 1695),
928            0.9,
929            [
930                Point::new(747.0, 14.0),
931                Point::new(747.4, 14.0),
932                Point::new(747.4, 1709.0),
933                Point::new(747.0, 1709.0),
934            ],
935        );
936
937        assert_eq!(region_target_width(&text_box, 48), 2);
938    }
939
940    #[test]
941    fn test_is_punctuation_common() {
942        // English punctuation
943        assert!(RecModel::is_punctuation(','));
944        assert!(RecModel::is_punctuation('.'));
945        assert!(RecModel::is_punctuation('!'));
946        assert!(RecModel::is_punctuation('?'));
947        assert!(RecModel::is_punctuation(';'));
948        assert!(RecModel::is_punctuation(':'));
949        assert!(RecModel::is_punctuation('"'));
950        assert!(RecModel::is_punctuation('\''));
951    }
952
953    #[test]
954    fn test_is_punctuation_chinese() {
955        // Chinese punctuation
956        assert!(RecModel::is_punctuation(','));
957        assert!(RecModel::is_punctuation('。'));
958        assert!(RecModel::is_punctuation('!'));
959        assert!(RecModel::is_punctuation('?'));
960        assert!(RecModel::is_punctuation(';'));
961        assert!(RecModel::is_punctuation(':'));
962        assert!(RecModel::is_punctuation('、'));
963        assert!(RecModel::is_punctuation('—'));
964        assert!(RecModel::is_punctuation('…'));
965    }
966
967    #[test]
968    fn test_is_punctuation_brackets() {
969        assert!(RecModel::is_punctuation('('));
970        assert!(RecModel::is_punctuation(')'));
971        assert!(RecModel::is_punctuation('['));
972        assert!(RecModel::is_punctuation(']'));
973        assert!(RecModel::is_punctuation('{'));
974        assert!(RecModel::is_punctuation('}'));
975        assert!(RecModel::is_punctuation('「'));
976        assert!(RecModel::is_punctuation('」'));
977        assert!(RecModel::is_punctuation('《'));
978        assert!(RecModel::is_punctuation('》'));
979    }
980
981    #[test]
982    fn test_is_punctuation_false() {
983        // Non-punctuation characters
984        assert!(!RecModel::is_punctuation('A'));
985        assert!(!RecModel::is_punctuation('z'));
986        assert!(!RecModel::is_punctuation('0'));
987        assert!(!RecModel::is_punctuation('中'));
988        assert!(!RecModel::is_punctuation('文'));
989        assert!(!RecModel::is_punctuation(' '));
990    }
991}