Skip to content

Commit e2fab98

Browse files
author
BuildTools
committed
draft implementation of iterators
1 parent 90c765f commit e2fab98

File tree

3 files changed

+218
-0
lines changed

3 files changed

+218
-0
lines changed

‎src/iterator.rs

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
use tesseract_plumbing::Text;
2+
use tesseract_sys::{PageIteratorLevel, PolyBlockType};
3+
4+
/// Iterate through the OCR results of an image
5+
///
6+
/// As the results are ephemeral, the `next` function only yields references to these results.
7+
/// ```
8+
/// use tesseract::{Tesseract, iterator::TextlineResult};
9+
/// use crate::tesseract::TesseractIteratorResult;
10+
/// use tesseract_sys::PolyBlockType;
11+
///
12+
/// let mut tesseract = Tesseract::new(None, Some("eng")).unwrap();
13+
/// let mut tesseract = tesseract.set_image("./img.png").expect("Failed to set image");
14+
/// let mut tesseract = tesseract.recognize().unwrap();
15+
/// let mut tesseract_iterator = tesseract.iterator().unwrap();
16+
17+
/// while let Some(textline) = tesseract_iterator.next::<TextlineResult>(None) {
18+
/// assert_eq!(PolyBlockType::PT_FLOWING_TEXT ,textline.block_type());
19+
/// }
20+
/// ```
21+
pub struct ResultIterator(tesseract_plumbing::ResultIterator, bool);
22+
23+
pub struct BlockResult<'a>(&'a tesseract_plumbing::ResultIterator);
24+
pub struct ParagraphResult<'a>(&'a tesseract_plumbing::ResultIterator);
25+
pub struct TextlineResult<'a>(&'a tesseract_plumbing::ResultIterator);
26+
pub struct WordResult<'a>(&'a tesseract_plumbing::ResultIterator);
27+
pub struct SymbolResult<'a>(&'a tesseract_plumbing::ResultIterator);
28+
29+
impl ResultIterator {
30+
pub(crate) fn new(iterator: tesseract_plumbing::ResultIterator) -> Self {
31+
Self(iterator, true)
32+
}
33+
pub fn next<'a, T>(&'a mut self, limit: Option<PageIteratorLevel>) -> Option<T>
34+
where
35+
T: TesseractIteratorResult<'a>,
36+
{
37+
let p = *self.0.as_ref();
38+
if self.1 {
39+
self.1 = false;
40+
return Some(T::from(&self.0));
41+
}
42+
let end_of_page_reached =
43+
unsafe { tesseract_sys::TessPageIteratorNext(p.cast(), T::LEVEL as u32) == 0 };
44+
if end_of_page_reached {
45+
return None;
46+
}
47+
Some(T::from(&self.0))
48+
}
49+
}
50+
51+
pub trait TesseractIteratorResult<'a>
52+
where
53+
Self: From<&'a tesseract_plumbing::ResultIterator>
54+
+ AsRef<tesseract_plumbing::ResultIterator>
55+
+ 'a,
56+
{
57+
/// The equivalent PageIteratorLevel of this result
58+
const LEVEL: PageIteratorLevel;
59+
/// Get the text contained of the iteration result
60+
fn get_text(&self) -> Option<Text> {
61+
let c_str = unsafe {
62+
tesseract_sys::TessResultIteratorGetUTF8Text(
63+
self.as_ref().as_ref().cast(),
64+
Self::LEVEL as u32,
65+
)
66+
};
67+
if c_str.is_null() {
68+
return None;
69+
}
70+
Some(unsafe { tesseract_plumbing::Text::new(c_str) })
71+
}
72+
73+
/// Get the bounding box of the iteration result
74+
fn bounding_box(&self) -> BoundingBox {
75+
let mut left = 0;
76+
let mut right = 0;
77+
let mut top = 0;
78+
let mut bottom = 0;
79+
// TODO: Use this to verify
80+
let _object_at_pos = unsafe {
81+
tesseract_sys::TessPageIteratorBoundingBox(
82+
self.as_ref().as_ref().cast(),
83+
Self::LEVEL as u32,
84+
&mut left,
85+
&mut top,
86+
&mut right,
87+
&mut bottom,
88+
)
89+
};
90+
BoundingBox {
91+
left,
92+
top,
93+
right,
94+
bottom,
95+
}
96+
}
97+
98+
/// Check what the current block type is
99+
fn block_type(&self) -> PolyBlockType {
100+
let block_type =
101+
unsafe { tesseract_sys::TessPageIteratorBlockType(self.as_ref().as_ref().cast()) };
102+
unsafe { std::mem::transmute(block_type) } // TODO: This doesn't check that the value is valid
103+
}
104+
}
105+
106+
#[derive(Debug)]
107+
pub struct BoundingBox {
108+
left: i32,
109+
top: i32,
110+
right: i32,
111+
bottom: i32,
112+
}
113+
114+
/// All results implement the same basic functionality, but with slight differences in the PageIteratorLevel
115+
macro_rules! result_impls {
116+
($name:ident -> $level:expr $(, $($tts:tt)*)?) => {
117+
118+
impl<'a> TesseractIteratorResult<'a> for $name<'a>
119+
where Self: From<&'a tesseract_plumbing::ResultIterator> + AsRef<tesseract_plumbing::ResultIterator> + 'a
120+
{
121+
const LEVEL: PageIteratorLevel = $level;
122+
}
123+
124+
impl<'a> From<&'a tesseract_plumbing::ResultIterator> for $name<'a> {
125+
fn from(value: &'a tesseract_plumbing::ResultIterator) -> $name<'a> {
126+
$name(value)
127+
}
128+
}
129+
130+
impl AsRef<tesseract_plumbing::ResultIterator> for $name<'_> {
131+
fn as_ref(&self) -> &tesseract_plumbing::ResultIterator {
132+
self.0
133+
}
134+
}
135+
136+
result_impls!($($($tts)*)?);
137+
};
138+
() => {};
139+
}
140+
141+
result_impls!(
142+
BlockResult -> PageIteratorLevel::RIL_BLOCK,
143+
ParagraphResult -> PageIteratorLevel::RIL_PARA,
144+
TextlineResult -> PageIteratorLevel::RIL_TEXTLINE,
145+
WordResult -> PageIteratorLevel::RIL_WORD,
146+
SymbolResult -> PageIteratorLevel::RIL_SYMBOL
147+
);

