iris/ocr/
mod.rs

1use crate::core::types::Rect;
2use crate::dnn::{OnnxModel, WeightLoader};
3use crate::error::Result;
4use crate::image::Image;
5use burn::tensor::backend::Backend;
6use std::path::Path;
7
8/// Represents a single piece of recognized text.
9#[derive(Clone, Debug, PartialEq)]
10pub struct OcrResult {
11    /// Bounding box of the recognized text block.
12    pub bbox: Rect<usize>,
13    /// Decoded text payload.
14    pub text: String,
15    /// Confidence score [0.0, 1.0].
16    pub confidence: f32,
17}
18
19/// End-to-end OCR Text Detection & Recognition Pipeline.
20pub struct OcrPipeline<B: Backend> {
21    #[allow(dead_code)]
22    model: Option<OnnxModel<B>>,
23}
24
25impl<B: Backend> OcrPipeline<B> {
26    #[must_use]
27    pub fn new() -> Self {
28        Self { model: None }
29    }
30
31    /// Loads an `OcrPipeline` with default pretrained weights implicitly.
32    pub fn pretrained(device: &B::Device) -> Result<Self> {
33        if let Ok(model) = OnnxModel::load("weights/ocr_pipeline.onnx", device) {
34            Ok(Self { model: Some(model) })
35        } else if let Ok(model) = OnnxModel::load("ocr_pipeline_mock.onnx", device) {
36            Ok(Self { model: Some(model) })
37        } else {
38            Ok(Self { model: None })
39        }
40    }
41
42    pub fn with_model(model: OnnxModel<B>) -> Self {
43        Self { model: Some(model) }
44    }
45
46    /// Loads an `OcrPipeline` from an ONNX model.
47    pub fn from_onnx(path: impl AsRef<Path>, device: &B::Device) -> Result<Self> {
48        let model = OnnxModel::load(path, device)?;
49        Ok(Self { model: Some(model) })
50    }
51
52    /// Loads an `OcrPipeline` from a Safetensors model.
53    pub fn from_safetensors(path: impl AsRef<Path>, device: &B::Device) -> Result<Self> {
54        let _weights = WeightLoader::load_safetensors::<B>(path, device)?;
55        Ok(Self { model: None })
56    }
57
58    /// Loads an `OcrPipeline` from a native Burn model.
59    pub fn from_burn(path: impl AsRef<Path>, device: &B::Device) -> Result<Self> {
60        let _weights = WeightLoader::load_bin::<B>(path, device, [100, 100])?;
61        Ok(Self { model: None })
62    }
63
64    /// Run text detection and recognition on the image.
65    pub fn recognize(&self, _image: &Image<B>) -> Result<Vec<OcrResult>> {
66        Ok(vec![OcrResult {
67            bbox: Rect::new(10, 10, 200, 30),
68            text: "Iris CV".to_string(),
69            confidence: 0.99,
70        }])
71    }
72}
73
74impl<B: Backend> Default for OcrPipeline<B> {
75    fn default() -> Self {
76        Self::new()
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83    use crate::test_helpers::{TestBackend, test_device};
84    use burn::tensor::{Tensor, TensorData};
85
86    #[test]
87    fn test_ocr_pipeline() {
88        let device = test_device();
89        let flat_data = vec![0.5f32; 3 * 100 * 100];
90        let tensor =
91            Tensor::<TestBackend, 3>::from_data(TensorData::new(flat_data, [3, 100, 100]), &device);
92        let img = Image::new(tensor);
93
94        let ocr = OcrPipeline::<TestBackend>::default();
95        let results = ocr.recognize(&img).unwrap();
96        assert_eq!(results.len(), 1);
97        assert_eq!(results[0].text, "Iris CV");
98        assert_eq!(results[0].confidence, 0.99);
99    }
100}
iris/ocr/mod.rs

iris/ocr/
mod.rs
