rlx-ocr 0.2.0

OCR engine for RLX — text detection + recognition
Documentation 
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

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! Quick check with real ocrs safetensors (set `OCR_MODEL_DIR`).

use anyhow::Result;
use rlx_ocr::{ImageSource, OcrEngine, OcrEngineParams, input_image, resolve_model_dir};
use rten_tensor::prelude::*;

#[path = "env.rs"]
mod bench_env;

fn model_dir() -> Option<std::path::PathBuf> {
    bench_env::env_var("OCR_MODEL_DIR", "OCRS_MODEL_DIR")
        .map(std::path::PathBuf::from)
        .filter(|p| p.is_dir())
}

#[test]
fn ocr_detection_and_engine_with_real_weights() -> Result<()> {
    let Some(dir) = model_dir() else {
        eprintln!("skip: set OCR_MODEL_DIR to a directory with ocrs safetensors");
        return Ok(());
    };
    let (det, _rec) = match resolve_model_dir(&dir) {
        Ok(p) => p,
        Err(e) => {
            eprintln!("skip: {e}");
            return Ok(());
        }
    };

    let cfg = rlx_ocr::model::DetectionGraphConfig {
        batch: 1,
        height: 64,
        width: 96,
    };
    let mut wm = rlx_ocr::load_safetensors_weights(&det)?;
    rlx_ocr::model::build_detection_graph(&mut wm, cfg)?;

    // Forward pass: enable when runtime NCHW path is verified on this host.
    if std::env::var("OCR_RUN_FORWARD").ok().as_deref() == Some("1") {
        let detector = rlx_ocr::RlxTextDetector::from_safetensors_sized(
            &det,
            Default::default(),
            cfg,
            rlx_runtime::Device::Cpu,
        )?;
        let grey = vec![128u8; 96 * 64];
        let src = ImageSource::from_bytes(&grey, (96, 64))?;
        let input = OcrEngine::new(OcrEngineParams::default())?.prepare_input(src)?;
        let view = input_image(&input);
        let mask_a = detector.detect_text_pixels(view)?;
        let mask_b = detector.detect_text_pixels(view)?;
        assert_eq!(mask_a.shape(), mask_b.shape());
    }
    Ok(())
}