oar_ocr_core/domain/tasks/
layout_detection.rs

1//! Concrete task implementations for layout detection.
2//!
3//! This module provides the layout detection task that identifies document layout elements.
4
5use super::validation::ensure_non_empty_images;
6use crate::ConfigValidator;
7use crate::core::OCRError;
8use crate::core::traits::TaskDefinition;
9use crate::core::traits::task::{ImageTaskInput, Task, TaskSchema, TaskType};
10use crate::processors::BoundingBox;
11use crate::utils::{ScoreValidator, validate_max_value};
12use serde::{Deserialize, Serialize};
13use std::collections::HashMap;
14
15/// Bounding box merge mode for layout elements.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
17pub enum MergeBboxMode {
18    /// Keep the larger bounding box
19    #[default]
20    Large,
21    /// Merge to union of bounding boxes
22    Union,
23    /// Keep the smaller bounding box
24    Small,
25}
26
27/// Unclip ratio configuration for layout detection.
28/// Controls how bounding boxes are scaled while keeping the center fixed.
29#[derive(Debug, Clone, Serialize, Deserialize)]
30pub enum UnclipRatio {
31    /// Single ratio applied to both width and height
32    Uniform(f32),
33    /// Separate ratios for (width, height)
34    Separate(f32, f32),
35    /// Per-class ratios: class_id -> (width_ratio, height_ratio)
36    PerClass(HashMap<usize, (f32, f32)>),
37}
38
39impl Default for UnclipRatio {
40    fn default() -> Self {
41        UnclipRatio::Separate(1.0, 1.0)
42    }
43}
44
45/// Configuration for layout detection task.
46#[derive(Debug, Clone, Serialize, Deserialize, ConfigValidator)]
47pub struct LayoutDetectionConfig {
48    /// Default score threshold for detection (default: 0.5, matches PaddleX's
49    /// `draw_threshold: 0.5` post-NMS visibility threshold from the published
50    /// `inference.yml` for the PP-DocLayout / PicoDet layout families).
51    #[validate(range(min = 0.0, max = 1.0))]
52    pub score_threshold: f32,
53    /// Maximum number of layout elements (default: 100)
54    #[validate(min = 1)]
55    pub max_elements: usize,
56    /// Per-class score thresholds (overrides score_threshold for specific classes)
57    /// PP-StructureV3 defaults:
58    /// - paragraph_title: 0.3
59    /// - formula: 0.3
60    /// - text: 0.4
61    /// - seal: 0.45
62    /// - others: 0.5
63    #[serde(default)]
64    pub class_thresholds: Option<HashMap<String, f32>>,
65    /// Per-class bounding box merge modes
66    #[serde(default)]
67    pub class_merge_modes: Option<HashMap<String, MergeBboxMode>>,
68    /// Enable NMS for layout detection (default: true)
69    #[serde(default = "default_layout_nms")]
70    pub layout_nms: bool,
71    /// NMS threshold (default: 0.5)
72    #[serde(default = "default_nms_threshold")]
73    pub nms_threshold: f32,
74    /// Unclip ratio for expanding/shrinking bounding boxes (PP-StructureV3)
75    /// Default: [1.0, 1.0] (no change)
76    #[serde(default)]
77    pub layout_unclip_ratio: Option<UnclipRatio>,
78}
79
80fn default_layout_nms() -> bool {
81    true
82}
83
84fn default_nms_threshold() -> f32 {
85    0.5
86}
87
88impl Default for LayoutDetectionConfig {
89    fn default() -> Self {
90        Self {
91            score_threshold: 0.5,
92            max_elements: 100,
93            class_thresholds: None,
94            class_merge_modes: None,
95            layout_nms: true,
96            nms_threshold: 0.5,
97            layout_unclip_ratio: None,
98        }
99    }
100}
101
102impl LayoutDetectionConfig {
103    /// Creates a config with PP-StructureV3 default class thresholds.
104    ///
105    /// PP-StructureV3 uses different thresholds for different element types:
106    /// - paragraph_title: 0.3
107    /// - formula: 0.3
108    /// - text: 0.4
109    /// - seal: 0.45
110    /// - others: 0.5 (default)
111    pub fn with_pp_structurev3_thresholds() -> Self {
112        let mut class_thresholds = HashMap::new();
113        class_thresholds.insert("paragraph_title".to_string(), 0.3);
114        class_thresholds.insert("formula".to_string(), 0.3);
115        class_thresholds.insert("text".to_string(), 0.4);
116        class_thresholds.insert("seal".to_string(), 0.45);
117
118        Self {
119            // Use the minimum of the per-class thresholds so the postprocessor doesn't drop
120            // candidates before the per-class filter runs.
121            score_threshold: 0.3,
122            max_elements: 100,
123            class_thresholds: Some(class_thresholds),
124            class_merge_modes: None,
125            layout_nms: true,
126            nms_threshold: 0.5,
127            layout_unclip_ratio: Some(UnclipRatio::Separate(1.0, 1.0)),
128        }
129    }
130
131    /// Creates a config with PP-DocLayoutV2 default thresholds and merge modes.
132    ///
133    /// These defaults follow the per-class thresholds and merge-mode settings
134    /// used by upstream PP-DocLayoutV2 deployments.
135    ///
136    /// Notes:
137    /// - The postprocessor applies `score_threshold` before per-class thresholds, so we set it
138    ///   to the minimum per-class threshold (0.4) to avoid dropping valid candidates early.
139    /// - `class_merge_modes` is populated for all labels so merge behavior is deterministic.
140    pub fn with_pp_doclayoutv2_defaults() -> Self {
141        let mut class_thresholds = HashMap::new();
142        class_thresholds.insert("abstract".to_string(), 0.5);
143        class_thresholds.insert("algorithm".to_string(), 0.5);
144        class_thresholds.insert("aside_text".to_string(), 0.5);
145        class_thresholds.insert("chart".to_string(), 0.5);
146        class_thresholds.insert("content".to_string(), 0.5);
147        class_thresholds.insert("display_formula".to_string(), 0.4);
148        class_thresholds.insert("doc_title".to_string(), 0.4);
149        class_thresholds.insert("figure_title".to_string(), 0.5);
150        class_thresholds.insert("footer".to_string(), 0.5);
151        class_thresholds.insert("footer_image".to_string(), 0.5);
152        class_thresholds.insert("footnote".to_string(), 0.5);
153        class_thresholds.insert("formula_number".to_string(), 0.5);
154        class_thresholds.insert("header".to_string(), 0.5);
155        class_thresholds.insert("header_image".to_string(), 0.5);
156        class_thresholds.insert("image".to_string(), 0.5);
157        class_thresholds.insert("inline_formula".to_string(), 0.4);
158        class_thresholds.insert("number".to_string(), 0.5);
159        class_thresholds.insert("paragraph_title".to_string(), 0.4);
160        class_thresholds.insert("reference".to_string(), 0.5);
161        class_thresholds.insert("reference_content".to_string(), 0.5);
162        class_thresholds.insert("seal".to_string(), 0.45);
163        class_thresholds.insert("table".to_string(), 0.5);
164        class_thresholds.insert("text".to_string(), 0.4);
165        class_thresholds.insert("vertical_text".to_string(), 0.4);
166        class_thresholds.insert("vision_footnote".to_string(), 0.5);
167
168        let mut merge_modes = HashMap::new();
169        merge_modes.insert("abstract".to_string(), MergeBboxMode::Union);
170        merge_modes.insert("algorithm".to_string(), MergeBboxMode::Union);
171        merge_modes.insert("aside_text".to_string(), MergeBboxMode::Union);
172        merge_modes.insert("chart".to_string(), MergeBboxMode::Large);
173        merge_modes.insert("content".to_string(), MergeBboxMode::Union);
174        merge_modes.insert("display_formula".to_string(), MergeBboxMode::Large);
175        merge_modes.insert("doc_title".to_string(), MergeBboxMode::Large);
176        merge_modes.insert("figure_title".to_string(), MergeBboxMode::Union);
177        merge_modes.insert("footer".to_string(), MergeBboxMode::Union);
178        merge_modes.insert("footer_image".to_string(), MergeBboxMode::Union);
179        merge_modes.insert("footnote".to_string(), MergeBboxMode::Union);
180        merge_modes.insert("formula_number".to_string(), MergeBboxMode::Union);
181        merge_modes.insert("header".to_string(), MergeBboxMode::Union);
182        merge_modes.insert("header_image".to_string(), MergeBboxMode::Union);
183        merge_modes.insert("image".to_string(), MergeBboxMode::Union);
184        merge_modes.insert("inline_formula".to_string(), MergeBboxMode::Large);
185        merge_modes.insert("number".to_string(), MergeBboxMode::Union);
186        merge_modes.insert("paragraph_title".to_string(), MergeBboxMode::Large);
187        merge_modes.insert("reference".to_string(), MergeBboxMode::Union);
188        merge_modes.insert("reference_content".to_string(), MergeBboxMode::Union);
189        merge_modes.insert("seal".to_string(), MergeBboxMode::Union);
190        merge_modes.insert("table".to_string(), MergeBboxMode::Union);
191        merge_modes.insert("text".to_string(), MergeBboxMode::Union);
192        merge_modes.insert("vertical_text".to_string(), MergeBboxMode::Union);
193        merge_modes.insert("vision_footnote".to_string(), MergeBboxMode::Union);
194
195        Self {
196            score_threshold: 0.4,
197            max_elements: 100,
198            class_thresholds: Some(class_thresholds),
199            class_merge_modes: Some(merge_modes),
200            layout_nms: true,
201            nms_threshold: 0.5,
202            layout_unclip_ratio: Some(UnclipRatio::Separate(1.0, 1.0)),
203        }
204    }
205
206    /// Creates a config with PP-DocLayoutV3 defaults (0.3 threshold + PP-DocLayout merge modes).
207    ///
208    /// PP-DocLayoutV3 keeps the same label set as PP-DocLayoutV2 but uses a lower
209    /// global threshold in the reference pipeline.
210    pub fn with_pp_doclayoutv3_defaults() -> Self {
211        let mut merge_modes = HashMap::new();
212        merge_modes.insert("abstract".to_string(), MergeBboxMode::Union);
213        merge_modes.insert("algorithm".to_string(), MergeBboxMode::Union);
214        merge_modes.insert("aside_text".to_string(), MergeBboxMode::Union);
215        merge_modes.insert("chart".to_string(), MergeBboxMode::Large);
216        merge_modes.insert("content".to_string(), MergeBboxMode::Union);
217        merge_modes.insert("display_formula".to_string(), MergeBboxMode::Large);
218        merge_modes.insert("doc_title".to_string(), MergeBboxMode::Large);
219        merge_modes.insert("figure_title".to_string(), MergeBboxMode::Union);
220        merge_modes.insert("footer".to_string(), MergeBboxMode::Union);
221        merge_modes.insert("footer_image".to_string(), MergeBboxMode::Union);
222        merge_modes.insert("footnote".to_string(), MergeBboxMode::Union);
223        merge_modes.insert("formula_number".to_string(), MergeBboxMode::Union);
224        merge_modes.insert("header".to_string(), MergeBboxMode::Union);
225        merge_modes.insert("header_image".to_string(), MergeBboxMode::Union);
226        merge_modes.insert("image".to_string(), MergeBboxMode::Union);
227        merge_modes.insert("inline_formula".to_string(), MergeBboxMode::Large);
228        merge_modes.insert("number".to_string(), MergeBboxMode::Union);
229        merge_modes.insert("paragraph_title".to_string(), MergeBboxMode::Large);
230        merge_modes.insert("reference".to_string(), MergeBboxMode::Union);
231        merge_modes.insert("reference_content".to_string(), MergeBboxMode::Union);
232        merge_modes.insert("seal".to_string(), MergeBboxMode::Union);
233        merge_modes.insert("table".to_string(), MergeBboxMode::Union);
234        merge_modes.insert("text".to_string(), MergeBboxMode::Union);
235        merge_modes.insert("vertical_text".to_string(), MergeBboxMode::Union);
236        merge_modes.insert("vision_footnote".to_string(), MergeBboxMode::Union);
237
238        Self {
239            score_threshold: 0.3,
240            max_elements: 100,
241            class_thresholds: None,
242            class_merge_modes: Some(merge_modes),
243            layout_nms: true,
244            nms_threshold: 0.5,
245            layout_unclip_ratio: Some(UnclipRatio::Separate(1.0, 1.0)),
246        }
247    }
248
249    /// Creates a config with PP-StructureV3 default thresholds, merge modes, and unclip ratio.
250    ///
251    /// Merge modes follow standard configuration:
252    /// - "large": paragraph_title, image, formula, chart
253    /// - "union": all other PP-DocLayout_plus-L classes
254    pub fn with_pp_structurev3_defaults() -> Self {
255        let mut cfg = Self::with_pp_structurev3_thresholds();
256
257        let mut merge_modes = HashMap::new();
258        merge_modes.insert("paragraph_title".to_string(), MergeBboxMode::Large);
259        merge_modes.insert("image".to_string(), MergeBboxMode::Large);
260        merge_modes.insert("text".to_string(), MergeBboxMode::Union);
261        merge_modes.insert("number".to_string(), MergeBboxMode::Union);
262        merge_modes.insert("abstract".to_string(), MergeBboxMode::Union);
263        merge_modes.insert("content".to_string(), MergeBboxMode::Union);
264        merge_modes.insert("figure_table_chart_title".to_string(), MergeBboxMode::Union);
265        merge_modes.insert("formula".to_string(), MergeBboxMode::Large);
266        merge_modes.insert("table".to_string(), MergeBboxMode::Union);
267        merge_modes.insert("reference".to_string(), MergeBboxMode::Union);
268        merge_modes.insert("doc_title".to_string(), MergeBboxMode::Union);
269        merge_modes.insert("footnote".to_string(), MergeBboxMode::Union);
270        merge_modes.insert("header".to_string(), MergeBboxMode::Union);
271        merge_modes.insert("algorithm".to_string(), MergeBboxMode::Union);
272        merge_modes.insert("footer".to_string(), MergeBboxMode::Union);
273        merge_modes.insert("seal".to_string(), MergeBboxMode::Union);
274        merge_modes.insert("chart".to_string(), MergeBboxMode::Large);
275        merge_modes.insert("formula_number".to_string(), MergeBboxMode::Union);
276        merge_modes.insert("aside_text".to_string(), MergeBboxMode::Union);
277        merge_modes.insert("reference_content".to_string(), MergeBboxMode::Union);
278
279        cfg.class_merge_modes = Some(merge_modes);
280        cfg.layout_unclip_ratio = Some(UnclipRatio::Separate(1.0, 1.0));
281        cfg
282    }
283
284    /// Gets the threshold for a specific class.
285    ///
286    /// Returns the class-specific threshold if configured, otherwise the default threshold.
287    pub fn get_class_threshold(&self, class_name: &str) -> f32 {
288        self.class_thresholds
289            .as_ref()
290            .and_then(|thresholds| thresholds.get(class_name).copied())
291            .unwrap_or(self.score_threshold)
292    }
293
294    /// Gets the merge mode for a specific class.
295    ///
296    /// Returns the class-specific merge mode if configured, otherwise Large (default).
297    pub fn get_class_merge_mode(&self, class_name: &str) -> MergeBboxMode {
298        self.class_merge_modes
299            .as_ref()
300            .and_then(|modes| modes.get(class_name).copied())
301            .unwrap_or(MergeBboxMode::Large)
302    }
303}
304
305/// A detected layout element from the layout detection model.
306///
307/// This represents the raw output from layout detection before conversion
308/// to the final `LayoutElement` in `domain::structure`.
309#[derive(Debug, Clone)]
310pub struct LayoutDetectionElement {
311    /// Bounding box of the element
312    pub bbox: BoundingBox,
313    /// Type of layout element (raw string label from model)
314    pub element_type: String,
315    /// Confidence score (0.0 to 1.0)
316    pub score: f32,
317}
318
319/// Output from layout detection task.
320#[derive(Debug, Clone)]
321pub struct LayoutDetectionOutput {
322    /// Detected layout elements per image
323    pub elements: Vec<Vec<LayoutDetectionElement>>,
324    /// Whether elements are already sorted by reading order (e.g., from PP-DocLayoutV2)
325    ///
326    /// When `true`, downstream consumers can skip reading order sorting algorithms
327    /// as the elements are already in the correct reading order based on model output.
328    pub is_reading_order_sorted: bool,
329}
330
331impl LayoutDetectionOutput {
332    /// Creates an empty layout detection output.
333    pub fn empty() -> Self {
334        Self {
335            elements: Vec::new(),
336            is_reading_order_sorted: false,
337        }
338    }
339
340    /// Creates a layout detection output with the given capacity.
341    pub fn with_capacity(capacity: usize) -> Self {
342        Self {
343            elements: Vec::with_capacity(capacity),
344            is_reading_order_sorted: false,
345        }
346    }
347
348    /// Sets the reading order sorted flag.
349    pub fn with_reading_order_sorted(mut self, sorted: bool) -> Self {
350        self.is_reading_order_sorted = sorted;
351        self
352    }
353}
354
355impl TaskDefinition for LayoutDetectionOutput {
356    const TASK_NAME: &'static str = "layout_detection";
357    const TASK_DOC: &'static str = "Layout detection/analysis";
358
359    fn empty() -> Self {
360        LayoutDetectionOutput::empty()
361    }
362}
363
364/// Layout detection task implementation.
365#[derive(Debug, Default)]
366pub struct LayoutDetectionTask {
367    config: LayoutDetectionConfig,
368}
369
370impl LayoutDetectionTask {
371    /// Creates a new layout detection task.
372    pub fn new(config: LayoutDetectionConfig) -> Self {
373        Self { config }
374    }
375}
376
377impl Task for LayoutDetectionTask {
378    type Config = LayoutDetectionConfig;
379    type Input = ImageTaskInput;
380    type Output = LayoutDetectionOutput;
381
382    fn task_type(&self) -> TaskType {
383        TaskType::LayoutDetection
384    }
385
386    fn schema(&self) -> TaskSchema {
387        TaskSchema::new(
388            TaskType::LayoutDetection,
389            vec!["image".to_string()],
390            vec!["layout_elements".to_string()],
391        )
392    }
393
394    fn validate_input(&self, input: &Self::Input) -> Result<(), OCRError> {
395        ensure_non_empty_images(&input.images, "No images provided for layout detection")?;
396
397        Ok(())
398    }
399
400    fn validate_output(&self, output: &Self::Output) -> Result<(), OCRError> {
401        let validator = ScoreValidator::new_unit_range("score");
402
403        for (idx, elements) in output.elements.iter().enumerate() {
404            // Validate element count
405            validate_max_value(
406                elements.len(),
407                self.config.max_elements,
408                "element count",
409                &format!("Image {}", idx),
410            )?;
411
412            // Validate scores
413            let scores: Vec<f32> = elements.iter().map(|e| e.score).collect();
414            validator.validate_scores_with(&scores, |elem_idx| {
415                format!("Image {}, element {}", idx, elem_idx)
416            })?;
417        }
418
419        Ok(())
420    }
421
422    fn empty_output(&self) -> Self::Output {
423        LayoutDetectionOutput::empty()
424    }
425}
426
427#[cfg(test)]
428mod tests {
429    use super::*;
430    use crate::processors::Point;
431    use image::RgbImage;
432
433    #[test]
434    fn test_layout_detection_task_creation() {
435        let task = LayoutDetectionTask::default();
436        assert_eq!(task.task_type(), TaskType::LayoutDetection);
437    }
438
439    #[test]
440    fn test_input_validation() {
441        let task = LayoutDetectionTask::default();
442
443        // Empty images should fail
444        let empty_input = ImageTaskInput::new(vec![]);
445        assert!(task.validate_input(&empty_input).is_err());
446
447        // Valid images should pass
448        let valid_input = ImageTaskInput::new(vec![RgbImage::new(100, 100)]);
449        assert!(task.validate_input(&valid_input).is_ok());
450    }
451
452    #[test]
453    fn test_output_validation() {
454        let task = LayoutDetectionTask::default();
455
456        // Valid output should pass
457        let box1 = BoundingBox::new(vec![
458            Point::new(0.0, 0.0),
459            Point::new(10.0, 0.0),
460            Point::new(10.0, 10.0),
461            Point::new(0.0, 10.0),
462        ]);
463        let element = LayoutDetectionElement {
464            bbox: box1,
465            element_type: "text".to_string(),
466            score: 0.95,
467        };
468        let output = LayoutDetectionOutput {
469            elements: vec![vec![element]],
470            is_reading_order_sorted: false,
471        };
472        assert!(task.validate_output(&output).is_ok());
473    }
474
475    #[test]
476    fn test_schema() {
477        let task = LayoutDetectionTask::default();
478        let schema = task.schema();
479        assert_eq!(schema.task_type, TaskType::LayoutDetection);
480        assert!(schema.input_types.contains(&"image".to_string()));
481        assert!(schema.output_types.contains(&"layout_elements".to_string()));
482    }
483}
oar_ocr_core/domain/tasks/layout_detection.rs

oar_ocr_core/domain/tasks/
layout_detection.rs
