Skip to main content

tesseract_ocr_static/
image.rs

1use core::ptr::NonNull;
2
3use crate::InvalidImage;
4
5fn rgb_to_rgba([r, g, b]: [u8; 3]) -> u32 {
6    u32::from(r) | (u32::from(g) << 8) | (u32::from(b) << 16)
7}
8
9fn rgba_to_rgba([r, g, b, a]: [u8; 4]) -> u32 {
10    u32::from(r) | (u32::from(g) << 8) | (u32::from(b) << 16) | (u32::from(a) << 24)
11}
12
13/// Tesseract-specific image.
14pub struct Image {
15    pub(crate) ptr: NonNull<c::PIX>,
16}
17
18impl Image {
19    /// Create an image from raw RGB bytes.
20    pub fn from_rgb(width: u32, height: u32, rgb: &[u8]) -> Result<Self, InvalidImage> {
21        assert!(
22            width <= i32::MAX as u32
23                && height <= i32::MAX as u32
24                && u64::from(width) * u64::from(height) * 3 == rgb.len() as u64
25        );
26        let bits_per_pixel = 32;
27        let mut image = Self::new(width, height, bits_per_pixel)?;
28        let pixels = image.as_pixels_mut();
29        for (pixel, rgb) in pixels.iter_mut().zip(rgb.chunks_exact(3)) {
30            *pixel = rgb_to_rgba([rgb[0], rgb[1], rgb[2]]);
31        }
32        Ok(image)
33    }
34
35    /// Create an image from raw RGBA bytes.
36    pub fn from_rgba(width: u32, height: u32, rgba: &[u8]) -> Result<Self, InvalidImage> {
37        assert!(
38            width <= i32::MAX as u32
39                && height <= i32::MAX as u32
40                && u64::from(width) * u64::from(height) * 4 == rgba.len() as u64
41        );
42        let bits_per_pixel = 32;
43        let mut image = Self::new(width, height, bits_per_pixel)?;
44        let pixels = image.as_pixels_mut();
45        for (pixel, rgba) in pixels.iter_mut().zip(rgba.chunks_exact(4)) {
46            *pixel = rgba_to_rgba([rgba[0], rgba[1], rgba[2], rgba[3]]);
47        }
48        Ok(image)
49    }
50
51    fn new(width: u32, height: u32, bits_per_pixel: u32) -> Result<Self, InvalidImage> {
52        let ptr = unsafe { c::pixCreate(width as i32, height as i32, bits_per_pixel as i32) };
53        let ptr = NonNull::new(ptr).ok_or(InvalidImage)?;
54        Ok(Self { ptr })
55    }
56
57    fn as_pixels_mut(&mut self) -> &mut [u32] {
58        let (ptr, len) = self.get_raw_pixels();
59        unsafe { core::slice::from_raw_parts_mut(ptr.cast(), len) }
60    }
61
62    fn get_raw_pixels(&self) -> (*mut u32, usize) {
63        let data_ptr = unsafe { c::pixGetData(self.ptr.as_ptr()) };
64        let wpl = unsafe { c::pixGetWpl(self.ptr.as_ptr()) };
65        let width = unsafe { c::pixGetWidth(self.ptr.as_ptr()) };
66        let height = unsafe { c::pixGetHeight(self.ptr.as_ptr()) };
67        assert!(wpl >= 0 && height >= 0 && width == wpl && !data_ptr.is_null());
68        let len: usize = (wpl as usize)
69            .checked_mul(height as usize)
70            .expect("Overflow");
71        (data_ptr, len)
72    }
73
74    /// Returns width and height of the image.
75    pub fn dimensions(&self) -> (u32, u32) {
76        let mut width = 0;
77        let mut height = 0;
78        let mut depth = 0;
79        let ret =
80            unsafe { c::pixGetDimensions(self.ptr.as_ptr(), &mut width, &mut height, &mut depth) };
81        assert!(ret == 0);
82        let _ = depth;
83        (width as u32, height as u32)
84    }
85}
86
87impl Drop for Image {
88    fn drop(&mut self) {
89        unsafe { c::pixDestroy(&mut self.ptr.as_ptr()) };
90    }
91}
92
93impl Clone for Image {
94    fn clone(&self) -> Self {
95        let ptr = unsafe { c::pixClone(self.ptr.as_ptr()) };
96        let ptr = NonNull::new(ptr).expect("pixClone returned NULL");
97        Self { ptr }
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use crate::TextRecognizer;
105
106    #[test]
107    fn ocr_test_image() {
108        let text = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/data/text.txt"));
109        let rgb = image::ImageReader::open(concat!(env!("CARGO_MANIFEST_DIR"), "/data/text.png"))
110            .unwrap()
111            .decode()
112            .unwrap()
113            .into_rgb8();
114        let image = Image::from_rgb(rgb.width(), rgb.height(), rgb.as_raw()).unwrap();
115        let mut recognizer = TextRecognizer::new().unwrap();
116        let results = recognizer.recognize_text(&image).unwrap();
117        assert_eq!(text, results.get_utf8_text().as_str());
118    }
119}