Skip to content

Load the models from memory? #37

@ftnfurina

Description

@ftnfurina

Hi! This is a really great project 💖, Bad english, excuse me 😅.

The project uses ort. I wonder if there are plans to support loading models from memory in the future?

Session::builder()?.commit_from_memory(include_bytes!("../onnx/PP-OCRv4_mobile_rec_infer.onnx"))

Below is the temporary solution I have used.

use std::{fmt, io::Write, path::Path};

use anyhow::{Result, anyhow};
use image::RgbImage;
use oar_ocr::{
    core::StandardPredictor,
    predictor::{
        TextDetPredictor,
        TextDetPredictorBuilder,
        TextRecPredictor,
        TextRecPredictorBuilder,
    },
    processors::{BoundingBox, LimitType},
};
use tempfile::NamedTempFile;
use tracing::debug;

pub struct Ocr {
    text_recognizer: TextRecPredictor,
    text_detector:   TextDetPredictor,
}

impl Ocr {
    pub fn new() -> Result<Self> {
        let mut rec_file = NamedTempFile::new()?;
        rec_file.write_all(include_bytes!("../onnx/PP-OCRv4_mobile_rec_infer.onnx"))?;

        let mut det_file = NamedTempFile::new()?;
        det_file.write_all(include_bytes!("../onnx/PP-OCRv4_mobile_det_infer.onnx"))?;

        let mut doc_dict = include_str!("../onnx/PP-OCRv4_doc_dict.txt")
            .lines()
            .map(|line| line.to_string())
            .collect::<Vec<String>>();
        doc_dict.push("".to_string());

        Ok(Self {
            text_recognizer: Self::build_text_recognizer(&rec_file.into_temp_path(), doc_dict)?,
            text_detector:   Self::build_text_detector(&det_file.into_temp_path())?,
        })
    }

    fn build_text_recognizer(model_path: &Path, doc_dict: Vec<String>) -> Result<TextRecPredictor> {
        TextRecPredictorBuilder::new()
            .batch_size(1)
            .model_name("crnn")
            .session_pool_size(1)
            .enable_logging(false)
            .model_input_shape([3, 48, 320])
            .character_dict(doc_dict)
            .build(model_path)
            .map_err(|e| anyhow!("Build text recognizer fail: {}", e))
    }

    fn build_text_detector(model_path: &Path) -> Result<TextDetPredictor> {
        TextDetPredictorBuilder::new()
            .batch_size(1)
            .model_name("crnn")
            .session_pool_size(1)
            .enable_logging(false)
            .limit_side_len(736)
            .limit_type(LimitType::Max)
            .max_side_limit(4000)
            .build(model_path)
            .map_err(|e| anyhow!("Build text detector fail: {}", e))
    }
    // ......
}

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions