oar_ocr_core/models/classification/
pp_lcnet.rs

1//! PP-LCNet Classification Model
2//!
3//! This module provides a pure implementation of the PP-LCNet model for image classification.
4//! PP-LCNet is a lightweight classification network that can be used for various classification
5//! tasks such as document orientation and text line orientation.
6
7use crate::core::OCRError;
8use crate::core::inference::{OrtInfer, TensorInput};
9use crate::domain::adapters::preprocessing::rgb_to_dynamic;
10use crate::processors::{NormalizeImage, TensorLayout};
11use crate::utils::topk::Topk;
12use image::{RgbImage, imageops::FilterType};
13
14/// Configuration for PP-LCNet model preprocessing.
15#[derive(Debug, Clone)]
16pub struct PPLCNetPreprocessConfig {
17    /// Input shape (height, width)
18    pub input_shape: (u32, u32),
19    /// Resizing filter to use
20    pub resize_filter: FilterType,
21    /// When set, resize by short edge to this size (keep ratio) then center-crop to `input_shape`.
22    ///
23    /// This matches `ResizeImage(resize_short=256)` + `CropImage(size=224)` used by
24    /// most PP-LCNet classifiers (e.g. `PP-LCNet_x1_0_doc_ori`, `PP-LCNet_x1_0_table_cls`).
25    ///
26    /// Some specialized PP-LCNet classifiers (e.g. `PP-LCNet_x1_0_textline_ori`) use a direct
27    /// resize to a fixed `(height,width)` without short-edge resize/crop; set this to `None`
28    /// to match `ResizeImage(size=[w,h])`.
29    pub resize_short: Option<u32>,
30    /// Scaling factor applied before normalization (defaults to 1.0 / 255.0)
31    pub normalize_scale: f32,
32    /// Mean values for normalization
33    pub normalize_mean: Vec<f32>,
34    /// Standard deviation values for normalization
35    pub normalize_std: Vec<f32>,
36    /// Tensor data layout (CHW or HWC)
37    pub tensor_layout: TensorLayout,
38}
39
40impl Default for PPLCNetPreprocessConfig {
41    fn default() -> Self {
42        Self {
43            input_shape: (224, 224),
44            // Use cv2.INTER_LINEAR for PP-LCNet resize.
45            resize_filter: FilterType::Triangle,
46            // PP-LCNet classifiers default to resize_short=256 then center-crop.
47            resize_short: Some(256),
48            normalize_scale: 1.0 / 255.0,
49            normalize_mean: vec![0.485, 0.456, 0.406],
50            normalize_std: vec![0.229, 0.224, 0.225],
51            tensor_layout: TensorLayout::CHW,
52        }
53    }
54}
55
56/// Configuration for PP-LCNet model postprocessing.
57#[derive(Debug, Clone)]
58pub struct PPLCNetPostprocessConfig {
59    /// Class labels
60    pub labels: Vec<String>,
61    /// Number of top predictions to return
62    pub topk: usize,
63}
64
65impl Default for PPLCNetPostprocessConfig {
66    fn default() -> Self {
67        Self {
68            labels: vec![],
69            topk: 1,
70        }
71    }
72}
73
74/// Output from PP-LCNet model.
75#[derive(Debug, Clone)]
76pub struct PPLCNetModelOutput {
77    /// Predicted class IDs per image
78    pub class_ids: Vec<Vec<usize>>,
79    /// Confidence scores for each prediction
80    pub scores: Vec<Vec<f32>>,
81    /// Label names for each prediction (if labels provided)
82    pub label_names: Option<Vec<Vec<String>>>,
83}
84
85/// Pure PP-LCNet model implementation.
86///
87/// This model performs image classification using the PP-LCNet architecture.
88#[derive(Debug)]
89pub struct PPLCNetModel {
90    /// ONNX Runtime inference engine
91    inference: OrtInfer,
92    /// Image normalizer for preprocessing
93    normalizer: NormalizeImage,
94    /// Top-k processor for postprocessing
95    topk_processor: Topk,
96    /// Input shape (height, width)
97    input_shape: (u32, u32),
98    /// Resizing filter
99    resize_filter: FilterType,
100    /// Optional short-edge resize (see `PPLCNetPreprocessConfig::resize_short`)
101    resize_short: Option<u32>,
102}
103
104impl PPLCNetModel {
105    /// Creates a new PP-LCNet model.
106    pub fn new(
107        inference: OrtInfer,
108        normalizer: NormalizeImage,
109        topk_processor: Topk,
110        input_shape: (u32, u32),
111        resize_filter: FilterType,
112        resize_short: Option<u32>,
113    ) -> Self {
114        Self {
115            inference,
116            normalizer,
117            topk_processor,
118            input_shape,
119            resize_filter,
120            resize_short,
121        }
122    }
123
124    /// Preprocesses images for classification.
125    ///
126    /// # Arguments
127    ///
128    /// * `images` - Input images to preprocess
129    ///
130    /// # Returns
131    ///
132    /// Preprocessed batch tensor
133    pub fn preprocess(&self, images: Vec<RgbImage>) -> Result<ndarray::Array4<f32>, OCRError> {
134        let image_refs: Vec<&RgbImage> = images.iter().collect();
135        self.preprocess_refs(&image_refs)
136    }
137
138    /// Preprocesses borrowed images without cloning their source pixel buffers.
139    pub fn preprocess_refs(&self, images: &[&RgbImage]) -> Result<ndarray::Array4<f32>, OCRError> {
140        let (crop_h, crop_w) = self.input_shape;
141
142        let resized_rgb: Vec<RgbImage> = if let Some(resize_short) = self.resize_short {
143            // PP-LCNet classifier preprocessing:
144            // 1) Resize by short edge (keep ratio)
145            // 2) Center crop to the model input size
146            //
147            // This matches `ResizeImage(resize_short=256)` + `CropImage(size=224)` used by
148            // models like `PP-LCNet_x1_0_doc_ori` / `PP-LCNet_x1_0_table_cls`.
149            images
150                .iter()
151                .filter_map(|&img| {
152                    let (w, h) = (img.width(), img.height());
153                    if w == 0 || h == 0 {
154                        return None;
155                    }
156
157                    let short = w.min(h) as f32;
158                    let scale = (resize_short as f32) / short;
159                    let new_w = ((w as f32) * scale).round().max(crop_w as f32) as u32;
160                    let new_h = ((h as f32) * scale).round().max(crop_h as f32) as u32;
161
162                    let resized = image::imageops::resize(img, new_w, new_h, self.resize_filter);
163
164                    // Center crop to (crop_w, crop_h)
165                    let x1 = (new_w.saturating_sub(crop_w)) / 2;
166                    let y1 = (new_h.saturating_sub(crop_h)) / 2;
167                    let cropped =
168                        image::imageops::crop_imm(&resized, x1, y1, crop_w, crop_h).to_image();
169                    Some(cropped)
170                })
171                .collect()
172        } else {
173            // Direct resize to input shape (height,width) without crop.
174            // This matches `ResizeImage(size=[w,h])` used by
175            // `PP-LCNet_x1_0_textline_ori` (80x160).
176            images
177                .iter()
178                .filter_map(|&img| {
179                    let (w, h) = (img.width(), img.height());
180                    if w == 0 || h == 0 {
181                        return None;
182                    }
183                    Some(image::imageops::resize(
184                        img,
185                        crop_w,
186                        crop_h,
187                        self.resize_filter,
188                    ))
189                })
190                .collect()
191        };
192
193        // Convert to dynamic images and normalize using common helper
194        let dynamic_images = rgb_to_dynamic(resized_rgb);
195        self.normalizer.normalize_batch_to(dynamic_images)
196    }
197
198    /// Runs inference on the preprocessed batch.
199    ///
200    /// # Arguments
201    ///
202    /// * `batch_tensor` - Preprocessed batch tensor
203    ///
204    /// # Returns
205    ///
206    /// Model predictions as a 2D tensor (batch_size x num_classes)
207    pub fn infer(
208        &self,
209        batch_tensor: &ndarray::Array4<f32>,
210    ) -> Result<ndarray::Array2<f32>, OCRError> {
211        let input_name = self.inference.input_name();
212        let inputs = vec![(input_name, TensorInput::Array4(batch_tensor))];
213
214        let outputs = self
215            .inference
216            .infer(&inputs)
217            .map_err(|e| OCRError::Inference {
218                model_name: "PP-LCNet".to_string(),
219                context: format!(
220                    "failed to run inference on batch with shape {:?}",
221                    batch_tensor.shape()
222                ),
223                source: Box::new(e),
224            })?;
225
226        let output = outputs
227            .into_iter()
228            .next()
229            .ok_or_else(|| OCRError::InvalidInput {
230                message: "PP-LCNet: no output returned from inference".to_string(),
231            })?;
232
233        output
234            .1
235            .try_into_array2_f32()
236            .map_err(|e| OCRError::Inference {
237                model_name: "PP-LCNet".to_string(),
238                context: "failed to convert output to 2D array".to_string(),
239                source: Box::new(e),
240            })
241    }
242
243    /// Postprocesses model predictions to class IDs and scores.
244    ///
245    /// # Arguments
246    ///
247    /// * `predictions` - Model predictions (batch_size x num_classes)
248    /// * `config` - Postprocessing configuration
249    ///
250    /// # Returns
251    ///
252    /// PPLCNetModelOutput containing class IDs, scores, and optional label names
253    pub fn postprocess(
254        &self,
255        predictions: &ndarray::Array2<f32>,
256        config: &PPLCNetPostprocessConfig,
257    ) -> Result<PPLCNetModelOutput, OCRError> {
258        let predictions_vec: Vec<Vec<f32>> =
259            predictions.outer_iter().map(|row| row.to_vec()).collect();
260
261        let topk_result = self
262            .topk_processor
263            .process(&predictions_vec, config.topk)
264            .unwrap_or_else(|_| crate::utils::topk::TopkResult {
265                indexes: vec![],
266                scores: vec![],
267                class_names: None,
268            });
269
270        let class_ids = topk_result.indexes;
271        let scores = topk_result.scores;
272
273        // Map class IDs to label names if labels are provided
274        let label_names = if !config.labels.is_empty() {
275            Some(
276                class_ids
277                    .iter()
278                    .map(|ids| {
279                        ids.iter()
280                            .map(|&id| {
281                                config
282                                    .labels
283                                    .get(id)
284                                    .cloned()
285                                    .unwrap_or_else(|| format!("class_{}", id))
286                            })
287                            .collect()
288                    })
289                    .collect(),
290            )
291        } else {
292            topk_result.class_names
293        };
294
295        Ok(PPLCNetModelOutput {
296            class_ids,
297            scores,
298            label_names,
299        })
300    }
301
302    /// Performs complete forward pass: preprocess -> infer -> postprocess.
303    ///
304    /// # Arguments
305    ///
306    /// * `images` - Input images to classify
307    /// * `config` - Postprocessing configuration
308    ///
309    /// # Returns
310    ///
311    /// PPLCNetModelOutput containing classification results
312    pub fn forward(
313        &self,
314        images: Vec<RgbImage>,
315        config: &PPLCNetPostprocessConfig,
316    ) -> Result<PPLCNetModelOutput, OCRError> {
317        let image_refs: Vec<&RgbImage> = images.iter().collect();
318        self.forward_refs(&image_refs, config)
319    }
320
321    /// Performs a complete forward pass from borrowed images.
322    pub fn forward_refs(
323        &self,
324        images: &[&RgbImage],
325        config: &PPLCNetPostprocessConfig,
326    ) -> Result<PPLCNetModelOutput, OCRError> {
327        let batch_tensor = self.preprocess_refs(images)?;
328        let predictions = self.infer(&batch_tensor)?;
329        self.postprocess(&predictions, config)
330    }
331}
332
333/// Builder for PP-LCNet model.
334#[derive(Debug, Default)]
335pub struct PPLCNetModelBuilder {
336    /// Preprocessing configuration
337    preprocess_config: PPLCNetPreprocessConfig,
338    /// ONNX Runtime session configuration
339    ort_config: Option<crate::core::config::OrtSessionConfig>,
340}
341
342impl PPLCNetModelBuilder {
343    /// Creates a new PP-LCNet model builder.
344    pub fn new() -> Self {
345        Self {
346            preprocess_config: PPLCNetPreprocessConfig::default(),
347            ort_config: None,
348        }
349    }
350
351    /// Sets the preprocessing configuration.
352    pub fn preprocess_config(mut self, config: PPLCNetPreprocessConfig) -> Self {
353        self.preprocess_config = config;
354        self
355    }
356
357    /// Sets the input image shape.
358    pub fn input_shape(mut self, shape: (u32, u32)) -> Self {
359        self.preprocess_config.input_shape = shape;
360        self
361    }
362
363    /// Sets the resizing filter.
364    pub fn resize_filter(mut self, filter: FilterType) -> Self {
365        self.preprocess_config.resize_filter = filter;
366        self
367    }
368
369    /// Sets the ONNX Runtime session configuration.
370    pub fn with_ort_config(mut self, config: crate::core::config::OrtSessionConfig) -> Self {
371        self.ort_config = Some(config);
372        self
373    }
374
375    /// Builds the PP-LCNet model.
376    ///
377    /// # Arguments
378    ///
379    /// * `model_source` - Path to the ONNX model file
380    ///
381    /// # Returns
382    ///
383    /// A configured PP-LCNet model instance
384    pub fn build(
385        self,
386        model_source: impl Into<crate::core::ModelSource>,
387    ) -> Result<PPLCNetModel, OCRError> {
388        // Create ONNX inference engine
389        let inference = if self.ort_config.is_some() {
390            use crate::core::config::ModelInferenceConfig;
391            let common_config = ModelInferenceConfig {
392                ort_session: self.ort_config,
393                ..Default::default()
394            };
395            OrtInfer::from_config(&common_config, model_source, None)?
396        } else {
397            OrtInfer::new(model_source, None)?
398        };
399
400        // Create normalizer (ImageNet normalization).
401        //
402        // PP-LCNet classifiers read images as **RGB** by default (no `DecodeImage`
403        // op in the official inference.yml), so we keep RGB order here.
404        let mean = self.preprocess_config.normalize_mean.clone();
405        let std = self.preprocess_config.normalize_std.clone();
406        let normalizer = NormalizeImage::with_color_order(
407            Some(self.preprocess_config.normalize_scale),
408            Some(mean),
409            Some(std),
410            Some(self.preprocess_config.tensor_layout),
411            Some(crate::processors::types::ColorOrder::RGB),
412        )?;
413
414        // Create top-k processor
415        let topk_processor = Topk::new(None);
416
417        Ok(PPLCNetModel::new(
418            inference,
419            normalizer,
420            topk_processor,
421            self.preprocess_config.input_shape,
422            self.preprocess_config.resize_filter,
423            self.preprocess_config.resize_short,
424        ))
425    }
426}
oar_ocr_core/models/classification/pp_lcnet.rs

oar_ocr_core/models/classification/
pp_lcnet.rs
