Result Iterators

Cargo.toml

@@ -8,6 +8,7 @@ description   = "Higher-level bindings for Tesseract OCR"
 license       = "MIT"
 keywords      = ["tesseract", "OCR", "bindings"]
 categories    = ["api-bindings", "multimedia::images"]
+edition = "2021"
 
 [dependencies]
 tesseract-sys = "~0.5"

src/iterator.rs

@@ -0,0 +1,147 @@
+use tesseract_plumbing::Text;
+use tesseract_sys::{PageIteratorLevel, PolyBlockType};
+
+/// Iterate through the OCR results of an image
+///
+/// As the results are ephemeral, the `next` function only yields references to these results.
+/// ```
+/// use tesseract::{Tesseract, iterator::TextlineResult};
+/// use crate::tesseract::TesseractIteratorResult;
+/// use tesseract_sys::PolyBlockType;
+///
+/// let mut tesseract = Tesseract::new(None, Some("eng")).unwrap();
+/// let mut tesseract = tesseract.set_image("./img.png").expect("Failed to set image");
+/// let mut tesseract = tesseract.recognize().unwrap();
+/// let mut tesseract_iterator = tesseract.iterator().unwrap();
+
+/// while let Some(textline) = tesseract_iterator.next::<TextlineResult>(None) {
+///     assert_eq!(PolyBlockType::PT_FLOWING_TEXT ,textline.block_type());
+/// }
+/// ```
+pub struct ResultIterator(tesseract_plumbing::ResultIterator, bool);
+
+pub struct BlockResult<'a>(&'a tesseract_plumbing::ResultIterator);
+pub struct ParagraphResult<'a>(&'a tesseract_plumbing::ResultIterator);
+pub struct TextlineResult<'a>(&'a tesseract_plumbing::ResultIterator);
+pub struct WordResult<'a>(&'a tesseract_plumbing::ResultIterator);
+pub struct SymbolResult<'a>(&'a tesseract_plumbing::ResultIterator);
+
+impl ResultIterator {
+    pub(crate) fn new(iterator: tesseract_plumbing::ResultIterator) -> Self {
+        Self(iterator, true)
+    }
+    pub fn next<'a, T>(&'a mut self, limit: Option<PageIteratorLevel>) -> Option<T>
+    where
+        T: TesseractIteratorResult<'a>,
+    {
+        let p = *self.0.as_ref();
+        if self.1 {
+            self.1 = false;
+            return Some(T::from(&self.0));
+        }
+        let end_of_page_reached =
+            unsafe { tesseract_sys::TessPageIteratorNext(p.cast(), T::LEVEL as u32) == 0 };
+        if end_of_page_reached {
+            return None;
+        }
+        Some(T::from(&self.0))
+    }
+}
+
+pub trait TesseractIteratorResult<'a>
+where
+    Self: From<&'a tesseract_plumbing::ResultIterator>
+        + AsRef<tesseract_plumbing::ResultIterator>
+        + 'a,
+{
+    /// The equivalent PageIteratorLevel of this result
+    const LEVEL: PageIteratorLevel;
+    /// Get the text contained of the iteration result
+    fn get_text(&self) -> Option<Text> {
+        let c_str = unsafe {
+            tesseract_sys::TessResultIteratorGetUTF8Text(
+                self.as_ref().as_ref().cast(),
+                Self::LEVEL as u32,
+            )
+        };
+        if c_str.is_null() {
+            return None;
+        }
+        Some(unsafe { tesseract_plumbing::Text::new(c_str) })
+    }
+
+    /// Get the bounding box of the iteration result
+    fn bounding_box(&self) -> BoundingBox {
+        let mut left = 0;
+        let mut right = 0;
+        let mut top = 0;
+        let mut bottom = 0;
+        // TODO: Use this to verify
+        let _object_at_pos = unsafe {
+            tesseract_sys::TessPageIteratorBoundingBox(
+                self.as_ref().as_ref().cast(),
+                Self::LEVEL as u32,
+                &mut left,
+                &mut top,
+                &mut right,
+                &mut bottom,
+            )
+        };
+        BoundingBox {
+            left,
+            top,
+            right,
+            bottom,
+        }
+    }
+
+    /// Check what the current block type is
+    fn block_type(&self) -> PolyBlockType {
+        let block_type =
+            unsafe { tesseract_sys::TessPageIteratorBlockType(self.as_ref().as_ref().cast()) };
+        unsafe { std::mem::transmute(block_type) } // TODO: This doesn't check that the value is valid
+    }
+}
+
+#[derive(Debug)]
+pub struct BoundingBox {
+    left: i32,
+    top: i32,
+    right: i32,
+    bottom: i32,
+}
+
+/// All results implement the same basic functionality, but with slight differences in the PageIteratorLevel
+macro_rules! result_impls {
+    ($name:ident -> $level:expr $(, $($tts:tt)*)?) => {
+
+        impl<'a> TesseractIteratorResult<'a> for $name<'a>
+        where Self: From<&'a tesseract_plumbing::ResultIterator> + AsRef<tesseract_plumbing::ResultIterator> + 'a
+        {
+            const LEVEL: PageIteratorLevel = $level;
+        }
+
+        impl<'a> From<&'a tesseract_plumbing::ResultIterator> for $name<'a> {
+            fn from(value: &'a tesseract_plumbing::ResultIterator) -> $name<'a> {
+                $name(value)
+            }
+        }
+
+        impl AsRef<tesseract_plumbing::ResultIterator> for $name<'_> {
+            fn as_ref(&self) -> &tesseract_plumbing::ResultIterator {
+                self.0
+            }
+        }
+
+        result_impls!($($($tts)*)?);
+    };
+    () => {};
+}
+
+result_impls!(
+    BlockResult -> PageIteratorLevel::RIL_BLOCK,
+    ParagraphResult -> PageIteratorLevel::RIL_PARA,
+    TextlineResult -> PageIteratorLevel::RIL_TEXTLINE,
+    WordResult -> PageIteratorLevel::RIL_WORD,
+    SymbolResult -> PageIteratorLevel::RIL_SYMBOL
+);

