use std::path::Path;
use image::{ImageBuffer, Pixel, Rgb};
use ort::{
environment::Environment, execution_providers::CPUExecutionProviderOptions, tensor::OrtOwnedTensor, value::Value, ExecutionProvider,
GraphOptimizationLevel, LoggingLevel, OrtResult, SessionBuilder
};
use test_log::test;
#[test]
fn upsample() -> OrtResult<()> {
const IMAGE_TO_LOAD: &str = "mushroom.png";
let environment = Environment::builder()
.with_name("integration_test")
.with_log_level(LoggingLevel::Warning)
.with_execution_providers([ExecutionProvider::CPU(CPUExecutionProviderOptions { use_arena: true })])
.build()?
.into_arc();
let session_data =
std::fs::read(Path::new(env!("CARGO_MANIFEST_DIR")).join("tests").join("data").join("upsample.onnx")).expect("Could not open model from file");
let session = SessionBuilder::new(&environment)?
.with_optimization_level(GraphOptimizationLevel::Level1)?
.with_intra_threads(1)?
.with_model_from_memory(&session_data)
.expect("Could not read model from memory");
let metadata = session.metadata()?;
assert_eq!(metadata.name()?, "tf2onnx");
assert_eq!(metadata.producer()?, "tf2onnx");
assert_eq!(session.inputs[0].dimensions().collect::<Vec<_>>(), [None, None, None, Some(3)]);
assert_eq!(session.outputs[0].dimensions().collect::<Vec<_>>(), [None, None, None, Some(3)]);
let image_buffer: ImageBuffer<Rgb<u8>, Vec<u8>> = image::open(Path::new(env!("CARGO_MANIFEST_DIR")).join("tests").join("data").join(IMAGE_TO_LOAD))
.unwrap()
.to_rgb8();
let array = ndarray::CowArray::from(
ndarray::Array::from_shape_fn((1, 224, 224, 3), |(_, j, i, c)| {
let pixel = image_buffer.get_pixel(i as u32, j as u32);
let channels = pixel.channels();
(channels[c] as f32) / 255.0
})
.into_dyn()
);
let input_tensor_values = vec![Value::from_array(session.allocator(), &array)?];
let outputs: Vec<Value> = session.run(input_tensor_values)?;
assert_eq!(outputs.len(), 1);
let output: OrtOwnedTensor<'_, f32, ndarray::Dim<ndarray::IxDynImpl>> = outputs[0].try_extract()?;
assert_eq!(output.view().shape(), [1, 448, 448, 3]);
Ok(())
}