rusto-rs 0.1.5

RustO! - Pure Rust OCR library based on RapidOCR with PaddleOCR engine
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

use ndarray::Array4;

#[cfg(feature = "use-opencv")]
use opencv::{core, imgproc, prelude::*};

#[cfg(feature = "use-opencv")]
use opencv::core::Mat;

#[cfg(not(feature = "use-opencv"))]
use crate::image_impl::Mat;

use crate::engine::EngineError;

pub struct DetPreProcess {
    pub limit_side_len: i32,
    pub limit_type: String,
    pub mean: [f32; 3],
    pub std: [f32; 3],
}

impl DetPreProcess {
    pub fn new(limit_side_len: i32, limit_type: String, mean: [f32; 3], std: [f32; 3]) -> Self {
        Self {
            limit_side_len,
            limit_type,
            mean,
            std,
        }
    }

    pub fn run(&self, img: &Mat) -> Result<Array4<f32>, EngineError> {
        let resized = self.resize(img)?;
        self.normalize_and_permute(&resized)
    }

    fn resize(&self, img: &Mat) -> Result<Mat, EngineError> {
        let h = img.rows();
        let w = img.cols();

        let ratio = if self.limit_type == "max" {
            let max_side = h.max(w) as f32;
            if max_side > self.limit_side_len as f32 {
                self.limit_side_len as f32 / max_side
            } else {
                1.0
            }
        } else {
            let min_side = h.min(w) as f32;
            if min_side < self.limit_side_len as f32 {
                self.limit_side_len as f32 / min_side
            } else {
                1.0
            }
        };
        
        let mut resize_h = (h as f32 * ratio) as i32;
        let mut resize_w = (w as f32 * ratio) as i32;

        resize_h = ((resize_h as f32 / 32.0).round() * 32.0) as i32;
        resize_w = ((resize_w as f32 / 32.0).round() * 32.0) as i32;

        if resize_h <= 0 || resize_w <= 0 {
            return Err(EngineError::Preprocess("resize_h or resize_w <= 0".to_string()));
        }

        #[cfg(feature = "use-opencv")]
        let dst = {
            let mut d = Mat::default();
            imgproc::resize(
                img,
                &mut d,
                core::Size::new(resize_w, resize_h),
                0.0,
                0.0,
                imgproc::INTER_LINEAR,
            )?;
            d
        };
        
        #[cfg(not(feature = "use-opencv"))]
        let dst = {
            let mut d = Mat::default();
            crate::image_impl::resize(
                img,
                &mut d,
                crate::image_impl::Size::new(resize_w, resize_h),
                crate::image_impl::INTER_LINEAR,
            )?;
            d
        };

        Ok(dst)
    }

    #[cfg(feature = "use-opencv")]
    fn normalize_and_permute(&self, img: &Mat) -> Result<Array4<f32>, EngineError> {
        let size = img.size()?;
        let h = size.height as usize;
        let w = size.width as usize;

        let mut out = Array4::<f32>::zeros((1, 3, h, w));
        let scale = 1.0 / 255.0;

        for y in 0..h {
            for x in 0..w {
                let pix = img.at_2d::<core::Vec3b>(y as i32, x as i32)?;
                let b = pix[0] as f32 * scale;
                let g = pix[1] as f32 * scale;
                let r = pix[2] as f32 * scale;

                out[[0, 0, y, x]] = (b - self.mean[0]) / self.std[0];
                out[[0, 1, y, x]] = (g - self.mean[1]) / self.std[1];
                out[[0, 2, y, x]] = (r - self.mean[2]) / self.std[2];
            }
        }

        Ok(out)
    }
    
    #[cfg(not(feature = "use-opencv"))]
    fn normalize_and_permute(&self, img: &Mat) -> Result<Array4<f32>, EngineError> {
        let size = img.size()?;
        let h = size.height as usize;
        let w = size.width as usize;

        let mut out = ndarray::Array4::<f32>::zeros((1, 3, h, w));
        let scale = 1.0 / 255.0;
        
        // Cache normalization parameters
        let mean_b = self.mean[0];
        let mean_g = self.mean[1];
        let mean_r = self.mean[2];
        let std_b = self.std[0];
        let std_g = self.std[1];
        let std_r = self.std[2];

        for y in 0..h {
            for x in 0..w {
                let pix = img.get_pixel(x as u32, y as u32);
                // image crate loads as RGB, convert to BGR for detection model
                let r = pix[0] as f32 * scale;  // Red channel
                let g = pix[1] as f32 * scale;  // Green channel
                let b = pix[2] as f32 * scale;  // Blue channel

                // Store in BGR order (detection model expects BGR)
                out[[0, 0, y, x]] = (b - mean_b) / std_b;  // Blue
                out[[0, 1, y, x]] = (g - mean_g) / std_g;  // Green  
                out[[0, 2, y, x]] = (r - mean_r) / std_r;  // Red
            }
        }

        Ok(out)
    }
}