easypdf_markdown/ocr/processor/
core.rs

1use crate::render::{ImageFormat, RenderBackend, RenderConfig};
2use crate::{MarkdownProcessorCapabilities, MarkdownWarning, PdfMarkdownProcessor};
3use easypdf_core::PdfInput;
4use easypdf_core::{PdfBlock, PdfBlockType, PdfDocumentModel, PdfPageModel, SourceLocation};
5use easypdf_core::{PdfError, Result};
6
7use crate::ocr::config::{OcrConfig, OcrTrigger};
8use crate::ocr::engine::{OcrEngine, OcrImage};
9
10use super::renderer::{StoredRendererAdapter, render_error_to_pdf};
11
12/// Markdown 处理器管道中的 OCR 处理器。
13///
14/// 扫描文档模型中需要 OCR 的页面（基于配置的 [`OcrTrigger`] 策略），
15/// 将这些页面渲染为图像，通过配置的 [`OcrEngine`] 执行 OCR，
16/// 并将识别文本作为新的 [`PdfBlock::Paragraph`] 块注入。
17///
18/// # Examples
19///
20/// ```
21/// use easypdf_markdown::PdfMarkdownProcessor;
22/// use easypdf_markdown::ocr::OcrProcessor;
23///
24/// let processor = OcrProcessor::with_mock_engine();
25/// let caps = processor.capabilities();
26/// assert!(caps.ocr());
27/// ```
28pub struct OcrProcessor {
29    engine: Box<dyn OcrEngine>,
30    renderer: Option<Box<dyn crate::render::PdfRenderer>>,
31    backend: RenderBackend,
32    config: OcrConfig,
33}
34
35impl std::fmt::Debug for OcrProcessor {
36    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37        f.debug_struct("OcrProcessor")
38            .field("engine", &self.engine.name())
39            .field("config", &self.config)
40            .finish_non_exhaustive()
41    }
42}
43
44impl OcrProcessor {
45    /// 使用给定引擎和渲染后端创建新的 OCR 处理器。
46    ///
47    /// 渲染后端用于在处理时从 PDF 输入路径构建
48    /// [`PdfRenderer`](crate::render::PdfRenderer)。
49    #[must_use]
50    pub fn new(engine: Box<dyn OcrEngine>, backend: RenderBackend) -> Self {
51        Self {
52            engine,
53            renderer: None,
54            backend,
55            config: OcrConfig::default(),
56        }
57    }
58
59    /// 使用模拟引擎创建新的 OCR 处理器（用于测试）。
60    #[must_use]
61    pub fn with_mock_engine() -> Self {
62        use crate::ocr::engines::MockOcrEngine;
63        Self::new(Box::new(MockOcrEngine::new()), RenderBackend::Text)
64    }
65
66    /// 设置预构建的渲染器（覆盖后端渲染方式）。
67    ///
68    /// 适用于使用不需要真实 PDF 文件的模拟渲染器进行测试。
69    #[must_use]
70    pub fn with_renderer(mut self, renderer: Box<dyn crate::render::PdfRenderer>) -> Self {
71        self.renderer = Some(renderer);
72        self
73    }
74
75    /// 设置 OCR 配置。
76    #[must_use]
77    pub fn with_config(mut self, config: OcrConfig) -> Self {
78        self.config = config;
79        self
80    }
81
82    /// 获取用于页面渲染的渲染器。
83    ///
84    /// 如果存在预构建的渲染器则使用它，否则从输入路径构建。
85    /// 对于字节输入且无预构建渲染器的情况，写入临时文件。
86    fn get_renderer<'a>(
87        &'a self,
88        input: &PdfInput,
89    ) -> Result<Box<dyn crate::render::PdfRenderer + 'a>> {
90        if let Some(ref renderer) = self.renderer {
91            // 已注入预构建渲染器（例如用于测试）。
92            // 由于无法克隆 trait 对象，将其包装在委托适配器中。
93            return Ok(Box::new(StoredRendererAdapter(renderer.as_ref())));
94        }
95
96        match input {
97            PdfInput::Path(path) => self
98                .backend
99                .build_renderer(path)
100                .map_err(|e| render_error_to_pdf(&e)),
101            PdfInput::Bytes(bytes) => {
102                let tmp = tempfile::NamedTempFile::new()?;
103                std::fs::write(tmp.path(), bytes)?;
104                self.backend
105                    .build_renderer(tmp.path())
106                    .map_err(|e| render_error_to_pdf(&e))
107            }
108            _ => Err(PdfError::Other("unsupported PDF input type".to_owned())),
109        }
110    }
111
112    /// 根据触发策略判断页面是否需要 OCR。
113    fn page_needs_ocr(&self, page: &PdfPageModel) -> bool {
114        match self.config.trigger {
115            OcrTrigger::Always => true,
116            OcrTrigger::OnEmptyPage => {
117                // 如果页面没有段落或标题块则执行 OCR。
118                !page.blocks().iter().any(|b| {
119                    matches!(
120                        b.block_type(),
121                        PdfBlockType::Paragraph | PdfBlockType::Heading
122                    )
123                })
124            }
125            OcrTrigger::WhenTextSparse { threshold } => {
126                let total = page.blocks().len();
127                if total == 0 {
128                    return true;
129                }
130                let text_count = page
131                    .blocks()
132                    .iter()
133                    .filter(|b| {
134                        matches!(
135                            b.block_type(),
136                            PdfBlockType::Paragraph
137                                | PdfBlockType::Heading
138                                | PdfBlockType::Code
139                                | PdfBlockType::Footnote
140                                | PdfBlockType::BlockQuote
141                        )
142                    })
143                    .count();
144                // 阈值比较：text_count / total < threshold。
145                // 改写为整数运算：text_count * 1000 < total * (threshold * 1000)。
146                // 页面计数始终很小，乘法不会溢出。
147                #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
148                let threshold_permille = (f64::from(threshold) * 1000.0).round() as usize;
149                text_count.saturating_mul(1000) < total.saturating_mul(threshold_permille)
150            }
151        }
152    }
153
154    /// 渲染页面并对其执行 OCR。
155    fn ocr_page(
156        &self,
157        renderer: &dyn crate::render::PdfRenderer,
158        page_index: usize,
159    ) -> Result<(String, Option<f32>)> {
160        let render_config = RenderConfig {
161            dpi: self.config.render_dpi,
162            format: ImageFormat::Png,
163            ..RenderConfig::default()
164        };
165
166        let rendered_img = renderer
167            .render_page(page_index, &render_config)
168            .map_err(|e| render_error_to_pdf(&e))?;
169        let ocr_image = OcrImage::from_rendered(&rendered_img);
170
171        let ocr_result = self
172            .engine
173            .recognize(&ocr_image)
174            .map_err(|e| PdfError::Other(format!("OCR engine error: {e}")))?;
175
176        let text = ocr_result.text;
177        let confidence = ocr_result.confidence;
178
179        // 按最小文本长度过滤。
180        if text.trim().len() < self.config.min_text_length {
181            return Ok((String::new(), confidence));
182        }
183
184        Ok((text, confidence))
185    }
186}
187
188impl PdfMarkdownProcessor for OcrProcessor {
189    fn capabilities(&self) -> MarkdownProcessorCapabilities {
190        MarkdownProcessorCapabilities::new().with_ocr()
191    }
192
193    fn process(
194        &self,
195        input: &PdfInput,
196        document: PdfDocumentModel,
197    ) -> Result<(PdfDocumentModel, Vec<MarkdownWarning>)> {
198        let renderer = self.get_renderer(input)?;
199        let mut new_pages = Vec::with_capacity(document.page_count());
200        let mut warnings = Vec::new();
201
202        for page in document.pages() {
203            if !self.page_needs_ocr(page) {
204                new_pages.push(page.clone());
205                continue;
206            }
207
208            match self.ocr_page(renderer.as_ref(), page.index().value()) {
209                Ok((text, confidence)) => {
210                    if text.is_empty() {
211                        // OCR 未返回可用文本。
212                        warnings.push(MarkdownWarning::OcrUnavailable {
213                            page_index: page.index(),
214                        });
215                        new_pages.push(page.clone());
216                        continue;
217                    }
218
219                    // 检查置信度阈值。
220                    if let Some(conf) = confidence
221                        && conf < self.config.min_confidence
222                    {
223                        warnings.push(MarkdownWarning::ProcessorFailed {
224                            message: format!(
225                                "OCR confidence {conf:.2} below threshold {:.2} on page {}",
226                                self.config.min_confidence,
227                                page.index().value() + 1
228                            ),
229                        });
230                    }
231
232                    // 构建注入 OCR 文本后的新页面。
233                    let mut new_page = PdfPageModel::new(page.index());
234                    if let (Some(w), Some(h)) = (page.width_pt(), page.height_pt()) {
235                        new_page = new_page.with_dimensions(w, h);
236                    }
237                    new_page = new_page.with_rotation(page.rotation());
238
239                    // 保留现有块。
240                    for block in page.blocks() {
241                        new_page = new_page.with_block(block.clone());
242                    }
243
244                    // 将 OCR 文本作为新段落注入。
245                    let ocr_source = SourceLocation::new(page.index(), confidence.unwrap_or(0.5));
246                    new_page = new_page.with_block(PdfBlock::paragraph(text, ocr_source));
247
248                    new_pages.push(new_page);
249                }
250                Err(e) => {
251                    warnings.push(MarkdownWarning::ProcessorFailed {
252                        message: format!("OCR failed on page {}: {e}", page.index().value() + 1),
253                    });
254                    new_pages.push(page.clone());
255                }
256            }
257        }
258
259        Ok((
260            PdfDocumentModel::new(document.metadata().clone(), new_pages),
261            warnings,
262        ))
263    }
264}
easypdf_markdown/ocr/processor/core.rs

easypdf_markdown/ocr/processor/
core.rs
