oar_ocr_core/processors/db_postprocess.rs
1//! Post-processing for DB (Differentiable Binarization) text detection models.
2//!
3//! The [`DBPostProcess`] struct converts raw detection heatmaps into geometric
4//! bounding boxes by thresholding, contour extraction, scoring, and optional
5//! polygonal post-processing. Supporting functionality (bitmap extraction,
6//! scoring, mask morphology) is split across helper modules within this
7//! directory.
8
9#[path = "db_bitmap.rs"]
10mod db_bitmap;
11#[path = "db_mask.rs"]
12mod db_mask;
13#[path = "db_score.rs"]
14mod db_score;
15
16use crate::processors::geometry::BoundingBox;
17use crate::processors::types::{BoxType, ImageScaleInfo, ScoreMode};
18use ndarray::Axis;
19use rayon::prelude::*;
20
21/// Runtime configuration for DB post-processing.
22///
23/// This struct contains parameters that may vary per inference call,
24/// such as detection thresholds and expansion ratios.
25#[derive(Debug, Clone)]
26pub struct DBPostProcessConfig {
27 /// Threshold for binarizing the prediction map.
28 pub thresh: f32,
29 /// Threshold for filtering bounding boxes based on their score.
30 pub box_thresh: f32,
31 /// Ratio for unclipping (expanding) bounding boxes.
32 pub unclip_ratio: f32,
33}
34
35impl DBPostProcessConfig {
36 /// Creates a new runtime config with specified values.
37 pub fn new(thresh: f32, box_thresh: f32, unclip_ratio: f32) -> Self {
38 Self {
39 thresh,
40 box_thresh,
41 unclip_ratio,
42 }
43 }
44}
45
46/// Post-processor for DB (Differentiable Binarization) text detection models.
47#[derive(Debug)]
48pub struct DBPostProcess {
49 /// Default threshold for binarizing the prediction map (default: 0.3).
50 pub thresh: f32,
51 /// Default threshold for filtering bounding boxes based on their score (default: 0.6).
52 pub box_thresh: f32,
53 /// Maximum number of candidate bounding boxes to consider (default: 1000).
54 pub max_candidates: usize,
55 /// Default ratio for unclipping (expanding) bounding boxes (default: 1.5).
56 pub unclip_ratio: f32,
57 /// Minimum side length for detected bounding boxes.
58 pub min_size: f32,
59 /// Method for calculating the score of a bounding box.
60 pub score_mode: ScoreMode,
61 /// Type of bounding box to generate (quadrilateral or polygon).
62 pub box_type: BoxType,
63 /// Whether to apply dilation to the segmentation mask before contour detection.
64 pub use_dilation: bool,
65}
66
67impl DBPostProcess {
68 /// Creates a new `DBPostProcess` instance with optional overrides.
69 pub fn new(
70 thresh: Option<f32>,
71 box_thresh: Option<f32>,
72 max_candidates: Option<usize>,
73 unclip_ratio: Option<f32>,
74 use_dilation: Option<bool>,
75 score_mode: Option<ScoreMode>,
76 box_type: Option<BoxType>,
77 ) -> Self {
78 Self {
79 thresh: thresh.unwrap_or(0.3),
80 box_thresh: box_thresh.unwrap_or(0.6),
81 max_candidates: max_candidates.unwrap_or(1000),
82 unclip_ratio: unclip_ratio.unwrap_or(1.5),
83 min_size: 3.0,
84 score_mode: score_mode.unwrap_or(ScoreMode::Fast),
85 box_type: box_type.unwrap_or(BoxType::Quad),
86 use_dilation: use_dilation.unwrap_or(false),
87 }
88 }
89
90 /// Applies post-processing to a batch of prediction maps.
91 ///
92 /// # Arguments
93 /// * `preds` - Model predictions (batch of heatmaps)
94 /// * `img_shapes` - Original image dimensions for each image in batch
95 /// * `config` - Runtime configuration for thresholds and ratios.
96 /// If `None`, uses the default values stored in this processor.
97 ///
98 /// # Returns
99 /// Tuple of (bounding_boxes, scores) for each image in batch
100 pub fn apply(
101 &self,
102 preds: &ndarray::Array4<f32>,
103 img_shapes: Vec<ImageScaleInfo>,
104 config: Option<&DBPostProcessConfig>,
105 ) -> (Vec<Vec<BoundingBox>>, Vec<Vec<f32>>) {
106 // Use provided config or fall back to stored defaults
107 let thresh = config.map(|c| c.thresh).unwrap_or(self.thresh);
108 let box_thresh = config.map(|c| c.box_thresh).unwrap_or(self.box_thresh);
109 let unclip_ratio = config.map(|c| c.unclip_ratio).unwrap_or(self.unclip_ratio);
110
111 // Process per-image: each batch entry is independent, so we can
112 // process them in parallel via rayon. The body of `process` itself
113 // parallelises the thresholding step below, so we keep this loop
114 // serial here — fanning out the contour/postprocess work for a
115 // small batch (<= a few images) rarely pays back the rayon setup
116 // cost, but the dominant cost is the thresholded bitmap build
117 // which is now multi-threaded.
118 let mut all_boxes = Vec::with_capacity(img_shapes.len());
119 let mut all_scores = Vec::with_capacity(img_shapes.len());
120
121 for (batch_idx, shape_batch) in img_shapes.iter().enumerate() {
122 let pred_slice = preds.index_axis(Axis(0), batch_idx);
123 let pred_channel = pred_slice.index_axis(Axis(0), 0);
124
125 let (boxes, scores) =
126 self.process(&pred_channel, shape_batch, thresh, box_thresh, unclip_ratio);
127 all_boxes.push(boxes);
128 all_scores.push(scores);
129 }
130
131 (all_boxes, all_scores)
132 }
133
134 fn process(
135 &self,
136 pred: &ndarray::ArrayView2<f32>,
137 img_shape: &ImageScaleInfo,
138 thresh: f32,
139 box_thresh: f32,
140 unclip_ratio: f32,
141 ) -> (Vec<BoundingBox>, Vec<f32>) {
142 let src_h = img_shape.src_h as u32;
143 let src_w = img_shape.src_w as u32;
144
145 let height = pred.shape()[0] as u32;
146 let width = pred.shape()[1] as u32;
147
148 tracing::debug!(
149 "DBPostProcess: pred {}x{}, src {}x{} (dest dimensions)",
150 height,
151 width,
152 src_h,
153 src_w
154 );
155
156 // Build binary mask directly as a `GrayImage` buffer. The previous
157 // implementation called `put_pixel` per element, which involves a
158 // bounds check and a function call for every pixel. Writing into
159 // the buffer with a single pass over the prediction map (chunked
160 // across rows) is several times faster and trivially parallelises
161 // over rows via rayon.
162 let mask_img = self.threshold_to_mask(pred, thresh);
163
164 // Apply dilation if needed
165 let mask_img = if self.use_dilation {
166 self.dilate_mask_img(&mask_img)
167 } else {
168 mask_img
169 };
170
171 match self.box_type {
172 BoxType::Poly => {
173 self.polygons_from_bitmap(pred, &mask_img, src_w, src_h, box_thresh, unclip_ratio)
174 }
175 BoxType::Quad => {
176 self.boxes_from_bitmap(pred, &mask_img, src_w, src_h, box_thresh, unclip_ratio)
177 }
178 }
179 }
180
181 /// Builds a binary segmentation mask from a prediction heatmap.
182 ///
183 /// Writes into the `GrayImage` row buffer directly (no per-pixel
184 /// `put_pixel` overhead) and parallelises the work over rows.
185 fn threshold_to_mask(&self, pred: &ndarray::ArrayView2<f32>, thresh: f32) -> image::GrayImage {
186 let height = pred.shape()[0] as u32;
187 let width = pred.shape()[1] as u32;
188 let mut mask_img = image::GrayImage::new(width, height);
189 let buf: &mut [u8] = mask_img.as_mut();
190 let row_bytes = width as usize;
191
192 // Fill the mask row by row. Pick between serial and parallel branches
193 // based on the number of rows: for small masks the rayon per-row
194 // overhead dominates; for typical 800x800 DB outputs (640x640
195 // predictions), the parallel branch wins by a wide margin.
196 let fill_row = |y: usize, row: &mut [u8]| {
197 let pred_row = pred.row(y);
198 // When the row is contiguous (typical row-major `pred`), iterate it
199 // as a slice so the compiler can drop bounds checks and vectorize.
200 if let Some(slice) = pred_row.as_slice() {
201 for (dst, &val) in row.iter_mut().zip(slice) {
202 *dst = if val > thresh { 255 } else { 0 };
203 }
204 } else {
205 for (x, dst) in row.iter_mut().enumerate() {
206 *dst = if pred_row[x] > thresh { 255 } else { 0 };
207 }
208 }
209 };
210 if height >= 64 {
211 buf.par_chunks_mut(row_bytes)
212 .enumerate()
213 .for_each(|(y, row)| fill_row(y, row));
214 } else {
215 buf.chunks_mut(row_bytes)
216 .enumerate()
217 .for_each(|(y, row)| fill_row(y, row));
218 }
219
220 mask_img
221 }
222}