ocr_rs/
engine.rs

1//! OCR Engine
2//!
3//! Provides complete OCR pipeline encapsulation, performs detection and recognition in one call
4
5use image::DynamicImage;
6use std::path::{Path, PathBuf};
7
8use crate::det::{DetModel, DetOptions};
9use crate::error::{OcrError, OcrResult};
10use crate::mnn::{Backend, InferenceConfig, PrecisionMode};
11use crate::ori::{OriModel, OriOptions};
12use crate::postprocess::TextBox;
13use crate::rec::{RecModel, RecOptions, RecognitionResult};
14
15const PARALLEL_RECOGNITION_MIN_REGIONS: usize = 5;
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18enum RegionRecognitionStrategy {
19    Batch,
20    ExactWidthParallel,
21}
22
23fn select_region_recognition_strategy(
24    enable_parallel: bool,
25    region_count: usize,
26) -> RegionRecognitionStrategy {
27    if enable_parallel && region_count >= PARALLEL_RECOGNITION_MIN_REGIONS {
28        RegionRecognitionStrategy::ExactWidthParallel
29    } else {
30        RegionRecognitionStrategy::Batch
31    }
32}
33
34/// OCR result
35#[derive(Debug, Clone)]
36pub struct OcrResult_ {
37    /// Recognized text
38    pub text: String,
39    /// Confidence score
40    pub confidence: f32,
41    /// Bounding box
42    pub bbox: TextBox,
43}
44
45impl OcrResult_ {
46    /// Create a new OCR result
47    pub fn new(text: String, confidence: f32, bbox: TextBox) -> Self {
48        Self {
49            text,
50            confidence,
51            bbox,
52        }
53    }
54}
55
56/// OCR engine configuration
57#[derive(Debug, Clone)]
58pub struct OcrEngineConfig {
59    /// Inference backend
60    pub backend: Backend,
61    /// Thread count
62    pub thread_count: i32,
63    /// Precision mode
64    pub precision_mode: PrecisionMode,
65    /// Detection options
66    pub det_options: DetOptions,
67    /// Recognition options
68    pub rec_options: RecOptions,
69    /// Orientation options (used when orientation model is enabled)
70    pub ori_options: OriOptions,
71    /// Whether to enable exact-width parallel recognition for multi-line images
72    pub enable_parallel: bool,
73    /// Minimum confidence threshold at result level (recognition results below this value will be filtered)
74    pub min_result_confidence: f32,
75    /// Minimum confidence threshold for orientation correction
76    pub ori_min_confidence: f32,
77}
78
79impl Default for OcrEngineConfig {
80    fn default() -> Self {
81        Self {
82            backend: Backend::CPU,
83            thread_count: 4,
84            precision_mode: PrecisionMode::Normal,
85            det_options: DetOptions::default(),
86            rec_options: RecOptions::default(),
87            ori_options: OriOptions::default(),
88            enable_parallel: true,
89            min_result_confidence: 0.5,
90            ori_min_confidence: 0.3,
91        }
92    }
93}
94
95impl OcrEngineConfig {
96    /// Create new configuration
97    pub fn new() -> Self {
98        Self::default()
99    }
100
101    /// Set inference backend
102    pub fn with_backend(mut self, backend: Backend) -> Self {
103        self.backend = backend;
104        self
105    }
106
107    /// Set thread count
108    pub fn with_threads(mut self, threads: i32) -> Self {
109        self.thread_count = threads;
110        self
111    }
112
113    /// Set precision mode
114    pub fn with_precision(mut self, precision: PrecisionMode) -> Self {
115        self.precision_mode = precision;
116        self
117    }
118
119    /// Set detection options
120    pub fn with_det_options(mut self, options: DetOptions) -> Self {
121        self.det_options = options;
122        self
123    }
124
125    /// Set recognition options
126    pub fn with_rec_options(mut self, options: RecOptions) -> Self {
127        self.rec_options = options;
128        self
129    }
130
131    /// Set orientation options
132    pub fn with_ori_options(mut self, options: OriOptions) -> Self {
133        self.ori_options = options;
134        self
135    }
136
137    /// Enable/disable parallel processing
138    ///
139    /// When at least five text regions are detected, preprocessing is parallelized and each
140    /// region keeps its exact tensor width to avoid padded batch inference.
141    pub fn with_parallel(mut self, enable: bool) -> Self {
142        self.enable_parallel = enable;
143        self
144    }
145
146    /// Set minimum confidence threshold at result level
147    ///
148    /// Recognition results below this threshold will be filtered out.
149    /// Recommended values: 0.5 (lenient), 0.7 (balanced), 0.9 (strict)
150    pub fn with_min_result_confidence(mut self, threshold: f32) -> Self {
151        self.min_result_confidence = threshold;
152        self
153    }
154
155    /// Set minimum confidence threshold for orientation correction
156    pub fn with_ori_min_confidence(mut self, threshold: f32) -> Self {
157        self.ori_min_confidence = threshold;
158        self
159    }
160
161    /// Fast mode preset
162    pub fn fast() -> Self {
163        Self {
164            precision_mode: PrecisionMode::Low,
165            det_options: DetOptions::fast(),
166            ..Default::default()
167        }
168    }
169
170    /// GPU mode preset (Metal)
171    #[cfg(any(target_os = "macos", target_os = "ios"))]
172    pub fn gpu() -> Self {
173        Self {
174            backend: Backend::Metal,
175            ..Default::default()
176        }
177    }
178
179    /// GPU mode preset (OpenCL)
180    #[cfg(not(any(target_os = "macos", target_os = "ios")))]
181    pub fn gpu() -> Self {
182        Self {
183            backend: Backend::OpenCL,
184            ..Default::default()
185        }
186    }
187
188    fn to_inference_config(&self) -> InferenceConfig {
189        InferenceConfig {
190            thread_count: self.thread_count,
191            precision_mode: self.precision_mode,
192            backend: self.backend,
193            ..Default::default()
194        }
195    }
196}
197
198/// OCR engine
199///
200/// Encapsulates complete OCR pipeline, including text detection and recognition
201///
202/// # Example
203///
204/// ```ignore
205/// use ocr_rs::{OcrEngine, OcrEngineConfig};
206///
207/// // Create engine
208/// let engine = OcrEngine::new(
209///     "det_model.mnn",
210///     "rec_model.mnn",
211///     "ppocr_keys.txt",
212///     None,
213/// )?;
214///
215/// // Recognize image
216/// let image = image::open("test.jpg")?;
217/// let results = engine.recognize(&image)?;
218///
219/// for result in results {
220///     println!("{}: {:.2}", result.text, result.confidence);
221/// }
222/// ```
223pub struct OcrEngine {
224    det_model: DetModel,
225    rec_model: RecModel,
226    ori_model: Option<OriModel>,
227    config: OcrEngineConfig,
228}
229
230impl OcrEngine {
231    fn build_with_paths(
232        det_model_path: &Path,
233        rec_model_path: &Path,
234        charset_path: &Path,
235        ori_model_path: Option<&Path>,
236        config: Option<OcrEngineConfig>,
237    ) -> OcrResult<Self> {
238        let config = config.unwrap_or_default();
239        let inference_config = config.to_inference_config();
240
241        // Optimization: Directly move the configuration to avoid multiple clones
242        let det_options = config.det_options.clone();
243        let rec_options = config.rec_options.clone();
244        let ori_options = config.ori_options.clone();
245
246        let det_model = DetModel::from_file(det_model_path, Some(inference_config.clone()))?
247            .with_options(det_options);
248
249        let rec_model =
250            RecModel::from_file(rec_model_path, charset_path, Some(inference_config.clone()))?
251                .with_options(rec_options);
252
253        let ori_model = match ori_model_path {
254            Some(path) => {
255                Some(OriModel::from_file(path, Some(inference_config))?.with_options(ori_options))
256            }
257            None => None,
258        };
259
260        Ok(Self {
261            det_model,
262            rec_model,
263            ori_model,
264            config,
265        })
266    }
267
268    /// Create OCR engine from model files
269    ///
270    /// # Parameters
271    /// - `det_model_path`: Detection model file path
272    /// - `rec_model_path`: Recognition model file path
273    /// - `charset_path`: Charset file path
274    /// - `config`: Optional engine configuration
275    pub fn new(
276        det_model_path: impl AsRef<Path>,
277        rec_model_path: impl AsRef<Path>,
278        charset_path: impl AsRef<Path>,
279        config: Option<OcrEngineConfig>,
280    ) -> OcrResult<Self> {
281        Self::build_with_paths(
282            det_model_path.as_ref(),
283            rec_model_path.as_ref(),
284            charset_path.as_ref(),
285            None,
286            config,
287        )
288    }
289
290    /// Create OCR engine from model files with orientation model
291    pub fn new_with_ori(
292        det_model_path: impl AsRef<Path>,
293        rec_model_path: impl AsRef<Path>,
294        charset_path: impl AsRef<Path>,
295        ori_model_path: impl AsRef<Path>,
296        config: Option<OcrEngineConfig>,
297    ) -> OcrResult<Self> {
298        Self::build_with_paths(
299            det_model_path.as_ref(),
300            rec_model_path.as_ref(),
301            charset_path.as_ref(),
302            Some(ori_model_path.as_ref()),
303            config,
304        )
305    }
306
307    /// Create OCR engine from model bytes
308    pub fn from_bytes(
309        det_model_bytes: &[u8],
310        rec_model_bytes: &[u8],
311        charset_bytes: &[u8],
312        config: Option<OcrEngineConfig>,
313    ) -> OcrResult<Self> {
314        let config = config.unwrap_or_default();
315        let inference_config = config.to_inference_config();
316
317        // Optimization: Directly move the configuration to avoid multiple clones
318        let det_options = config.det_options.clone();
319        let rec_options = config.rec_options.clone();
320
321        let det_model = DetModel::from_bytes(det_model_bytes, Some(inference_config.clone()))?
322            .with_options(det_options);
323
324        let rec_model = RecModel::from_bytes_with_charset(
325            rec_model_bytes,
326            charset_bytes,
327            Some(inference_config.clone()),
328        )?
329        .with_options(rec_options);
330
331        Ok(Self {
332            det_model,
333            rec_model,
334            ori_model: None,
335            config,
336        })
337    }
338
339    /// Create OCR engine from model bytes with orientation model
340    pub fn from_bytes_with_ori(
341        det_model_bytes: &[u8],
342        rec_model_bytes: &[u8],
343        charset_bytes: &[u8],
344        ori_model_bytes: &[u8],
345        config: Option<OcrEngineConfig>,
346    ) -> OcrResult<Self> {
347        let config = config.unwrap_or_default();
348        let inference_config = config.to_inference_config();
349
350        let det_options = config.det_options.clone();
351        let rec_options = config.rec_options.clone();
352        let ori_options = config.ori_options.clone();
353
354        let det_model = DetModel::from_bytes(det_model_bytes, Some(inference_config.clone()))?
355            .with_options(det_options);
356
357        let rec_model = RecModel::from_bytes_with_charset(
358            rec_model_bytes,
359            charset_bytes,
360            Some(inference_config.clone()),
361        )?
362        .with_options(rec_options);
363
364        let ori_model = OriModel::from_bytes(ori_model_bytes, Some(inference_config))?
365            .with_options(ori_options);
366
367        Ok(Self {
368            det_model,
369            rec_model,
370            ori_model: Some(ori_model),
371            config,
372        })
373    }
374
375    /// Create detection-only engine
376    pub fn det_only(
377        det_model_path: impl AsRef<Path>,
378        config: Option<OcrEngineConfig>,
379    ) -> OcrResult<DetOnlyEngine> {
380        let config = config.unwrap_or_default();
381        let inference_config = config.to_inference_config();
382
383        let det_model = DetModel::from_file(det_model_path, Some(inference_config))?
384            .with_options(config.det_options);
385
386        Ok(DetOnlyEngine { det_model })
387    }
388
389    /// Create recognition-only engine
390    pub fn rec_only(
391        rec_model_path: impl AsRef<Path>,
392        charset_path: impl AsRef<Path>,
393        config: Option<OcrEngineConfig>,
394    ) -> OcrResult<RecOnlyEngine> {
395        let config = config.unwrap_or_default();
396        let inference_config = config.to_inference_config();
397
398        let rec_model = RecModel::from_file(rec_model_path, charset_path, Some(inference_config))?
399            .with_options(config.rec_options);
400
401        Ok(RecOnlyEngine { rec_model })
402    }
403
404    /// Perform complete OCR recognition
405    ///
406    /// # Parameters
407    /// - `image`: Input image
408    ///
409    /// # Returns
410    /// List of OCR results, each result contains text, confidence and bounding box
411    pub fn recognize(&self, image: &DynamicImage) -> OcrResult<Vec<OcrResult_>> {
412        // 0. Orientation correction for full image (optional)
413        let corrected_image = if let Some(ori_model) = self.ori_model.as_ref() {
414            self.correct_orientation_with_model(ori_model, image.clone())
415        } else {
416            image.clone()
417        };
418
419        // 1. Detect text regions
420        let boxes = self.det_model.detect_expanded(&corrected_image)?;
421
422        if boxes.is_empty() {
423            return Ok(Vec::new());
424        }
425
426        // 2. Render regions directly into recognition tensors. Large pages use
427        // exact-width tensors to avoid padding every line to the widest region.
428        let rec_results =
429            match select_region_recognition_strategy(self.config.enable_parallel, boxes.len()) {
430                RegionRecognitionStrategy::Batch => {
431                    self.rec_model.recognize_regions(&corrected_image, &boxes)?
432                }
433                RegionRecognitionStrategy::ExactWidthParallel => self
434                    .rec_model
435                    .recognize_regions_exact_parallel(&corrected_image, &boxes)?,
436            };
437
438        // 3. Combine results and filter low confidence
439        let results: Vec<OcrResult_> = rec_results
440            .into_iter()
441            .zip(boxes)
442            .filter(|(rec, _)| {
443                !rec.text.is_empty() && rec.confidence >= self.config.min_result_confidence
444            })
445            .map(|(rec, bbox)| OcrResult_::new(rec.text, rec.confidence, bbox))
446            .collect();
447
448        Ok(results)
449    }
450
451    /// Perform detection only
452    pub fn detect(&self, image: &DynamicImage) -> OcrResult<Vec<TextBox>> {
453        self.det_model.detect(image)
454    }
455
456    /// Perform recognition only (requires pre-cropped text line images)
457    pub fn recognize_text(&self, image: &DynamicImage) -> OcrResult<RecognitionResult> {
458        self.rec_model.recognize(image)
459    }
460
461    /// Batch recognize text line images
462    pub fn recognize_batch(&self, images: &[DynamicImage]) -> OcrResult<Vec<RecognitionResult>> {
463        self.rec_model.recognize_batch(images)
464    }
465
466    /// Get orientation model reference (if enabled)
467    pub fn ori_model(&self) -> Option<&OriModel> {
468        self.ori_model.as_ref()
469    }
470
471    /// Get detection model reference
472    pub fn det_model(&self) -> &DetModel {
473        &self.det_model
474    }
475
476    /// Get recognition model reference
477    pub fn rec_model(&self) -> &RecModel {
478        &self.rec_model
479    }
480
481    /// Get configuration
482    pub fn config(&self) -> &OcrEngineConfig {
483        &self.config
484    }
485
486    fn correct_orientation_with_model(
487        &self,
488        ori_model: &OriModel,
489        image: DynamicImage,
490    ) -> DynamicImage {
491        let result = match ori_model.classify(&image) {
492            Ok(result) => result,
493            Err(_) => return image,
494        };
495
496        if !result.is_valid(self.config.ori_min_confidence) {
497            return image;
498        }
499
500        if result.angle.rem_euclid(360) == 0 {
501            return image;
502        }
503
504        rotate_by_angle(&image, result.angle)
505    }
506}
507
508/// Builder for OCR engine
509pub struct OcrEngineBuilder {
510    det_model_path: Option<PathBuf>,
511    rec_model_path: Option<PathBuf>,
512    charset_path: Option<PathBuf>,
513    ori_model_path: Option<PathBuf>,
514    config: Option<OcrEngineConfig>,
515}
516
517impl Default for OcrEngineBuilder {
518    fn default() -> Self {
519        Self::new()
520    }
521}
522
523impl OcrEngineBuilder {
524    /// Create a new builder
525    pub fn new() -> Self {
526        Self {
527            det_model_path: None,
528            rec_model_path: None,
529            charset_path: None,
530            ori_model_path: None,
531            config: None,
532        }
533    }
534
535    /// Set detection model path
536    pub fn with_det_model_path(mut self, path: impl AsRef<Path>) -> Self {
537        self.det_model_path = Some(path.as_ref().to_path_buf());
538        self
539    }
540
541    /// Set recognition model path
542    pub fn with_rec_model_path(mut self, path: impl AsRef<Path>) -> Self {
543        self.rec_model_path = Some(path.as_ref().to_path_buf());
544        self
545    }
546
547    /// Set charset path
548    pub fn with_charset_path(mut self, path: impl AsRef<Path>) -> Self {
549        self.charset_path = Some(path.as_ref().to_path_buf());
550        self
551    }
552
553    /// Set orientation model path
554    pub fn with_ori_model_path(mut self, path: impl AsRef<Path>) -> Self {
555        self.ori_model_path = Some(path.as_ref().to_path_buf());
556        self
557    }
558
559    /// Set engine configuration
560    pub fn with_config(mut self, config: OcrEngineConfig) -> Self {
561        self.config = Some(config);
562        self
563    }
564
565    /// Build OCR engine
566    pub fn build(self) -> OcrResult<OcrEngine> {
567        let det_model_path = self
568            .det_model_path
569            .ok_or_else(|| OcrError::InvalidParameter("Missing det_model_path".to_string()))?;
570        let rec_model_path = self
571            .rec_model_path
572            .ok_or_else(|| OcrError::InvalidParameter("Missing rec_model_path".to_string()))?;
573        let charset_path = self
574            .charset_path
575            .ok_or_else(|| OcrError::InvalidParameter("Missing charset_path".to_string()))?;
576
577        OcrEngine::build_with_paths(
578            det_model_path.as_path(),
579            rec_model_path.as_path(),
580            charset_path.as_path(),
581            self.ori_model_path.as_deref(),
582            self.config,
583        )
584    }
585}
586
587/// Detection-only engine
588pub struct DetOnlyEngine {
589    det_model: DetModel,
590}
591
592impl DetOnlyEngine {
593    /// Detect text regions in image
594    pub fn detect(&self, image: &DynamicImage) -> OcrResult<Vec<TextBox>> {
595        self.det_model.detect(image)
596    }
597
598    /// Detect and return cropped images
599    pub fn detect_and_crop(&self, image: &DynamicImage) -> OcrResult<Vec<(DynamicImage, TextBox)>> {
600        self.det_model.detect_and_crop(image)
601    }
602
603    /// Get detection model reference
604    pub fn model(&self) -> &DetModel {
605        &self.det_model
606    }
607}
608
609/// Recognition-only engine
610pub struct RecOnlyEngine {
611    rec_model: RecModel,
612}
613
614impl RecOnlyEngine {
615    /// Recognize a single image
616    pub fn recognize(&self, image: &DynamicImage) -> OcrResult<RecognitionResult> {
617        self.rec_model.recognize(image)
618    }
619
620    /// Return text only
621    pub fn recognize_text(&self, image: &DynamicImage) -> OcrResult<String> {
622        self.rec_model.recognize_text(image)
623    }
624
625    /// Batch recognition
626    pub fn recognize_batch(&self, images: &[DynamicImage]) -> OcrResult<Vec<RecognitionResult>> {
627        self.rec_model.recognize_batch(images)
628    }
629
630    /// Get recognition model reference
631    pub fn model(&self) -> &RecModel {
632        &self.rec_model
633    }
634}
635
636/// Convenience function: recognize from file
637///
638/// # Example
639///
640/// ```ignore
641/// let results = ocr_rs::ocr_file(
642///     "test.jpg",
643///     "det_model.mnn",
644///     "rec_model.mnn",
645///     "ppocr_keys.txt",
646/// )?;
647/// ```
648pub fn ocr_file(
649    image_path: impl AsRef<Path>,
650    det_model_path: impl AsRef<Path>,
651    rec_model_path: impl AsRef<Path>,
652    charset_path: impl AsRef<Path>,
653) -> OcrResult<Vec<OcrResult_>> {
654    let image = image::open(image_path)?;
655    let engine = OcrEngine::new(det_model_path, rec_model_path, charset_path, None)?;
656    engine.recognize(&image)
657}
658
659/// Convenience function: recognize from file with orientation model
660pub fn ocr_file_with_ori(
661    image_path: impl AsRef<Path>,
662    det_model_path: impl AsRef<Path>,
663    rec_model_path: impl AsRef<Path>,
664    charset_path: impl AsRef<Path>,
665    ori_model_path: impl AsRef<Path>,
666) -> OcrResult<Vec<OcrResult_>> {
667    let image = image::open(image_path)?;
668    let engine = OcrEngine::new_with_ori(
669        det_model_path,
670        rec_model_path,
671        charset_path,
672        ori_model_path,
673        None,
674    )?;
675    engine.recognize(&image)
676}
677
678fn rotate_by_angle(image: &DynamicImage, angle: i32) -> DynamicImage {
679    // The model reports rotation from horizontal; rotate back to correct.
680    match angle.rem_euclid(360) {
681        90 => DynamicImage::ImageRgb8(image::imageops::rotate270(&image.to_rgb8())),
682        180 => DynamicImage::ImageRgb8(image::imageops::rotate180(&image.to_rgb8())),
683        270 => DynamicImage::ImageRgb8(image::imageops::rotate90(&image.to_rgb8())),
684        _ => image.clone(),
685    }
686}
687
688#[cfg(test)]
689mod tests {
690    use super::*;
691
692    #[test]
693    fn parallel_strategy_uses_exact_width_only_for_multi_line_pages() {
694        assert_eq!(
695            select_region_recognition_strategy(true, 5),
696            RegionRecognitionStrategy::ExactWidthParallel
697        );
698        assert_eq!(
699            select_region_recognition_strategy(true, 4),
700            RegionRecognitionStrategy::Batch
701        );
702        assert_eq!(
703            select_region_recognition_strategy(false, 12),
704            RegionRecognitionStrategy::Batch
705        );
706    }
707
708    #[test]
709    fn test_ocr_result() {
710        let bbox = TextBox::new(imageproc::rect::Rect::at(0, 0).of_size(100, 20), 0.9);
711        let result = OcrResult_::new("Hello".to_string(), 0.95, bbox);
712
713        assert_eq!(result.text, "Hello");
714        assert_eq!(result.confidence, 0.95);
715    }
716}
ocr_rs/engine.rs

ocr_rs/
engine.rs
