diff --git a/crates/liteparse/src/ocr_merge.rs b/crates/liteparse/src/ocr_merge.rs index acd4e248..9cf983e1 100644 --- a/crates/liteparse/src/ocr_merge.rs +++ b/crates/liteparse/src/ocr_merge.rs @@ -4,7 +4,13 @@ use crate::error::LiteParseError; use crate::ocr::{OcrEngine, OcrOptions, OcrResult}; use crate::types::{Page, TextItem}; use image::{ImageBuffer, Rgba}; -use pdfium::Document; +use pdfium::{Bitmap, Document}; + +/// Minimum native text length before skipping full-page OCR. +const SPARSE_PAGE_TEXT_THRESHOLD: usize = 100; +/// Match `Page::image_bounds` filters in the OCR trigger path. +const MIN_IMAGE_SIZE_PT: f32 = 25.0; +const MAX_IMAGE_PAGE_COVERAGE: f32 = 0.9; /// Owned page bitmap prepared for OCR. Indices refer to positions in the `pages` slice. pub(crate) struct RenderedPage { @@ -12,10 +18,27 @@ pub(crate) struct RenderedPage { pub rgb_bytes: Vec, pub width: u32, pub height: u32, + pub coord_map: OcrCoordMap, +} + +#[derive(Clone, Copy, Debug)] +pub(crate) enum OcrCoordMap { + FullPage { + dpi: f32, + }, + ImageRegion { + page_x: f32, + page_y: f32, + page_w: f32, + page_h: f32, + }, } /// Render pages that need OCR from an already-open document. /// +/// OCR is triggered when a page has fewer than [`SPARSE_PAGE_TEXT_THRESHOLD`] characters +/// of native text (full-page OCR) or has qualifying embedded images (per-image OCR). +/// /// The pdfium `Document` holds raw pointers that are not `Send`, so callers must /// drop it before awaiting the OCR engine. pub(crate) fn render_pages_for_ocr( @@ -27,45 +50,66 @@ pub(crate) fn render_pages_for_ocr( for (idx, page) in pages.iter().enumerate() { let text_length: usize = page.text_items.iter().map(|item| item.text.len()).sum(); let page_obj = document.page((page.page_number - 1) as i32)?; - let has_images = !page_obj.image_bounds(25.0, 0.9).is_empty(); + let images = page_obj.image_bounds(MIN_IMAGE_SIZE_PT, MAX_IMAGE_PAGE_COVERAGE); - if text_length >= 100 && !has_images { + if text_length >= SPARSE_PAGE_TEXT_THRESHOLD && images.is_empty() { + continue; + } + + if text_length < SPARSE_PAGE_TEXT_THRESHOLD { + let bitmap = page_obj.render(dpi)?; + let (rgb_bytes, width, height) = bitmap_to_rgb(&bitmap)?; + rendered.push(RenderedPage { + idx, + rgb_bytes, + width, + height, + coord_map: OcrCoordMap::FullPage { dpi }, + }); continue; } - let bitmap = page_obj.render(dpi)?; - let width = bitmap.width() as u32; - let height = bitmap.height() as u32; - let rgba = bitmap.to_rgba(); - - let img: ImageBuffer, Vec> = ImageBuffer::from_raw(width, height, rgba) - .ok_or(LiteParseError::Other( - "failed to create image buffer".into(), - ))?; - let rgb_img = image::DynamicImage::ImageRgba8(img).to_rgb8(); - let rgb_bytes = rgb_img.into_raw(); - - rendered.push(RenderedPage { - idx, - rgb_bytes, - width, - height, - }); + for (image_idx, bounds) in images.iter().enumerate() { + let bitmap = page_obj.render_image_object(image_idx)?; + let (rgb_bytes, width, height) = bitmap_to_rgb(&bitmap)?; + rendered.push(RenderedPage { + idx, + rgb_bytes, + width, + height, + coord_map: OcrCoordMap::ImageRegion { + page_x: bounds.x, + page_y: bounds.y, + page_w: bounds.width, + page_h: bounds.height, + }, + }); + } } Ok(rendered) } +fn bitmap_to_rgb(bitmap: &Bitmap) -> Result<(Vec, u32, u32), LiteParseError> { + let width = bitmap.width() as u32; + let height = bitmap.height() as u32; + let rgba = bitmap.to_rgba(); + + let img: ImageBuffer, Vec> = ImageBuffer::from_raw(width, height, rgba).ok_or( + LiteParseError::Other("failed to create image buffer".into()), + )?; + let rgb_img = image::DynamicImage::ImageRgba8(img).to_rgb8(); + Ok((rgb_img.into_raw(), width, height)) +} + /// Run OCR on pre-rendered page bitmaps and merge results into `pages`. pub(crate) async fn ocr_and_merge_rendered( pages: &mut [Page], rendered: Vec, - dpi: f32, + _dpi: f32, ocr_engine: Arc, ocr_language: &str, num_workers: usize, ) -> Result<(), LiteParseError> { - // Phase 1: spawn OCR tasks onto the tokio runtime so they run on - // separate threads. A semaphore limits concurrency to `num_workers`. let num_workers = num_workers.max(1); let semaphore = Arc::new(tokio::sync::Semaphore::new(num_workers)); let mut handles = Vec::with_capacity(rendered.len()); @@ -78,12 +122,17 @@ pub(crate) async fn ocr_and_merge_rendered( let language = ocr_language.to_string(); let page_number = pages[r.idx].page_number; let rt_handle = handle.clone(); + let coord_map = r.coord_map; + let bw = r.width; + let bh = r.height; handles.push(( r.idx, page_number, + coord_map, + bw, + bh, tokio::task::spawn_blocking(move || { - // Block this thread until a permit is available. let _permit = rt_handle .block_on(sem.acquire_owned()) .expect("semaphore closed"); @@ -93,10 +142,7 @@ pub(crate) async fn ocr_and_merge_rendered( )); } - // Phase 3: collect results and merge into pages. - let scale_factor = 72.0 / dpi; - - for (idx, page_number, handle) in handles { + for (idx, page_number, coord_map, bitmap_width, bitmap_height, handle) in handles { let ocr_results: Vec = match handle.await { Ok(Ok(results)) => results, Ok(Err(e)) => { @@ -114,17 +160,19 @@ pub(crate) async fn ocr_and_merge_rendered( } let page = &mut pages[idx]; + let skip_overlap_check = matches!(coord_map, OcrCoordMap::ImageRegion { .. }); + for r in &ocr_results { if r.confidence <= 0.1 { continue; } - let ocr_x = r.bbox[0] * scale_factor; - let ocr_y = r.bbox[1] * scale_factor; - let ocr_w = (r.bbox[2] - r.bbox[0]) * scale_factor; - let ocr_h = (r.bbox[3] - r.bbox[1]) * scale_factor; + let (ocr_x, ocr_y, ocr_w, ocr_h) = + map_ocr_bbox_to_page(r.bbox, bitmap_width, bitmap_height, &coord_map); - if overlaps_existing_text(&page.text_items, ocr_x, ocr_y, ocr_w, ocr_h, 2.0) { + if !skip_overlap_check + && overlaps_existing_text(&page.text_items, ocr_x, ocr_y, ocr_w, ocr_h, 2.0) + { continue; } @@ -150,6 +198,43 @@ pub(crate) async fn ocr_and_merge_rendered( Ok(()) } +fn map_ocr_bbox_to_page( + bbox: [f32; 4], + bitmap_width: u32, + bitmap_height: u32, + coord_map: &OcrCoordMap, +) -> (f32, f32, f32, f32) { + match coord_map { + OcrCoordMap::FullPage { dpi } => { + let scale = 72.0 / dpi; + ( + bbox[0] * scale, + bbox[1] * scale, + (bbox[2] - bbox[0]) * scale, + (bbox[3] - bbox[1]) * scale, + ) + } + OcrCoordMap::ImageRegion { + page_x, + page_y, + page_w, + page_h, + } => { + if bitmap_width == 0 || bitmap_height == 0 { + return (0.0, 0.0, 0.0, 0.0); + } + let sx = page_w / bitmap_width as f32; + let sy = page_h / bitmap_height as f32; + ( + page_x + bbox[0] * sx, + page_y + bbox[1] * sy, + (bbox[2] - bbox[0]) * sx, + (bbox[3] - bbox[1]) * sy, + ) + } + } +} + /// Check if an OCR bounding box overlaps with any existing text item. fn overlaps_existing_text( items: &[TextItem], @@ -265,4 +350,30 @@ mod tests { assert_eq!(clean_ocr_table_artifacts(" "), ""); assert_eq!(clean_ocr_table_artifacts(" 123 "), "123"); } + + #[test] + fn test_map_ocr_bbox_full_page() { + let map = OcrCoordMap::FullPage { dpi: 150.0 }; + let (x, y, w, h) = map_ocr_bbox_to_page([30.0, 60.0, 90.0, 120.0], 1000, 1000, &map); + let scale = 72.0 / 150.0; + assert!((x - 30.0 * scale).abs() < 1e-5); + assert!((y - 60.0 * scale).abs() < 1e-5); + assert!((w - 60.0 * scale).abs() < 1e-5); + assert!((h - 60.0 * scale).abs() < 1e-5); + } + + #[test] + fn test_map_ocr_bbox_image_region() { + let map = OcrCoordMap::ImageRegion { + page_x: 100.0, + page_y: 50.0, + page_w: 200.0, + page_h: 100.0, + }; + let (x, y, w, h) = map_ocr_bbox_to_page([10.0, 20.0, 110.0, 70.0], 200, 100, &map); + assert!((x - 110.0).abs() < 1e-5); + assert!((y - 70.0).abs() < 1e-5); + assert!((w - 100.0).abs() < 1e-5); + assert!((h - 50.0).abs() < 1e-5); + } }