ppocr_rs/
angle_net.rs

1use crate::{base_net::BaseNet, ocr_error::OcrError, ocr_result::Angle, ocr_utils::OcrUtils};
2
3use ort::{
4    inputs,
5    session::{Session, SessionOutputs},
6    value::Tensor,
7};
8
9const MEAN_VALUES: [f32; 3] = [127.5, 127.5, 127.5];
10const NORM_VALUES: [f32; 3] = [1.0 / 127.5, 1.0 / 127.5, 1.0 / 127.5];
11const ANGLE_DST_WIDTH: u32 = 192;
12const ANGLE_DST_HEIGHT: u32 = 48;
13const ANGLE_COLS: usize = 2;
14
15#[derive(Debug)]
16pub struct AngleNet {
17    session: Option<Session>,
18    input_names: Vec<String>,
19}
20
21impl BaseNet for AngleNet {
22    fn new() -> Self {
23        Self {
24            session: None,
25            input_names: Vec::new(),
26        }
27    }
28
29    fn set_input_names(&mut self, input_names: Vec<String>) {
30        self.input_names = input_names;
31    }
32
33    fn set_session(&mut self, session: Option<Session>) {
34        self.session = session;
35    }
36}
37
38impl AngleNet {
39    pub fn get_angles(
40        &mut self,
41        part_imgs: &[image::RgbImage],
42        do_angle: bool,
43        most_angle: bool,
44    ) -> Result<Vec<Angle>, OcrError> {
45        let mut angles = Vec::new();
46
47        if do_angle {
48            for img in part_imgs {
49                let angle = self.get_angle(img)?;
50                angles.push(angle);
51            }
52        } else {
53            angles.extend(part_imgs.iter().map(|_| Angle::default()));
54        }
55
56        if do_angle && most_angle {
57            let sum: i32 = angles.iter().map(|x| x.index).sum();
58            let half_percent = angles.len() as f32 / 2.0;
59            let most_angle_index = if (sum as f32) < half_percent { 0 } else { 1 };
60
61            for angle in angles.iter_mut() {
62                angle.index = most_angle_index;
63            }
64        }
65
66        Ok(angles)
67    }
68
69    fn get_angle(&mut self, img_src: &image::RgbImage) -> Result<Angle, OcrError> {
70        let Some(session) = &mut self.session else {
71            return Err(OcrError::SessionNotInitialized);
72        };
73
74        let angle_img = image::imageops::resize(
75            img_src,
76            ANGLE_DST_WIDTH,
77            ANGLE_DST_HEIGHT,
78            image::imageops::FilterType::Triangle,
79        );
80
81        let input_tensors =
82            OcrUtils::substract_mean_normalize(&angle_img, &MEAN_VALUES, &NORM_VALUES);
83
84        let input_tensors = Tensor::from_array(input_tensors)?;
85
86        // [downport rc.11→rc.9] inputs! macro ritorna Result in rc.9
87        let outputs = session.run(inputs![self.input_names[0].clone() => input_tensors]?)?;
88
89        let angle = Self::score_to_angle(&outputs, ANGLE_COLS)?;
90
91        Ok(angle)
92    }
93
94    fn score_to_angle(
95        output_tensor: &SessionOutputs,
96        angle_cols: usize,
97    ) -> Result<Angle, OcrError> {
98        let (_, red_data) = output_tensor.iter().next().unwrap();
99
100        // [downport rc.11→rc.9]
101        let src_data: Vec<f32> = crate::compat::tensor_to_vec_f32(&red_data)?;
102
103        let mut angle = Angle::default();
104        let mut max_value = f32::MIN;
105        let mut angle_index = 0;
106
107        for (i, value) in src_data.iter().take(angle_cols).enumerate() {
108            if i == 0 || value > &max_value {
109                max_value = *value;
110                angle_index = i as i32;
111            }
112        }
113
114        angle.index = angle_index;
115        angle.score = max_value;
116        Ok(angle)
117    }
118}
ppocr_rs/angle_net.rs

ppocr_rs/
angle_net.rs
