feat: structured ocr parse and paragraph search

Cargo.toml

@@ -10,7 +10,7 @@ in-memory = []
 [dependencies]
 chrono = "0.4"
 screenshots = "0.6.0"
-rusty-tesseract = { version = "*", git = "https://github.com/STRRL/rusty-tesseract", rev = "84418ef" }
+rusty-tesseract = { version = "1.1.7" }
 image = "0.24.6"
 anyhow = { version = "1.0", features = ["backtrace"] }
 async-trait = "0.1.71"
@@ -33,3 +33,4 @@ rust-embed = { version = "6.8.1", features = ["axum-ex"] }
 mime_guess = "2"
 imageproc = "0.23"
 colorsys = "0.6"
+xml = "0.8"

Makefile

@@ -11,3 +11,7 @@ frontend-embbed: webui-export
 .PHONY: webui-export
 webui-export:
 	cd webui && pnpm install && cp next.config.js.export next.config.js && npx next build && cp next.config.js.dev next.config.js
+
+.PHONY: clean
+clean:
+	cargo clean

src/analysis.rs

@@ -7,7 +7,7 @@ use serde::{Deserialize, Serialize};
 use crate::{
     image_archive::{ImageArchiver},
     ocr::{CharacterRecognizer, RecognizeItem},
-    repository::{EntityImage, EntityText, Repository},
+    repository::{EntityImage, EntityWord, Repository},
     screenshot::Screenshot,
 };
 
@@ -36,16 +36,16 @@ impl Analysis {
         let entity_image = self.repo.save_image(&entity_image).await?;
 
         let ocr_result: Vec<RecognizeItem> = self.ocr.recognize(&screenshot.image).await?;
-        let entity_texts: Vec<EntityText> = ocr_result
+        let entity_texts: Vec<EntityWord> = ocr_result
             .iter()
             .filter(|it| it.level == 5)
-            .filter_map(|it: &RecognizeItem| -> Option<EntityText> { it.try_into().ok() })
+            .filter_map(|it: &RecognizeItem| -> Option<EntityWord> { it.try_into().ok() })
             .map(|mut it| {
                 it.image_id = entity_image.id;
                 it
             })
             .collect();
-        self.repo.save_texts(&entity_texts).await?;
+        self.repo.save_words(&entity_texts).await?;
         Ok(())
     }
 