src/lib.rs

@@ -7,8 +7,10 @@ use std::ffi::CString;
 use std::ffi::NulError;
 use std::os::raw::c_int;
 use std::str;
+pub mod iterator;
 mod page_seg_mode;
 
+pub use iterator::{ResultIterator, TesseractIteratorResult};
 pub use page_seg_mode::PageSegMode;
 
 use self::tesseract_sys::{
@@ -237,6 +239,21 @@ impl Tesseract {
     pub fn set_page_seg_mode(&mut self, mode: PageSegMode) {
         self.0.set_page_seg_mode(mode.as_tess_page_seg_mode());
     }
+
+    pub fn iterator(&mut self) -> Option<ResultIterator> {
+        let result_iter =
+            unsafe { tesseract_sys::TessBaseAPIGetIterator(self.0.get_raw().cast_mut()) };
+        if result_iter.is_null() {
+            return None;
+        }
+        Some(ResultIterator::new(
+            tesseract_plumbing::ResultIterator::new(result_iter),
+        ))
+    }
+
+    pub unsafe fn get_raw(&mut self) -> *const tesseract_sys::TessBaseAPI {
+        self.0.get_raw()
+    }
 }
 
 pub fn ocr(filename: &str, language: &str) -> Result<String, TesseractError> {
@@ -259,81 +276,3 @@ pub fn ocr_from_frame(
         .recognize()?
         .get_text()?)
 }
-
-#[test]
-fn ocr_test() -> Result<(), TesseractError> {
-    assert_eq!(
-        ocr("img.png", "eng")?,
-        include_str!("../img.txt").to_string()
-    );
-    Ok(())
-}
-
-#[test]
-fn ocr_from_frame_test() -> Result<(), TesseractError> {
-    assert_eq!(
-        ocr_from_frame(include_bytes!("../img.tiff"), 2256, 324, 3, 2256 * 3, "eng")?,
-        include_str!("../img.txt").to_string()
-    );
-    Ok(())
-}
-
-#[test]
-fn ocr_from_mem_with_ppi() -> Result<(), TesseractError> {
-    let mut cube = Tesseract::new(None, Some("eng"))?
-        .set_image_from_mem(include_bytes!("../img.tiff"))?
-        .set_source_resolution(70);
-    assert_eq!(&cube.get_text()?, include_str!("../img.txt"));
-    Ok(())
-}
-
-#[test]
-fn expanded_test() -> Result<(), TesseractError> {
-    let mut cube = Tesseract::new(None, Some("eng"))?
-        .set_image("img.png")?
-        .set_variable("tessedit_char_blacklist", "z")?
-        .recognize()?;
-    assert_eq!(&cube.get_text()?, include_str!("../img.txt"));
-    Ok(())
-}
-
-#[test]
-fn hocr_test() -> Result<(), TesseractError> {
-    let mut cube = Tesseract::new(None, Some("eng"))?.set_image("img.png")?;
-    assert!(&cube.get_hocr_text(0)?.contains("<div class='ocr_page'"));
-    Ok(())
-}
-
-#[test]
-#[ignore] // Many systems do not have legacy Tesseract data available
-fn oem_test() -> Result<(), TesseractError> {
-    let only_tesseract_str =
-        Tesseract::new_with_oem(None, Some("eng"), OcrEngineMode::TesseractOnly)?
-            .set_image("img.png")?
-            .recognize()?
-            .get_text()?;
-
-    let only_lstm_str = Tesseract::new_with_oem(None, Some("eng"), OcrEngineMode::LstmOnly)?
-        .set_image("img.png")?
-        .recognize()?
-        .get_text()?;
-
-    assert_ne!(only_tesseract_str, only_lstm_str);
-    Ok(())
-}
-
-#[test]
-fn oem_ltsm_only_test() -> Result<(), TesseractError> {
-    let only_lstm_str = Tesseract::new_with_oem(None, Some("eng"), OcrEngineMode::LstmOnly)?
-        .set_image("img.png")?
-        .recognize()?
-        .get_text()?;
-
-    assert_eq!(only_lstm_str, include_str!("../img.txt"));
-    Ok(())
-}
-
-#[test]
-fn initialize_with_none() {
-    assert!(Tesseract::new(None, None).is_ok());
-}

tests/iterations.rs

@@ -0,0 +1,54 @@
+use tesseract::iterator::BlockResult;
+use tesseract::{iterator::TextlineResult, Tesseract};
+use tesseract::{TesseractError, TesseractIteratorResult};
+use tesseract_sys::PolyBlockType;
+
+fn init_tesseract(path: &str) -> Result<Tesseract, TesseractError> {
+    let tesseract = Tesseract::new(None, Some("eng"))?;
+    let tesseract = tesseract.set_image(path).expect("Failed to set image");
+    Ok(tesseract)
+}
+
+#[test]
+fn iterate_textlines() -> Result<(), TesseractError> {
+    let tesseract = init_tesseract("./img.png")?;
+    let mut tesseract = tesseract.recognize()?;
+
+    let mut tesseract_iterator = tesseract.iterator().unwrap();
+    let text = include_str!("../img.txt");
+    let mut lines = text.lines();
+    while let Some(textline) = tesseract_iterator.next::<TextlineResult>(None) {
+        assert_eq!(PolyBlockType::PT_FLOWING_TEXT, textline.block_type());
+        assert_eq!(
+            textline
+                .get_text()
+                .expect("Textline had no text")
+                .to_string()
+                .trim(),
+            lines
+                .next()
+                .expect("expected image text had fewer lines than detected in image")
+        );
+    }
+    Ok(())
+}
+
+#[test]
+fn iterate_block() -> Result<(), TesseractError> {
+    let tesseract = init_tesseract("./img.png")?;
+    let mut tesseract = tesseract.recognize()?;
+
+    let mut tesseract_iterator = tesseract.iterator().unwrap();
+    let text = include_str!("../img.txt");
+    while let Some(textline) = tesseract_iterator.next::<BlockResult>(None) {
+        assert_eq!(PolyBlockType::PT_FLOWING_TEXT, textline.block_type());
+        assert_eq!(
+            textline
+                .get_text()
+                .expect("Textline had no text")
+                .to_string(),
+            text
+        );
+    }
+    Ok(())
+}