‎src/lib.rs

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,10 @@ use std::ffi::CString;
77
use std::ffi::NulError;
88
use std::os::raw::c_int;
99
use std::str;
10+
pub mod iterator;
1011
mod page_seg_mode;
1112

13+
pub use iterator::{ResultIterator, TesseractIteratorResult};
1214
pub use page_seg_mode::PageSegMode;
1315

1416
use self::tesseract_sys::{
@@ -237,6 +239,21 @@ impl Tesseract {
237239
pub fn set_page_seg_mode(&mut self, mode: PageSegMode) {
238240
self.0.set_page_seg_mode(mode.as_tess_page_seg_mode());
239241
}
242+
243+
pub fn iterator(&mut self) -> Option<ResultIterator> {
244+
let result_iter =
245+
unsafe { tesseract_sys::TessBaseAPIGetIterator(self.0.get_raw().cast_mut()) };
246+
if result_iter.is_null() {
247+
return None;
248+
}
249+
Some(ResultIterator::new(
250+
tesseract_plumbing::ResultIterator::new(result_iter),
251+
))
252+
}
253+
254+
pub unsafe fn get_raw(&mut self) -> *const tesseract_sys::TessBaseAPI {
255+
self.0.get_raw()
256+
}
240257
}
241258

242259
pub fn ocr(filename: &str, language: &str) -> Result<String, TesseractError> {

‎tests/iterations.rs

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
use tesseract::iterator::BlockResult;
2+
use tesseract::{iterator::TextlineResult, Tesseract};
3+
use tesseract::{TesseractError, TesseractIteratorResult};
4+
use tesseract_sys::PolyBlockType;
5+
6+
fn init_tesseract(path: &str) -> Result<Tesseract, TesseractError> {
7+
let tesseract = Tesseract::new(None, Some("eng"))?;
8+
let tesseract = tesseract.set_image(path).expect("Failed to set image");
9+
Ok(tesseract)
10+
}
11+
12+
#[test]
13+
fn iterate_textlines() -> Result<(), TesseractError> {
14+
let tesseract = init_tesseract("./img.png")?;
15+
let mut tesseract = tesseract.recognize()?;
16+
17+
let mut tesseract_iterator = tesseract.iterator().unwrap();
18+
let text = include_str!("../img.txt");
19+
let mut lines = text.lines();
20+
while let Some(textline) = tesseract_iterator.next::<TextlineResult>(None) {
21+
assert_eq!(PolyBlockType::PT_FLOWING_TEXT, textline.block_type());
22+
assert_eq!(
23+
textline
24+
.get_text()
25+
.expect("Textline had no text")
26+
.to_string()
27+
.trim(),
28+
lines
29+
.next()
30+
.expect("expected image text had fewer lines than detected in image")
31+
);
32+
}
33+
Ok(())
34+
}
35+
36+
#[test]
37+
fn iterate_block() -> Result<(), TesseractError> {
38+
let tesseract = init_tesseract("./img.png")?;
39+
let mut tesseract = tesseract.recognize()?;
40+
41+
let mut tesseract_iterator = tesseract.iterator().unwrap();
42+
let text = include_str!("../img.txt");
43+
while let Some(textline) = tesseract_iterator.next::<BlockResult>(None) {
44+
assert_eq!(PolyBlockType::PT_FLOWING_TEXT, textline.block_type());
45+
assert_eq!(
46+
textline
47+
.get_text()
48+
.expect("Textline had no text")
49+
.to_string(),
50+
text
51+
);
52+
}
53+
Ok(())
54+
}

0 commit comments

Comments
 (0)