oneocr_rs/
ocr_engine.rs

1use crate::errors::OneOcrError;
2use crate::ffi::{
3    CreateOcrInitOptions, CreateOcrPipeline, CreateOcrProcessOptions,
4    OcrInitOptionsSetUseModelDelayLoad, OcrProcessOptionsGetMaxRecognitionLineCount,
5    OcrProcessOptionsGetResizeResolution, OcrProcessOptionsSetMaxRecognitionLineCount,
6    OcrProcessOptionsSetResizeResolution, RawImage, ReleaseOcrInitOptions, ReleaseOcrPipeline,
7    ReleaseOcrProcessOptions, RunOcrPipeline,
8};
9use crate::ocr_result::OcrResult;
10use crate::{ImageInput, ONE_OCR_MODEL_FILE_NAME, ONE_OCR_MODEL_KEY, OcrOptions};
11use image::{DynamicImage, ImageBuffer, Rgba};
12use std::ffi::{CString, c_void};
13use std::ptr;
14
15// Macros
16use crate::check_ocr_call;
17
18/// The `OcrEngine` struct represents the OneOcr processing engine.
19#[derive(Debug)]
20pub struct OcrEngine {
21    init_options: *mut c_void,
22    pipeline: *mut c_void,
23    process_options: *mut c_void,
24    ocr_options: OcrOptions,
25}
26
27impl OcrEngine {
28    /// Creates a new instance of the OCR engine with specified options.
29    /// This function loads the necessary library and initializes the OCR pipeline with the provided options.
30    pub fn new_with_options(ocr_options: OcrOptions) -> Result<Self, OneOcrError> {
31        let mut init_options: *mut c_void = ptr::null_mut();
32        check_ocr_call!(
33            unsafe { CreateOcrInitOptions(&mut init_options) },
34            "Failed to create init options"
35        );
36
37        // Disable model delay load
38        check_ocr_call!(
39            unsafe { OcrInitOptionsSetUseModelDelayLoad(init_options, 0) },
40            "Failed to set model delay load"
41        );
42
43        let model_path = Self::get_model_path()?;
44        let model_path_cstr = CString::new(model_path).map_err(|e| {
45            OneOcrError::ModelFileLoadError(format!("Failed to convert model path to CString: {e}"))
46        })?;
47
48        let key_cstr = CString::new(ONE_OCR_MODEL_KEY).map_err(|e| {
49            OneOcrError::InvalidModelKey(format!("Failed to convert model key to CString: {e}"))
50        })?;
51
52        let mut pipeline: *mut c_void = ptr::null_mut();
53        check_ocr_call!(
54            unsafe {
55                CreateOcrPipeline(
56                    model_path_cstr.as_ptr(),
57                    key_cstr.as_ptr(),
58                    init_options,
59                    &mut pipeline,
60                )
61            },
62            "Failed to create OCR pipeline"
63        );
64
65        let mut process_options: *mut c_void = ptr::null_mut();
66        check_ocr_call!(
67            unsafe { CreateOcrProcessOptions(&mut process_options) },
68            "Failed to create OCR process options"
69        );
70
71        check_ocr_call!(
72            unsafe {
73                OcrProcessOptionsSetMaxRecognitionLineCount(
74                    process_options,
75                    ocr_options.max_recognition_line_count,
76                )
77            },
78            "Failed to set max recognition line count"
79        );
80
81        check_ocr_call!(
82            unsafe {
83                OcrProcessOptionsSetResizeResolution(
84                    process_options,
85                    ocr_options.resize_resolution.width,
86                    ocr_options.resize_resolution.height,
87                )
88            },
89            "Failed to set resize resolution"
90        );
91
92        Ok(Self {
93            init_options,
94            pipeline,
95            process_options,
96            ocr_options,
97        })
98    }
99
100    /// Creates a new instance of the OCR engine with default options.
101    /// This function loads the necessary library and initializes the OCR pipeline.
102    pub fn new() -> Result<Self, OneOcrError> {
103        Self::new_with_options(OcrOptions::default())
104    }
105
106    /// Retrieves the maximum number of lines that can be recognized.
107    /// Default is 100.
108    pub fn get_max_recognition_line_count(&self) -> Result<i32, OneOcrError> {
109        let mut count: i32 = 0;
110        check_ocr_call!(
111            unsafe {
112                OcrProcessOptionsGetMaxRecognitionLineCount(self.process_options, &mut count)
113            },
114            "Failed to get max recognition line count"
115        );
116        Ok(count)
117    }
118
119    /// Sets the maximum number of lines that can be recognized.
120    /// Default is 100, range is 0-1000.
121    pub fn set_max_recognition_line_count(&self, count: i32) -> Result<(), OneOcrError> {
122        check_ocr_call!(
123            unsafe { OcrProcessOptionsSetMaxRecognitionLineCount(self.process_options, count) },
124            "Failed to set max recognition line count"
125        );
126        Ok(())
127    }
128
129    /// Retrieves the maximum internal resize resolution.
130    ///
131    /// The `resize resolution` defines the maximum dimensions to which an image will be automatically scaled internally before OCR processing.
132    /// It’s a performance and accuracy trade-off rather than a restriction on the original image’s resolution.
133    ///
134    /// Default is 1152*768.
135    pub fn get_resize_resolution(&self) -> Result<(i64, i64), OneOcrError> {
136        let mut width: i64 = 0;
137        let mut height: i64 = 0;
138        check_ocr_call!(
139            unsafe {
140                OcrProcessOptionsGetResizeResolution(self.process_options, &mut width, &mut height)
141            },
142            "Failed to get resize resolution"
143        );
144        Ok((width, height))
145    }
146
147    /// Sets the maximum internal resize resolution.
148    ///
149    /// The `resize resolution` defines the maximum dimensions to which an image will be automatically scaled internally before OCR processing.
150    /// It’s a performance and accuracy trade-off rather than a restriction on the original image’s resolution.
151    ///
152    /// The maximum resolution is 1152*768.
153    pub fn set_resize_resolution(&self, width: i32, height: i32) -> Result<(), OneOcrError> {
154        check_ocr_call!(
155            unsafe { OcrProcessOptionsSetResizeResolution(self.process_options, width, height) },
156            "Failed to set resize resolution"
157        );
158        Ok(())
159    }
160
161    /// Run OCR processing on an image.
162    ///
163    /// This method accepts various input types through the `ImageInput` enum
164    /// and allows configuration through `OcrOptions`.
165    ///
166    /// # Arguments
167    ///
168    /// * `input` - The image input source (file path, image buffer, or dynamic image)
169    ///
170    /// # Returns
171    ///
172    /// Returns an `OcrResult` containing the recognized text and associated metadata,
173    /// or an error if the OCR processing fails.
174    ///
175    /// # Examples
176    ///
177    /// ```no_run
178    /// use oneocr_rs::{OcrEngine, OcrOptions, ImageInput};
179    /// use std::path::Path;
180    /// let engine = OcrEngine::new().unwrap();
181    ///
182    /// // Process from file path
183    /// let result = engine.run(Path::new("image.jpg").into()).unwrap();
184    /// ```
185    ///
186    /// ```ignore
187    /// // Process from in-memory image buffer
188    /// let img_buffer: ImageBuffer<Rgba<u8>, Vec<u8>> = capture_screenshot(); // Your screenshot function
189    /// let result = engine.run(img_buffer.into()).unwrap();
190    /// ```
191    pub fn run(&self, input: ImageInput) -> Result<OcrResult, OneOcrError> {
192        let img_rgba = self.load_image(input)?;
193        self.run_ocr_on_rgba_image(&img_rgba, self.ocr_options.include_word_level_details)
194    }
195
196    /// Loads an image from various input sources and converts it to RGBA format.
197    fn load_image(&self, input: ImageInput) -> Result<ImageBuffer<Rgba<u8>, Vec<u8>>, OneOcrError> {
198        match input {
199            ImageInput::FilePath(path) => {
200                let img = image::open(path)?;
201                Ok(self.convert_to_rgba(img))
202            }
203            ImageInput::Buffer(buffer) => Ok(buffer),
204            ImageInput::Dynamic(img) => Ok(self.convert_to_rgba(img)),
205        }
206    }
207
208    /// Converts a DynamicImage to RGBA format.
209    fn convert_to_rgba(&self, img: DynamicImage) -> ImageBuffer<Rgba<u8>, Vec<u8>> {
210        match img {
211            DynamicImage::ImageRgba8(i) => i,
212            _ => img.to_rgba8(),
213        }
214    }
215
216    /// Performs OCR on an RGBA image buffer.
217    fn run_ocr_on_rgba_image(
218        &self,
219        img_rgba: &ImageBuffer<Rgba<u8>, Vec<u8>>,
220        word_level_detail: bool,
221    ) -> Result<OcrResult, OneOcrError> {
222        let (rows, cols) = (img_rgba.height() as i32, img_rgba.width() as i32);
223        let step = (img_rgba.sample_layout().height_stride) as i64;
224        let data_ptr = img_rgba.as_ptr() as i64;
225        let image = RawImage {
226            t: 3, // RGBA format identifier expected by the C API
227            col: cols,
228            row: rows,
229            _unk: 0,
230            step,
231            data_ptr,
232        };
233
234        let mut ocr_result: *mut c_void = ptr::null_mut();
235        check_ocr_call!(
236            unsafe { RunOcrPipeline(self.pipeline, &image, self.process_options, &mut ocr_result) },
237            "Failed to run OCR pipeline"
238        );
239
240        OcrResult::new(ocr_result, word_level_detail)
241    }
242
243    /// Retrieves the path to the model file.
244    fn get_model_path() -> Result<String, OneOcrError> {
245        let exe_path = std::env::current_exe().map_err(|e| {
246            OneOcrError::ModelFileLoadError(format!("Failed to get current executable path: {e}"))
247        })?;
248        let model_path_buf = exe_path
249            .parent()
250            .ok_or_else(|| {
251                OneOcrError::ModelFileLoadError(
252                    "Failed to get parent directory of current executable".to_string(),
253                )
254            })?
255            .join(ONE_OCR_MODEL_FILE_NAME);
256        let model_path_string = model_path_buf.to_string_lossy().to_string();
257
258        Ok(model_path_string)
259    }
260}
261
262impl Drop for OcrEngine {
263    fn drop(&mut self) {
264        unsafe {
265            ReleaseOcrPipeline(self.pipeline);
266            ReleaseOcrInitOptions(self.init_options);
267            ReleaseOcrProcessOptions(self.process_options);
268        };
269    }
270}
oneocr_rs/ocr_engine.rs

oneocr_rs/
ocr_engine.rs
