layoutparser_ort/
layout_element.rs

1use geo_types::{coord, Coord, Rect};
2
3#[derive(Debug, Clone)]
4#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
5/// A detected element in a document's layout.
6pub struct LayoutElement {
7    /// Bounding box of the element.
8    pub bbox: Rect<f32>,
9    /// Type of element. This value depends on the labels of the model used to detect this element.
10    pub element_type: String,
11    // Confidence for the detection.
12    pub confidence: f32,
13    /// Source of the detection (the name of module which detected this element).
14    pub source: String,
15    /// Text within this element. This field is filled after OCR.
16    pub text: Option<String>,
17}
18
19impl LayoutElement {
20    /// Constructs a [`LayoutElement`] instance.
21    pub fn new(
22        x1: f32,
23        y1: f32,
24        x2: f32,
25        y2: f32,
26        element_type: &str,
27        confidence: f32,
28        source: &str,
29    ) -> Self {
30        let bbox = Rect::new(coord! { x: x1, y: y1 }, coord! { x: x2, y: y2 });
31
32        Self {
33            bbox,
34            element_type: element_type.to_string(),
35            confidence,
36            text: None,
37            source: source.to_string(),
38        }
39    }
40
41    /// Constructs a [`LayoutElement`] instance with text.
42    pub fn new_with_text(
43        x1: f32,
44        y1: f32,
45        x2: f32,
46        y2: f32,
47        element_type: &str,
48        text: String,
49        confidence: f32,
50        source: &str,
51    ) -> Self {
52        let bbox = Rect::new(coord! { x: x1, y: y1 }, coord! { x: x2, y: y2 });
53
54        Self {
55            bbox,
56            element_type: element_type.to_string(),
57            confidence,
58            text: Some(text),
59            source: source.to_string(),
60        }
61    }
62
63    /// Pads the bounding box of a [`LayoutElement`]. Useful for OCRing the element.
64    pub fn pad(&mut self, padding: f32) {
65        self.bbox
66            .set_min(self.bbox.min() - coord! { x: padding, y: padding });
67        self.bbox
68            .set_max(self.bbox.max() + coord! { x: padding, y: padding });
69    }
70
71    /// Crop the section of the image according to the [`LayoutElement`]'s bounding box.
72    pub fn crop_from_image(&self, img: &image::DynamicImage) -> image::DynamicImage {
73        let (x1, y1) = (self.bbox.min().x as u32, self.bbox.min().y as u32);
74        let (width, height) = (self.bbox.width() as u32, self.bbox.height() as u32);
75
76        img.clone().crop(x1, y1, width, height)
77    }
78
79    /// Apply a transformation to the bounding box points.
80    pub fn transform(&mut self, transform: impl Fn(Coord<f32>) -> Coord<f32>) {
81        self.bbox.set_min(transform(self.bbox.min()));
82        self.bbox.set_max(transform(self.bbox.max()));
83    }
84}