use ndarray::ArrayD;
use crate::config::Backend;
use crate::error::{OcrError, Result};
#[cfg(feature = "ort")]
mod ort_backend;
#[cfg(feature = "tract")]
mod tract_backend;
pub type Tensor = ArrayD<f32>;
pub trait ModelBackend: Send + Sync {
fn name(&self) -> &str;
fn run(&self, input: Tensor) -> Result<Tensor>;
}
pub fn load_backend(
backend: Backend,
model_bytes: &[u8],
threads: usize,
fixed_input: Option<&[usize]>,
) -> Result<Box<dyn ModelBackend>> {
let _ = (model_bytes, threads, fixed_input);
match backend {
#[cfg(feature = "ort")]
Backend::Ort => Ok(Box::new(ort_backend::OrtBackend::load(model_bytes, threads)?)),
#[cfg(feature = "tract")]
Backend::Tract => Ok(Box::new(tract_backend::TractBackend::load(model_bytes, fixed_input)?)),
other => Err(OcrError::inference(format!(
"backend {other:?} is not compiled in (enable the matching cargo feature)"
))),
}
}