oar_ocr_core/models/detection/
rtdetr.rs

1//! RT-DETR Layout Detection Model
2//!
3//! This module provides a pure implementation of the RT-DETR model for layout detection.
4//! The model is independent of any specific task and can be reused in different contexts.
5
6use crate::core::OCRError;
7use crate::core::inference::{OrtInfer, TensorInput};
8use crate::processors::{
9    DetResizeForTest, ImageScaleInfo, LimitType, NormalizeImage, TensorLayout,
10};
11use image::{DynamicImage, RgbImage};
12use ndarray::Array2;
13
14type RTDetrPreprocessArtifacts = (
15    ndarray::Array4<f32>,
16    Vec<ImageScaleInfo>,
17    Vec<[f32; 2]>,
18    Vec<[f32; 2]>,
19);
20type RTDetrPreprocessResult = Result<RTDetrPreprocessArtifacts, OCRError>;
21
22/// Preprocessing configuration for RT-DETR model.
23#[derive(Debug, Clone)]
24pub struct RTDetrPreprocessConfig {
25    /// Target image shape (height, width)
26    pub image_shape: (u32, u32),
27    /// Whether to keep aspect ratio when resizing
28    pub keep_ratio: bool,
29    /// Limit side length
30    pub limit_side_len: u32,
31    /// Normalization scale factor
32    pub scale: f32,
33    /// Normalization mean values (RGB)
34    pub mean: Vec<f32>,
35    /// Normalization std values (RGB)
36    pub std: Vec<f32>,
37}
38
39impl Default for RTDetrPreprocessConfig {
40    fn default() -> Self {
41        Self {
42            image_shape: (640, 640),
43            keep_ratio: false,
44            limit_side_len: 640,
45            scale: 1.0 / 255.0,
46            // Paddle's RT-DETR exports expect inputs scaled to [0,1] without mean/std shift
47            mean: vec![0.0, 0.0, 0.0],
48            std: vec![1.0, 1.0, 1.0],
49        }
50    }
51}
52
53/// Postprocessing configuration for RT-DETR model.
54#[derive(Debug, Clone)]
55pub struct RTDetrPostprocessConfig {
56    /// Number of classes
57    pub num_classes: usize,
58}
59
60/// Output from RT-DETR model.
61#[derive(Debug, Clone)]
62pub struct RTDetrModelOutput {
63    /// Detection predictions tensor [batch_size, num_detections, 6]
64    /// Each detection: [x1, y1, x2, y2, score, class_id]
65    pub predictions: ndarray::Array4<f32>,
66}
67
68/// RT-DETR layout detection model.
69///
70/// This is a pure model implementation that handles:
71/// - Preprocessing: Image resizing and normalization
72/// - Inference: Running the ONNX model
73/// - Postprocessing: Returning raw predictions
74///
75/// The model is independent of any specific task or adapter.
76#[derive(Debug)]
77pub struct RTDetrModel {
78    inference: OrtInfer,
79    resizer: DetResizeForTest,
80    normalizer: NormalizeImage,
81    _preprocess_config: RTDetrPreprocessConfig,
82}
83
84impl RTDetrModel {
85    /// Creates a new RT-DETR model.
86    pub fn new(
87        inference: OrtInfer,
88        preprocess_config: RTDetrPreprocessConfig,
89    ) -> Result<Self, OCRError> {
90        // Create resizer
91        let resizer = DetResizeForTest::new(
92            None,
93            Some((
94                preprocess_config.image_shape.0,
95                preprocess_config.image_shape.1,
96            )),
97            Some(preprocess_config.keep_ratio),
98            Some(preprocess_config.limit_side_len),
99            Some(LimitType::Max),
100            None,
101            None,
102        );
103
104        // Create normalizer.
105        // Paddle models expect BGR input; treat config mean/std as RGB and reorder.
106        let normalizer = NormalizeImage::with_color_order_from_rgb_stats(
107            Some(preprocess_config.scale),
108            preprocess_config.mean.clone(),
109            preprocess_config.std.clone(),
110            Some(TensorLayout::CHW),
111            crate::processors::types::ColorOrder::BGR,
112        )?;
113
114        Ok(Self {
115            inference,
116            resizer,
117            normalizer,
118            _preprocess_config: preprocess_config,
119        })
120    }
121
122    /// Preprocesses images for RT-DETR model.
123    ///
124    /// Returns:
125    /// - Batch tensor ready for inference
126    /// - Image shapes after resizing [h, w, ratio_h, ratio_w]
127    /// - Original shapes [h, w]
128    /// - Resized shapes [h, w]
129    pub fn preprocess(&self, images: Vec<RgbImage>) -> RTDetrPreprocessResult {
130        // Store original dimensions
131        let orig_shapes: Vec<[f32; 2]> = images
132            .iter()
133            .map(|img| [img.height() as f32, img.width() as f32])
134            .collect();
135
136        // Convert to DynamicImage
137        let dynamic_images: Vec<DynamicImage> =
138            images.into_iter().map(DynamicImage::ImageRgb8).collect();
139
140        // Resize images
141        let (resized_images, img_shapes) = self.resizer.apply(
142            dynamic_images,
143            None, // Use configured limit_side_length
144            None, // Use configured limit_type
145            None,
146        );
147
148        // Get resized dimensions
149        let resized_shapes: Vec<[f32; 2]> = resized_images
150            .iter()
151            .map(|img| [img.height() as f32, img.width() as f32])
152            .collect();
153
154        // Normalize and convert to tensor
155        let batch_tensor = self.normalizer.normalize_batch_to(resized_images)?;
156
157        Ok((batch_tensor, img_shapes, orig_shapes, resized_shapes))
158    }
159
160    /// Runs inference on the preprocessed batch tensor.
161    ///
162    /// RT-DETR requires both `scale_factor` and `im_shape` inputs.
163    pub fn infer(
164        &self,
165        batch_tensor: &ndarray::Array4<f32>,
166        scale_factor: &Array2<f32>,
167        im_shape: &Array2<f32>,
168    ) -> Result<ndarray::Array4<f32>, OCRError> {
169        let inputs = vec![
170            ("image", TensorInput::Array4(batch_tensor)),
171            ("scale_factor", TensorInput::Array2(scale_factor)),
172            ("im_shape", TensorInput::Array2(im_shape)),
173        ];
174
175        let outputs = self
176            .inference
177            .infer(&inputs)
178            .map_err(|e| OCRError::Inference {
179                model_name: "RT-DETR".to_string(),
180                context: "failed to run inference".to_string(),
181                source: Box::new(e),
182            })?;
183
184        // Find primary output (RT-DETR uses "fetch_name_0", fall back to first output)
185        let output = outputs
186            .iter()
187            .find(|(name, _)| name == "fetch_name_0")
188            .or_else(|| outputs.first())
189            .ok_or_else(|| OCRError::InvalidInput {
190                message: "RT-DETR: no outputs found from model".to_string(),
191            })?;
192
193        let output_shape = output.1.shape();
194
195        // RT-DETR can output either 2D or 4D format
196        match output_shape.len() {
197            2 => {
198                // 2D format: [num_boxes, N] where N >= 6
199                // Convert to 4D format [1, num_boxes, 1, N]
200                let output_array =
201                    output
202                        .1
203                        .clone()
204                        .try_into_array_f32()
205                        .map_err(|e| OCRError::InvalidInput {
206                            message: format!("Failed to extract output tensor: {}", e),
207                        })?;
208                let num_boxes = output_shape[0] as usize;
209                let box_dim = output_shape[1] as usize;
210                let (data, _offset) = output_array.into_raw_vec_and_offset();
211                ndarray::Array::from_shape_vec((1, num_boxes, 1, box_dim), data).map_err(|e| {
212                    OCRError::InvalidInput {
213                        message: format!("Failed to reshape 2D output to 4D: {}", e),
214                    }
215                })
216            }
217            4 => {
218                // Standard 4D output format
219                output
220                    .1
221                    .clone()
222                    .try_into_array4_f32()
223                    .map_err(|e| OCRError::InvalidInput {
224                        message: format!("Failed to convert to 4D array: {}", e),
225                    })
226            }
227            _ => Err(OCRError::InvalidInput {
228                message: format!(
229                    "RT-DETR inference: expected 2D or 4D output, got {}D with shape {:?}",
230                    output_shape.len(),
231                    output_shape
232                ),
233            }),
234        }
235    }
236
237    /// Postprocesses model predictions.
238    ///
239    /// For RT-DETR, we just return the raw predictions.
240    /// The adapter layer will handle converting these to task-specific outputs.
241    pub fn postprocess(
242        &self,
243        predictions: ndarray::Array4<f32>,
244        _config: &RTDetrPostprocessConfig,
245    ) -> Result<RTDetrModelOutput, OCRError> {
246        Ok(RTDetrModelOutput { predictions })
247    }
248
249    /// Runs the complete forward pass: preprocess -> infer -> postprocess.
250    pub fn forward(
251        &self,
252        images: Vec<RgbImage>,
253        config: &RTDetrPostprocessConfig,
254    ) -> Result<(RTDetrModelOutput, Vec<ImageScaleInfo>), OCRError> {
255        let (batch_tensor, img_shapes, _orig_shapes, resized_shapes) = self.preprocess(images)?;
256
257        let batch_size = batch_tensor.shape()[0];
258
259        // Build scale_factor array [ratio_h, ratio_w]
260        let scale_data: Vec<f32> = img_shapes
261            .iter()
262            .flat_map(|shape| [shape.ratio_h, shape.ratio_w])
263            .collect();
264        let scale_factor = Array2::from_shape_vec((batch_size, 2), scale_data).map_err(|e| {
265            OCRError::InvalidInput {
266                message: format!("Failed to create scale_factor array: {}", e),
267            }
268        })?;
269
270        // Build im_shape array using resized dimensions
271        let im_shape_data: Vec<f32> = resized_shapes
272            .iter()
273            .flat_map(|shape| [shape[0], shape[1]])
274            .collect();
275        let im_shape = Array2::from_shape_vec((batch_size, 2), im_shape_data).map_err(|e| {
276            OCRError::InvalidInput {
277                message: format!("Failed to create im_shape array: {}", e),
278            }
279        })?;
280
281        let predictions = self.infer(&batch_tensor, &scale_factor, &im_shape)?;
282        let output = self.postprocess(predictions, config)?;
283        Ok((output, img_shapes))
284    }
285}
286
287/// Builder for RT-DETR model.
288#[derive(Debug, Default)]
289pub struct RTDetrModelBuilder {
290    preprocess_config: Option<RTDetrPreprocessConfig>,
291}
292
293impl RTDetrModelBuilder {
294    /// Creates a new builder.
295    pub fn new() -> Self {
296        Self::default()
297    }
298
299    /// Sets the preprocessing configuration.
300    pub fn preprocess_config(mut self, config: RTDetrPreprocessConfig) -> Self {
301        self.preprocess_config = Some(config);
302        self
303    }
304
305    /// Sets the image shape.
306    pub fn image_shape(mut self, height: u32, width: u32) -> Self {
307        let mut config = self.preprocess_config.unwrap_or_default();
308        config.image_shape = (height, width);
309        config.limit_side_len = height.max(width);
310        self.preprocess_config = Some(config);
311        self
312    }
313
314    /// Builds the RT-DETR model.
315    pub fn build(self, inference: OrtInfer) -> Result<RTDetrModel, OCRError> {
316        let preprocess_config = self.preprocess_config.unwrap_or_default();
317        RTDetrModel::new(inference, preprocess_config)
318    }
319}
oar_ocr_core/models/detection/rtdetr.rs

oar_ocr_core/models/detection/
rtdetr.rs
