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

oar_ocr_core/models/detection/
db.rs
