oar_ocr_core/core/inference/
mod.rs

1//! Structures and helpers for ONNX Runtime inference.
2//!
3//! This module provides a unified inference interface that is backend-agnostic
4//! and does not make assumptions about input/output semantics.
5
6use crate::core::errors::OCRError;
7use ort::{session::Session, value::ValueType};
8use std::sync::Mutex;
9
10pub mod session;
11mod tensor_output;
12
13// OrtInfer implementation modules
14#[path = "ort_infer_builders.rs"]
15mod ort_infer_builders;
16#[path = "ort_infer_config.rs"]
17mod ort_infer_config;
18#[path = "ort_infer_execution.rs"]
19mod ort_infer_execution;
20
21pub use ort_infer_builders::ensure_cuda_launch_blocking;
22pub use ort_infer_execution::TensorInput;
23pub use session::load_session;
24pub use tensor_output::TensorOutput;
25
26/// Core ONNX Runtime inference engine with support for pooling and configurable sessions.
27pub struct OrtInfer {
28    pub(self) sessions: Vec<Mutex<Session>>,
29    pub(self) next_idx: std::sync::atomic::AtomicUsize,
30    pub(self) input_name: String,
31    pub(self) model_path: std::path::PathBuf,
32    pub(self) model_name: String,
33}
34
35impl std::fmt::Debug for OrtInfer {
36    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37        f.debug_struct("OrtInfer")
38            .field("sessions", &self.sessions.len())
39            .field("input_name", &self.input_name)
40            .field("model_path", &self.model_path)
41            .field("model_name", &self.model_name)
42            .finish()
43    }
44}
45
46impl OrtInfer {
47    /// Returns the input tensor name.
48    pub fn input_name(&self) -> &str {
49        &self.input_name
50    }
51
52    /// Gets a session from the pool.
53    pub fn get_session(&self, idx: usize) -> Result<std::sync::MutexGuard<'_, Session>, OCRError> {
54        self.sessions[idx % self.sessions.len()]
55            .lock()
56            .map_err(|_| OCRError::ConfigError {
57                message: "Failed to acquire session lock".to_string(),
58            })
59    }
60
61    /// Returns the declared input names from the model.
62    pub fn input_names_from_model(&self) -> Vec<String> {
63        let Some(session_mutex) = self.sessions.first() else {
64            return Vec::new();
65        };
66        let Ok(session_guard) = session_mutex.lock() else {
67            return Vec::new();
68        };
69        session_guard
70            .inputs()
71            .iter()
72            .map(|i| i.name().to_string())
73            .collect()
74    }
75
76    /// Attempts to retrieve the primary input tensor shape from the first session.
77    ///
78    /// Returns a vector of dimensions if available. Dynamic dimensions (e.g., -1) are returned as-is.
79    pub fn primary_input_shape(&self) -> Option<Vec<i64>> {
80        let session_mutex = self.sessions.first()?;
81        let session_guard = session_mutex.lock().ok()?;
82        let input = session_guard.inputs().first()?;
83        match input.dtype() {
84            ValueType::Tensor { shape, .. } => Some(shape.iter().copied().collect()),
85            _ => None,
86        }
87    }
88
89    /// Returns the declared output names and tensor shapes from the first session.
90    ///
91    /// This is intended for model adapters that need to choose among multiple
92    /// ONNX outputs before interpreting tensors semantically.
93    pub fn output_shapes(&self) -> Vec<(String, Vec<i64>)> {
94        let Some(session_mutex) = self.sessions.first() else {
95            return Vec::new();
96        };
97        let Ok(session_guard) = session_mutex.lock() else {
98            return Vec::new();
99        };
100        session_guard
101            .outputs()
102            .iter()
103            .filter_map(|output| match output.dtype() {
104                ValueType::Tensor { shape, .. } => Some((
105                    output.name().to_string(),
106                    shape.iter().copied().collect::<Vec<_>>(),
107                )),
108                _ => None,
109            })
110            .collect()
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117    use crate::core::config::ModelInferenceConfig;
118
119    #[test]
120    fn test_from_config_with_ort_session() {
121        let common = ModelInferenceConfig::new();
122        let result = OrtInfer::from_config(&common, "dummy_path.onnx", None);
123        assert!(result.is_err()); // File doesn't exist
124    }
125}
oar_ocr_core/core/inference/mod.rs

oar_ocr_core/core/inference/
mod.rs
