rusto-rs 0.1.6

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

use std::time::Instant;

use ndarray::{Array4, ArrayD};

use crate::engine::{EngineError, MnnSession};
use crate::postprocess::{DBPostProcess, TextDetOutput};
use crate::preprocess::DetPreProcess;
use crate::types::DetConfig;

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

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

pub struct TextDetector {
    pub cfg: DetConfig,
    pub session: MnnSession,
    pub postprocess: DBPostProcess,
}

impl TextDetector {
    pub fn new(cfg: DetConfig) -> Result<Self, EngineError> {
        let session = MnnSession::from_det_config(&cfg)?;
        let postprocess = DBPostProcess::new(
            cfg.thresh,
            cfg.box_thresh,
            cfg.max_candidates,
            cfg.unclip_ratio,
            cfg.use_dilation,
        );
        Ok(Self {
            cfg,
            session,
            postprocess,
        })
    }

    pub fn run(&mut self, img: &Mat) -> Result<TextDetOutput, EngineError> {
        let start = Instant::now();

        let ori_h = img.rows();
        let ori_w = img.cols();
        let max_wh = ori_h.max(ori_w);

        let limit_side_len = if self.cfg.limit_type == "min" {
            self.cfg.limit_side_len
        } else if max_wh < 960 {
            960
        } else if max_wh < 1500 {
            1500
        } else {
            2000
        };

        let pre = DetPreProcess::new(
            limit_side_len,
            self.cfg.limit_type.clone(),
            self.cfg.mean,
            self.cfg.std,
        );
        let input = pre.run(img)?;
        let input_dyn: ArrayD<f32> = input.into_dyn();
        let preds_dyn = self.session.run(input_dyn)?;
        let preds: Array4<f32> = preds_dyn
            .into_dimensionality()
            .map_err(|_| EngineError::InvalidInputShape)?;
        let (mut boxes, scores) = self.postprocess.process(&preds, ori_h, ori_w)?;
        if boxes.is_empty() {
            return Ok(TextDetOutput::empty());
        }

        self.sorted_boxes(&mut boxes);
        let elapse = start.elapsed().as_secs_f64();

        Ok(TextDetOutput {
            img: None,
            boxes: Some(boxes),
            scores: Some(scores),
            elapse,
        })
    }

    #[cfg(feature = "use-opencv")]
    fn sorted_boxes(&self, dt_boxes: &mut Vec<[opencv::core::Point2f; 4]>) {
        dt_boxes.sort_by(|a, b| {
            let ay = a[0].y as i32;
            let by = b[0].y as i32;
            if ay != by {
                ay.cmp(&by)
            } else {
                let ax = a[0].x as i32;
                let bx = b[0].x as i32;
                ax.cmp(&bx)
            }
        });
    }

    #[cfg(not(feature = "use-opencv"))]
    fn sorted_boxes(&self, dt_boxes: &mut Vec<[crate::image_impl::Point2f; 4]>) {
        dt_boxes.sort_by(|a, b| {
            let ay = a[0].y as i32;
            let by = b[0].y as i32;
            if ay != by {
                ay.cmp(&by)
            } else {
                let ax = a[0].x as i32;
                let bx = b[0].x as i32;
                ax.cmp(&bx)
            }
        });

    }
}