ak/ocr postprocess

chunkmydocs/src/models/server/segment.rs

@@ -1,3 +1,4 @@
+use crate::utils::configs::extraction_config;
 use postgres_types::{ FromSql, ToSql };
 use serde::{ Deserialize, Serialize };
 use strum_macros::{ Display, EnumString };
@@ -130,7 +131,7 @@ impl Segment {
                         .trim();
                     format!("<ul><li>{}</li></ul>", cleaned_content)
                 }
-            },
+            }
             SegmentType::Picture => "<img src='' alt='{}' />".to_string(),
             _ =>
                 format!(
@@ -163,7 +164,24 @@ impl Segment {
         }
     }
 
+    fn update_content(&mut self) {
+        if let Some(ocr) = &self.ocr {
+            let config = match extraction_config::Config::from_env() {
+                Ok(config) => config,
+                Err(e) => {
+                    eprintln!("Error getting extraction config: {}", e);
+                    return;
+                }
+            };
+            let avg_confidence = ocr.iter().map(|ocr_result| ocr_result.confidence.unwrap_or(0.0)).sum::<f32>() / ocr.len() as f32;
+            if avg_confidence >= config.ocr_confidence_threshold {
+                self.content = ocr.iter().map(|ocr_result| ocr_result.text.clone()).collect::<Vec<String>>().join(" ");
+            }
+        }
+    }
+
     pub fn finalize(&mut self) {
+        self.update_content();
         self.html = Some(self.to_html());
         self.markdown = Some(self.to_markdown());
     }

chunkmydocs/src/utils/configs/extraction_config.rs

@@ -22,6 +22,8 @@ pub struct Config {
     pub base_url: String,
     #[serde(default = "default_ocr_concurrency")]
     pub ocr_concurrency: usize,
+    #[serde(default = "default_ocr_confidence_threshold")]
+    pub ocr_confidence_threshold: f32,
     #[serde(default = "default_pdf_density")]
     pub pdf_density: f32,
     #[serde(default = "default_page_image_density")]
@@ -34,6 +36,10 @@ fn default_ocr_concurrency() -> usize {
     10
 }
 
+fn default_ocr_confidence_threshold() -> f32 {
+    0.85
+}
+
 fn default_pdf_density() -> f32 {
     72.0
 }

chunkmydocs/src/utils/services/ocr.rs

@@ -13,7 +13,7 @@ pub async fn download_and_ocr(
     s3_client: &S3Client,
     reqwest_client: &ReqwestClient,
     image_location: &str,
-) -> Result<(Vec<OCRResult>, String), Box<dyn std::error::Error>> {
+) -> Result<(Vec<OCRResult>, String, String), Box<dyn std::error::Error>> {
     let original_file =
         download_to_tempfile(s3_client, reqwest_client, image_location, None).await?;
     // let preprocessed_file = preprocess_image(&original_file.path()).await?;
@@ -24,14 +24,14 @@ pub async fn download_and_ocr(
             return Err(e.to_string().into());
         }
     };
-    Ok((ocr_results, "".to_string()))
+    Ok((ocr_results, "".to_string(), "".to_string()))
 }
 
 pub async fn download_and_table_ocr(
     s3_client: &S3Client,
     reqwest_client: &ReqwestClient,
     image_location: &str,
-) -> Result<(Vec<OCRResult>, String), Box<dyn std::error::Error>> {
+) -> Result<(Vec<OCRResult>, String, String), Box<dyn std::error::Error>> {
     let original_file =
         download_to_tempfile(s3_client, reqwest_client, image_location, None).await?;
     let original_file_path = original_file.path().to_owned();
@@ -59,8 +59,9 @@ pub async fn download_and_table_ocr(
     let table_structures_with_content =
         add_content_to_table_structure(table_structures, ocr_results);
     let html = get_table_html(table_structures_with_content.clone());
+    let markdown = get_table_markdown(table_structures_with_content.clone());
     let table_ocr_results = table_structures_to_ocr_results(table_structures_with_content);
-    Ok((table_ocr_results, html))
+    Ok((table_ocr_results, html, markdown))
 }
 
 fn add_content_to_table_structure(
@@ -102,15 +103,32 @@ fn get_table_html(table_structures: Vec<TableStructure>) -> String {
     for row in table_structures {
         html.push_str("<tr>");
         for cell in row.cells {
-            html.push_str(&format!("<td>{}</td>", cell.content.unwrap_or_default()));
+            html.push_str(&format!("<td colspan='{}' rowspan='{}'>{}</td>", cell.col_span, cell.row_span, cell.content.unwrap_or_default()));
         }
         html.push_str("</tr>");
     }
     html.push_str("</table>");
     return html;
 }
 
-fn table_structures_to_ocr_results(table_structures: Vec<TableStructure>) -> Vec<OCRResult> {
+fn get_table_markdown(
+    table_structures: Vec<TableStructure>
+) -> String {
+    let mut markdown = String::new();
+    markdown.push_str("|");
+    for row in table_structures {
+        markdown.push_str("|");
+        for cell in row.cells {
+            markdown.push_str(&format!("{} |", cell.content.unwrap_or_default()));
+        }
+    }
+    markdown.push_str("\n");
+    return markdown;
+}
+
+fn table_structures_to_ocr_results(
+    table_structures: Vec<TableStructure>
+) -> Vec<OCRResult> {
     table_structures
         .iter()
         .flat_map(|table| table.cells.iter())

chunkmydocs/src/utils/workers/ocr.rs

@@ -91,11 +91,15 @@ pub async fn process(payload: QueuePayload) -> Result<(), Box<dyn std::error::Er
                             };
                             drop(_permit);
                             match ocr_result {
-                                Ok((ocr_result, html)) => {
+                                Ok((ocr_result, html, markdown)) => {
                                     segment.ocr = Some(ocr_result);
+                                    segment.finalize();
                                     if !html.is_empty() {
                                         segment.html = Some(html);
                                     }
+                                    if !markdown.is_empty() {
+                                        segment.markdown = Some(markdown);
+                                    }
                                     Ok::<_, Box<dyn std::error::Error>>(())
                                 }
                                 Err(e) => {