This commit is contained in:
Matthew Deville 2026-01-22 01:28:12 +01:00
parent 75dd44f27a
commit 1fae113c5c

View file

@ -63,12 +63,14 @@ fn main() -> Result<(), Box<dyn Error>> {
// Create ONNX Runtime session // Create ONNX Runtime session
let mut session = create_session(&model_path)?; let mut session = create_session(&model_path)?;
let input_name = session.inputs()[0].name().to_string();
// Preprocess image // Preprocess image
let input_tensor = preprocess_image(&img, model_info.input_size.0, model_info.input_size.1)?; let input_tensor = preprocess_image(&img, model_info.input_size.0, model_info.input_size.1)?;
// Run inference // Run inference
println!("Running inference..."); println!("Running inference...");
let outputs = session.run(ort::inputs!["input" => Tensor::from_array(input_tensor)?])?; let outputs = session.run(ort::inputs![input_name => Tensor::from_array(input_tensor)?])?;
// Extract mask output // Extract mask output
let mask_output = &outputs[0]; let mask_output = &outputs[0];