@@ -64,11 +64,11 @@ impl Analysis {
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct SearchResult {
     pub image_id: u32,
-    pub texts: Vec<EntityText>,
+    pub texts: Vec<EntityWord>,
 }
 
 impl SearchResult {
-    pub fn new(image_id: u32, texts: Vec<EntityText>) -> Self {
+    pub fn new(image_id: u32, texts: Vec<EntityWord>) -> Self {
         Self { image_id, texts }
     }
 }

src/http/service.rs

@@ -50,7 +50,7 @@ impl Service {
         let mut markups = Vec::new();
 
         for text_id in text_ids {
-            let entity_text = self.repo.get_text_by_id(*text_id).await?;
+            let entity_text = self.repo.get_word_by_id(*text_id).await?;
             let markup_box = MarkupBox::new(
                 entity_text.left,
                 entity_text.top,

src/ocr/hocr_parse.rs

@@ -0,0 +1,168 @@
+
+use xml::reader::{EventReader, XmlEvent};
+use super::model::{Paragraph, Line, Word, BoundingBox};
+
+pub fn parse_hocr_xml(hocr: &str) -> Vec<Paragraph> {
+    let mut result = Vec::new();
+    let reader = EventReader::new(hocr.as_bytes());
+
+    let mut current_par: Option<Paragraph> = None;
+    let mut current_line: Option<Line> = None;
+    let mut current_word: Option<Word> = None;
+
+    for event in reader {
+        match event {
+            Ok(XmlEvent::StartElement {
+                name,
+                attributes,
+                namespace,
+            }) => {
+                let par = attributes
+                    .iter()
+                    .find(|item| item.name.local_name == "class" && item.value.contains("ocr_par"));
+                if par.is_some() {
+                    let id = attributes
+                        .iter()
+                        .find(|item| item.name.local_name == "id")
+                        .unwrap()
+                        .value
+                        .clone();
+                    let title = attributes
+                        .iter()
+                        .find(|item| item.name.local_name == "title")
+                        .unwrap()
+                        .value
+                        .clone();
+                    let bbox = title
+                        .split(";")
+                        .take(1)
+                        .collect::<String>()
+                        .split_whitespace()
+                        .skip(1)
+                        .map(|item| item.parse::<i32>().unwrap())
+                        .collect::<Vec<i32>>();
+                    let bounding_box = BoundingBox::new_i32(bbox[0], bbox[1], bbox[2], bbox[3]);
+                    let lang = attributes
+                        .iter()
+                        .find(|item| item.name.local_name == "lang")
+                        .unwrap()
+                        .value
+                        .clone();
+                    current_par = Some(Paragraph::new(id, bounding_box, lang, Vec::new()));
+                }
+
+                let line = attributes.iter().find(|item| {
+                    item.name.local_name == "class" && item.value.contains("ocr_line")
+                });
+                if line.is_some() {
+                    let id = attributes
+                        .iter()
+                        .find(|item| item.name.local_name == "id")
+                        .unwrap()
+                        .value
+                        .clone();
+                    let title = attributes
+                        .iter()
+                        .find(|item| item.name.local_name == "title")
+                        .unwrap()
+                        .value
+                        .clone();
+                    let bbox = title
+                        .split(";")
+                        .take(1)
+                        .collect::<String>()
+                        .split_whitespace()
+                        .skip(1)
+                        .map(|item| item.parse::<i32>().unwrap())
+                        .collect::<Vec<i32>>();
+                    let bounding_box = BoundingBox::new_i32(bbox[0], bbox[1], bbox[2], bbox[3]);
+                    let words = Vec::new();
+                    current_line = Some(Line::new(id, bounding_box, words))
+                }
+                let word = attributes.iter().find(|item| {
+                    item.name.local_name == "class" && item.value.contains("ocrx_word")
+                });
+                if word.is_some() {
+                    let id = attributes
+                        .iter()
+                        .find(|item| item.name.local_name == "id")
+                        .unwrap()
+                        .value
+                        .clone();
+                    let title = attributes
+                        .iter()
+                        .find(|item| item.name.local_name == "title")
+                        .unwrap()
+                        .value
+                        .clone();
+                    let bbox = title
+                        .split(";")
+                        .take(1)
+                        .collect::<String>()
+                        .split_whitespace()
+                        .skip(1)
+                        .map(|item| item.parse::<i32>().unwrap())
+                        .collect::<Vec<i32>>();
+                    let bounding_box = BoundingBox::new_i32(bbox[0], bbox[1], bbox[2], bbox[3]);
+                    let content = String::new();
+                    current_word = Some(Word::new(id, bounding_box, content))
+                }
+            }
+            Ok(XmlEvent::Characters(content)) => {
+                if let Some(word) = current_word.as_mut() {
+                    word.content = content;
+                }
+            }
+            Ok(XmlEvent::EndElement { name }) => {
+                if name.local_name == "p" {
+                    if let Some(par) = current_par {
+                        result.push(par);
+                        current_par = None;
+                    }
+                }
+                if name.local_name == "span" {
+                    if let Some(word) = current_word {
+                        // closing word
+                        if let Some(line) = current_line.as_mut() {
+                            line.words.push(word);
+                        }
+                        current_word = None;
+                    } else if let Some(line) = current_line {
+                        // closing line
+                        if let Some(par) = current_par.as_mut() {
+                            par.lines.push(line);
+                        }
+                        current_line = None;
+                    }
+                }
+            }
+            Ok(_) => {}
+            Err(e) => {
+                panic!("Error: {}", e)
+            }
+        }
+    }
+    return result;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::parse_hocr_xml;
+    use std::fs::File;
+    use std::io::{BufReader, Read};
+
+    #[test]
+    fn test_parse_xml() -> anyhow::Result<()> {
+        let file = File::open("static/tesseract.out.hocr")?;
+        let mut file = BufReader::new(file);
+        let mut str = String::new();
+        file.read_to_string(&mut str);
+        let result = parse_hocr_xml(&str);
+
+        for p in result {
+            let text = p.text();
+            println!("Paragraph: {:}", text);
+        }
+        Ok(())
+    }
+}

src/ocr.rs → src/ocr/mod.rs

@@ -1,6 +1,13 @@
 use anyhow::Ok;
 use async_trait::async_trait;
 
+use self::{hocr_parse::parse_hocr_xml, model::Paragraph};
+
+mod hocr_parse;
+/// Structured recognized data, basically refers from hOCR spec.
+/// https://kba.github.io/hocr-spec/1.2
+mod model;
+
 #[derive(Debug, Clone)]
 pub struct RecognizeItem {
     pub text: String,
@@ -47,7 +54,9 @@ impl MarkupBox {
 
 #[async_trait]
 pub trait CharacterRecognizer {
+    #[deprecated]
     async fn recognize(&self, image: &image::DynamicImage) -> anyhow::Result<Vec<RecognizeItem>>;
+    async fn recognize_hocr(&self, image: &image::DynamicImage) -> anyhow::Result<Vec<Paragraph>>;
 }
 
 pub struct TesseractOCR {}
@@ -61,7 +70,7 @@ impl TesseractOCR {
 #[async_trait]
 impl CharacterRecognizer for TesseractOCR {
     async fn recognize(&self, image: &image::DynamicImage) -> anyhow::Result<Vec<RecognizeItem>> {
-        let default_args = rusty_tesseract::Args::default();
+        let mut default_args = rusty_tesseract::Args::default();
         let ri = rusty_tesseract::Image::from_dynamic_image(image)?;
         let output = rusty_tesseract::image_to_data(&ri, &default_args)?;
         let result: Vec<RecognizeItem> = output
@@ -75,4 +84,15 @@ impl CharacterRecognizer for TesseractOCR {
             .collect();
         Ok(result)
     }
+
+    async fn recognize_hocr(&self, image: &image::DynamicImage) -> anyhow::Result<Vec<Paragraph>> {
+        let mut default_args = rusty_tesseract::Args::default();
+        default_args
+            .config_variables
+            .insert("tessedit_create_hocr".into(), "1".into());
+        let ri = rusty_tesseract::Image::from_dynamic_image(image)?;
+        let output_hocr = rusty_tesseract::image_to_string(&ri, &default_args)?;
+        let result = parse_hocr_xml(&output_hocr);
+        Ok(result)
+    }
 }