oar_ocr/processors/postprocess/
db_score.rs

1use super::DBPostProcess;
2use crate::processors::geometry::{BoundingBox, ScanlineBuffer};
3use itertools::Itertools;
4use rayon::prelude::*;
5
6#[derive(Debug, Clone, Copy)]
7struct Region {
8    start_y: usize,
9    end_y: usize,
10    start_x: usize,
11    end_x: usize,
12}
13
14impl Region {
15    fn new(start_y: usize, end_y: usize, start_x: usize, end_x: usize) -> Self {
16        Self {
17            start_y,
18            end_y,
19            start_x,
20            end_x,
21        }
22    }
23
24    fn height(&self) -> usize {
25        self.end_y - self.start_y
26    }
27
28    fn width(&self) -> usize {
29        self.end_x - self.start_x
30    }
31}
32
33impl DBPostProcess {
34    /// Calculates the score of a bounding box using a fast approximation method.
35    pub fn box_score_fast(&self, pred: &ndarray::Array2<f32>, bbox: &BoundingBox) -> f32 {
36        let height = pred.shape()[0];
37        let width = pred.shape()[1];
38
39        let (min_x, max_x) = bbox
40            .points
41            .iter()
42            .map(|p| p.x)
43            .minmax()
44            .into_option()
45            .unwrap_or((0.0, 0.0));
46        let (min_y, max_y) = bbox
47            .points
48            .iter()
49            .map(|p| p.y)
50            .minmax()
51            .into_option()
52            .unwrap_or((0.0, 0.0));
53
54        let min_x = min_x.max(0.0).min(width as f32 - 1.0);
55        let max_x = max_x.max(0.0).min(width as f32 - 1.0);
56        let min_y = min_y.max(0.0).min(height as f32 - 1.0);
57        let max_y = max_y.max(0.0).min(height as f32 - 1.0);
58
59        let start_y = min_y as usize;
60        let end_y = max_y as usize + 1;
61        let start_x = min_x as usize;
62        let end_x = max_x as usize + 1;
63
64        self.box_score_fast_contour(pred, bbox, start_y, end_y, start_x, end_x)
65    }
66
67    fn box_score_fast_contour(
68        &self,
69        pred: &ndarray::Array2<f32>,
70        bbox: &BoundingBox,
71        start_y: usize,
72        end_y: usize,
73        start_x: usize,
74        end_x: usize,
75    ) -> f32 {
76        let region = Region::new(start_y, end_y, start_x, end_x);
77        self.box_score_fast_contour_with_policy(pred, bbox, region, None)
78    }
79
80    fn box_score_fast_contour_with_policy(
81        &self,
82        pred: &ndarray::Array2<f32>,
83        bbox: &BoundingBox,
84        region: Region,
85        policy: Option<&crate::core::config::ParallelPolicy>,
86    ) -> f32 {
87        let region_height = region.height();
88        let region_width = region.width();
89
90        let max_polygon_points = bbox.points.len();
91        let mut scanline_buffer = ScanlineBuffer::new(max_polygon_points);
92
93        let pixel_threshold = policy
94            .map(|p| p.postprocess_pixel_threshold)
95            .unwrap_or(8_000);
96
97        if region_height * region_width < pixel_threshold {
98            let mut total_score = 0.0;
99            let mut total_pixels = 0;
100
101            for y in region.start_y..region.end_y {
102                let scanline_y = y as f32 + 0.5;
103                let (line_score, line_pixels) = scanline_buffer.process_scanline(
104                    scanline_y,
105                    bbox,
106                    region.start_x,
107                    region.end_x,
108                    pred,
109                );
110                total_score += line_score;
111                total_pixels += line_pixels;
112            }
113
114            if total_pixels > 0 {
115                total_score / total_pixels as f32
116            } else {
117                0.0
118            }
119        } else {
120            let scanline_results: Vec<(f32, usize)> = (region.start_y..region.end_y)
121                .into_par_iter()
122                .map(|y| {
123                    let scanline_y = y as f32 + 0.5;
124
125                    let mut thread_buffer = ScanlineBuffer::new(max_polygon_points);
126                    thread_buffer.process_scanline(
127                        scanline_y,
128                        bbox,
129                        region.start_x,
130                        region.end_x,
131                        pred,
132                    )
133                })
134                .collect();
135
136            let total_score: f32 = scanline_results.iter().map(|(score, _)| score).sum();
137            let total_pixels: usize = scanline_results.iter().map(|(_, pixels)| pixels).sum();
138
139            if total_pixels > 0 {
140                total_score / total_pixels as f32
141            } else {
142                0.0
143            }
144        }
145    }
146
147    pub(super) fn box_score_slow(
148        &self,
149        pred: &ndarray::Array2<f32>,
150        contour: &imageproc::contours::Contour<u32>,
151    ) -> f32 {
152        let mut total_score = 0.0;
153        let mut pixel_count = 0;
154
155        for point in &contour.points {
156            let x = point.x as usize;
157            let y = point.y as usize;
158
159            if y < pred.shape()[0] && x < pred.shape()[1] {
160                total_score += pred[[y, x]];
161                pixel_count += 1;
162            }
163        }
164
165        if pixel_count > 0 {
166            total_score / pixel_count as f32
167        } else {
168            0.0
169        }
170    }
171}
oar_ocr/processors/postprocess/db_score.rs

oar_ocr/processors/postprocess/
db_score.rs
