pipeless-ai 1.11.0

An open-source computer vision framework to build and deploy applications in minutes
use log::warn;

use crate as pipeless;

use super::{
    onnx::{OnnxSession, OnnxSessionParams},
    openvino::{OpenvinoSession, OpenvinoSessionParams}
};

pub trait SessionTrait {
    fn infer(&self, frame: pipeless::data::Frame) -> pipeless::data::Frame;
}

pub enum InferenceSession {
    Onnx(OnnxSession),
    Openvino(OpenvinoSession)
}
impl InferenceSession {
    pub fn infer(&self, frame: pipeless::data::Frame) -> pipeless::data::Frame {
        match self {
            InferenceSession::Onnx(onnx_session) => onnx_session.infer(frame),
            InferenceSession::Openvino(_) => unimplemented!(),
        }
    }
}

pub enum SessionParams {
    Onnx(OnnxSessionParams),
    Openvino(OpenvinoSessionParams)
}
impl SessionParams {
    pub fn from_raw_data(stage_name: &str, runtime: &super::runtime::InferenceRuntime, data: &serde_json::Value) -> Self {
        match runtime {
            super::runtime::InferenceRuntime::Onnx => {
                let execution_provider = data["execution_provider"].as_str().unwrap_or_else(|| {
                    warn!("'execution_provider' not specified, defaulting to 'CPU'");
                    "cpu"
                });
                let execution_mode = data["execution_mode"].as_str();
                let inter_threads = data["inter_threads"].as_i64();
                if inter_threads.is_some() && execution_mode.is_none() {
                    warn!("'execution_mode' must be set to 'Parallel' for 'inter_threads' to take effect");
                }
                let intra_threads = data["intra_threads"].as_i64();
                let custom_op_lib_path = data["custom_op_lib_path"].as_str();
                SessionParams::Onnx(
                    OnnxSessionParams::new(
                        stage_name,
                        execution_provider, execution_mode,
                        inter_threads.map(|t| t as i16),
                        intra_threads.map(|t| t as i16),
                        custom_op_lib_path,
                ))
            },
            super::runtime::InferenceRuntime::Openvino => unimplemented!(),
        }
    }
}