oar_ocr_core/models/detection/
db.rs

1//! DB (Differentiable Binarization) Model
2//!
3//! This module provides a pure implementation of the DB text detection model.
4//! The model handles preprocessing, inference, and postprocessing independently of tasks.
5
6use crate::core::inference::{OrtInfer, TensorInput};
7use crate::core::{OCRError, validate_positive, validate_range};
8use crate::processors::{
9    BoundingBox, BoxType, DBPostProcess, DBPostProcessConfig, DetResizeForTest, ImageScaleInfo,
10    LimitType, NormalizeImage, ScoreMode, TensorLayout,
11};
12use image::{DynamicImage, GenericImageView, RgbImage};
13use std::path::Path;
14use tracing::debug;
15
16/// Configuration for DB model preprocessing.
17#[derive(Debug, Clone, Default)]
18pub struct DBPreprocessConfig {
19    /// Limit for the side length of the image
20    pub limit_side_len: Option<u32>,
21    /// Type of limit to apply
22    pub limit_type: Option<LimitType>,
23    /// Maximum side limit for the image
24    pub max_side_limit: Option<u32>,
25    /// Resize long dimension (alternative to limit_side_len)
26    pub resize_long: Option<u32>,
27}
28
29/// Configuration for DB model postprocessing.
30#[derive(Debug, Clone)]
31pub struct DBPostprocessConfig {
32    /// Pixel-level threshold for text detection
33    pub score_threshold: f32,
34    /// Box-level threshold for filtering detections
35    pub box_threshold: f32,
36    /// Expansion ratio for detected regions using Vatti clipping
37    pub unclip_ratio: f32,
38    /// Maximum number of candidate detections
39    pub max_candidates: usize,
40    /// Whether to use dilation
41    pub use_dilation: bool,
42    /// Score calculation mode
43    pub score_mode: ScoreMode,
44    /// Type of bounding box (Quad or Poly)
45    pub box_type: BoxType,
46}
47
48impl Default for DBPostprocessConfig {
49    fn default() -> Self {
50        Self {
51            score_threshold: 0.3,
52            box_threshold: 0.7,
53            unclip_ratio: 1.5,
54            max_candidates: 1000,
55            use_dilation: false,
56            score_mode: ScoreMode::Fast,
57            box_type: BoxType::Quad,
58        }
59    }
60}
61
62impl DBPostprocessConfig {
63    /// Validates the configuration parameters.
64    pub fn validate(&self) -> Result<(), OCRError> {
65        // Validate score_threshold is in [0, 1]
66        validate_range(self.score_threshold, 0.0, 1.0, "score_threshold")?;
67
68        // Validate box_threshold is in [0, 1]
69        validate_range(self.box_threshold, 0.0, 1.0, "box_threshold")?;
70
71        // Validate unclip_ratio is positive
72        validate_positive(self.unclip_ratio, "unclip_ratio")?;
73
74        // Validate max_candidates is positive
75        validate_positive(self.max_candidates, "max_candidates")?;
76
77        Ok(())
78    }
79}
80
81/// DB model output containing bounding boxes and confidence scores.
82#[derive(Debug, Clone)]
83pub struct DBModelOutput {
84    /// Detected bounding boxes for each image in the batch
85    pub boxes: Vec<Vec<BoundingBox>>,
86    /// Confidence scores for each bounding box
87    pub scores: Vec<Vec<f32>>,
88}
89
90/// Pure DB model implementation.
91///
92/// This model implements the core DB architecture and can be configured
93/// for different detection tasks through preprocessing and postprocessing configs.
94#[derive(Debug)]
95pub struct DBModel {
96    /// ONNX Runtime inference engine
97    inference: OrtInfer,
98    /// Image resizer for preprocessing
99    resizer: DetResizeForTest,
100    /// Image normalizer for preprocessing
101    normalizer: NormalizeImage,
102    /// Postprocessor for converting predictions to bounding boxes
103    postprocessor: DBPostProcess,
104}
105
106impl DBModel {
107    /// Creates a new DB model.
108    pub fn new(
109        inference: OrtInfer,
110        resizer: DetResizeForTest,
111        normalizer: NormalizeImage,
112        postprocessor: DBPostProcess,
113    ) -> Self {
114        Self {
115            inference,
116            resizer,
117            normalizer,
118            postprocessor,
119        }
120    }
121
122    /// Preprocesses images for detection.
123    pub fn preprocess(
124        &self,
125        images: Vec<RgbImage>,
126    ) -> Result<(ndarray::Array4<f32>, Vec<ImageScaleInfo>), OCRError> {
127        // Convert to DynamicImage
128        let dynamic_images: Vec<DynamicImage> =
129            images.into_iter().map(DynamicImage::ImageRgb8).collect();
130
131        // Apply detection resizing
132        let (resized_images, img_shapes) = self.resizer.apply(
133            dynamic_images,
134            None, // Use default limit_side_len
135            None, // Use default limit_type
136            None, // Use default max_side_limit
137        );
138
139        debug!("After resize: {} images", resized_images.len());
140        for (i, (img, shape)) in resized_images.iter().zip(&img_shapes).enumerate() {
141            debug!(
142                "  Image {}: {}x{}, shape=[src_h={:.0}, src_w={:.0}, ratio_h={:.3}, ratio_w={:.3}]",
143                i,
144                img.width(),
145                img.height(),
146                shape.src_h,
147                shape.src_w,
148                shape.ratio_h,
149                shape.ratio_w
150            );
151        }
152
153        // Apply ImageNet normalization and convert to tensor.
154        //
155        // Note: External models often decode images as BGR and then normalize with
156        // mean/std as provided in their configs. In this repo, input images are
157        // loaded as RGB; we keep them in RGB here and rely on `NormalizeImage`
158        // with `ColorOrder::BGR` to map channels (RGB -> BGR) without a manual swap.
159        let batch_tensor = self.normalizer.normalize_batch_to(resized_images)?;
160        debug!("Batch tensor shape: {:?}", batch_tensor.shape());
161
162        Ok((batch_tensor, img_shapes))
163    }
164
165    fn resize_images(&self, images: Vec<RgbImage>) -> Vec<(usize, DynamicImage, ImageScaleInfo)> {
166        let dynamic_images: Vec<DynamicImage> =
167            images.into_iter().map(DynamicImage::ImageRgb8).collect();
168        let (resized_images, img_shapes) = self.resizer.apply(
169            dynamic_images,
170            None, // Use default limit_side_len
171            None, // Use default limit_type
172            None, // Use default max_side_limit
173        );
174
175        resized_images
176            .into_iter()
177            .zip(img_shapes)
178            .enumerate()
179            .map(|(idx, (image, shape))| (idx, image, shape))
180            .collect()
181    }
182
183    /// Runs inference on the preprocessed batch.
184    pub fn infer(
185        &self,
186        batch_tensor: &ndarray::Array4<f32>,
187    ) -> Result<ndarray::Array4<f32>, OCRError> {
188        let input_name = self.inference.input_name();
189        let inputs = vec![(input_name, TensorInput::Array4(batch_tensor))];
190
191        let outputs = self
192            .inference
193            .infer(&inputs)
194            .map_err(|e| OCRError::Inference {
195                model_name: "DB".to_string(),
196                context: format!(
197                    "failed to run inference on batch with shape {:?}",
198                    batch_tensor.shape()
199                ),
200                source: Box::new(e),
201            })?;
202
203        let output = outputs
204            .into_iter()
205            .next()
206            .ok_or_else(|| OCRError::InvalidInput {
207                message: "DB: no output returned from inference".to_string(),
208            })?;
209
210        output
211            .1
212            .try_into_array4_f32()
213            .map_err(|e| OCRError::Inference {
214                model_name: "DB".to_string(),
215                context: "failed to convert output to 4D array".to_string(),
216                source: Box::new(e),
217            })
218    }
219
220    /// Postprocesses model predictions to bounding boxes.
221    pub fn postprocess(
222        &self,
223        predictions: &ndarray::Array4<f32>,
224        img_shapes: Vec<ImageScaleInfo>,
225        score_threshold: f32,
226        box_threshold: f32,
227        unclip_ratio: f32,
228    ) -> DBModelOutput {
229        let config = DBPostProcessConfig::new(score_threshold, box_threshold, unclip_ratio);
230        let (boxes, scores) = self
231            .postprocessor
232            .apply(predictions, img_shapes, Some(&config));
233        DBModelOutput { boxes, scores }
234    }
235
236    fn forward_resized_batch(
237        &self,
238        batch: Vec<(usize, DynamicImage, ImageScaleInfo)>,
239        boxes_by_image: &mut [Vec<BoundingBox>],
240        scores_by_image: &mut [Vec<f32>],
241        score_threshold: f32,
242        box_threshold: f32,
243        unclip_ratio: f32,
244    ) -> Result<(), OCRError> {
245        let mut original_indices = Vec::with_capacity(batch.len());
246        let mut resized_images = Vec::with_capacity(batch.len());
247        let mut img_shapes = Vec::with_capacity(batch.len());
248
249        for (original_idx, resized_image, img_shape) in batch {
250            original_indices.push(original_idx);
251            resized_images.push(resized_image);
252            img_shapes.push(img_shape);
253        }
254
255        let batch_tensor = self.normalizer.normalize_batch_to(resized_images)?;
256        let predictions = self.infer(&batch_tensor)?;
257        let group_output = self.postprocess(
258            &predictions,
259            img_shapes,
260            score_threshold,
261            box_threshold,
262            unclip_ratio,
263        );
264
265        for (group_idx, original_idx) in original_indices.into_iter().enumerate() {
266            boxes_by_image[original_idx] = group_output
267                .boxes
268                .get(group_idx)
269                .cloned()
270                .unwrap_or_default();
271            scores_by_image[original_idx] = group_output
272                .scores
273                .get(group_idx)
274                .cloned()
275                .unwrap_or_default();
276        }
277
278        Ok(())
279    }
280
281    /// Runs the complete forward pass: preprocess -> infer -> postprocess.
282    pub fn forward(
283        &self,
284        images: Vec<RgbImage>,
285        score_threshold: f32,
286        box_threshold: f32,
287        unclip_ratio: f32,
288    ) -> Result<DBModelOutput, OCRError> {
289        if images.is_empty() {
290            return Ok(DBModelOutput {
291                boxes: Vec::new(),
292                scores: Vec::new(),
293            });
294        }
295
296        let image_count = images.len();
297        let resized = self.resize_images(images);
298        let mut groups: Vec<Vec<(usize, DynamicImage, ImageScaleInfo)>> = Vec::new();
299
300        for item in resized {
301            let dims = item.1.dimensions();
302            if let Some(group) = groups
303                .iter_mut()
304                .find(|group| group.first().map(|(_, img, _)| img.dimensions()) == Some(dims))
305            {
306                group.push(item);
307            } else {
308                groups.push(vec![item]);
309            }
310        }
311
312        debug!(
313            "DB forward: {} images grouped into {} shape batch(es)",
314            image_count,
315            groups.len()
316        );
317
318        let mut boxes_by_image = vec![Vec::new(); image_count];
319        let mut scores_by_image = vec![Vec::new(); image_count];
320
321        for group in groups {
322            self.forward_resized_batch(
323                group,
324                &mut boxes_by_image,
325                &mut scores_by_image,
326                score_threshold,
327                box_threshold,
328                unclip_ratio,
329            )?;
330        }
331
332        Ok(DBModelOutput {
333            boxes: boxes_by_image,
334            scores: scores_by_image,
335        })
336    }
337}
338
339/// Builder for DB model.
340pub struct DBModelBuilder {
341    /// Preprocessing configuration
342    preprocess_config: DBPreprocessConfig,
343    /// Postprocessing configuration
344    postprocess_config: DBPostprocessConfig,
345    /// ONNX Runtime session configuration
346    ort_config: Option<crate::core::config::OrtSessionConfig>,
347}
348
349impl DBModelBuilder {
350    /// Creates a new DB model builder with default settings.
351    pub fn new() -> Self {
352        Self {
353            preprocess_config: DBPreprocessConfig::default(),
354            postprocess_config: DBPostprocessConfig::default(),
355            ort_config: None,
356        }
357    }
358
359    /// Sets the preprocessing configuration.
360    pub fn preprocess_config(mut self, config: DBPreprocessConfig) -> Self {
361        self.preprocess_config = config;
362        self
363    }
364
365    /// Sets the postprocessing configuration.
366    pub fn postprocess_config(mut self, config: DBPostprocessConfig) -> Self {
367        self.postprocess_config = config;
368        self
369    }
370
371    /// Sets the ONNX Runtime session configuration.
372    pub fn with_ort_config(mut self, config: crate::core::config::OrtSessionConfig) -> Self {
373        self.ort_config = Some(config);
374        self
375    }
376
377    /// Builds the DB model.
378    pub fn build(self, model_path: &Path) -> Result<DBModel, OCRError> {
379        // Create ONNX inference engine
380        let inference = if self.ort_config.is_some() {
381            use crate::core::config::ModelInferenceConfig;
382            let common_config = ModelInferenceConfig {
383                ort_session: self.ort_config,
384                ..Default::default()
385            };
386            OrtInfer::from_config(&common_config, model_path, Some("x"))?
387        } else {
388            OrtInfer::new(model_path, Some("x"))?
389        };
390
391        // Create resizer
392        let resizer = DetResizeForTest::new(
393            None,                                  // input_shape
394            None,                                  // image_shape
395            None,                                  // keep_ratio
396            self.preprocess_config.limit_side_len, // limit_side_len
397            self.preprocess_config.limit_type,     // limit_type
398            self.preprocess_config.resize_long,    // resize_long
399            self.preprocess_config.max_side_limit, // max_side_limit
400        );
401
402        // Create normalizer.
403        // External models read images in BGR. Their configs use ImageNet stats
404        // in that *same* channel order (B, G, R). Our images are loaded as RGB,
405        // so we keep them in RGB and use `ColorOrder::BGR` to map channels
406        // into BGR order during normalization.
407        let normalizer = NormalizeImage::with_color_order(
408            Some(1.0 / 255.0),               // scale
409            Some(vec![0.485, 0.456, 0.406]), // mean
410            Some(vec![0.229, 0.224, 0.225]), // std
411            Some(TensorLayout::CHW),         // order
412            Some(crate::processors::types::ColorOrder::BGR),
413        )?;
414
415        // Create postprocessor
416        let postprocessor = DBPostProcess::new(
417            Some(self.postprocess_config.score_threshold),
418            Some(self.postprocess_config.box_threshold),
419            Some(self.postprocess_config.max_candidates),
420            Some(self.postprocess_config.unclip_ratio),
421            Some(self.postprocess_config.use_dilation),
422            Some(self.postprocess_config.score_mode),
423            Some(self.postprocess_config.box_type),
424        );
425
426        Ok(DBModel::new(inference, resizer, normalizer, postprocessor))
427    }
428}
429
430impl Default for DBModelBuilder {
431    fn default() -> Self {
432        Self::new()
433    }
434}
oar_ocr_core/models/detection/db.rs

oar_ocr_core/models/detection/
db.rs
