问题描述
我正在尝试使用tract_onnx 将使用onnxruntime 的Python 脚本转换为Rust。我试图实现的特定 POC 是来自 ONNX Model Zoo 的 the rothe_vgg.py
script。这个脚本使用了三个模型:
- 超人脸检测 (
version-RFB-320.onnx
) -
vgg_ilsvrc_16_age_imdb_wiki.onnx
和vgg_ilsvrc_16_gender_imdb_wiki.onnx
age and gender models
目前,我只是尝试第一个检测人脸的模型。我可以让the example Python code上班:
face_detector_onnx = "models/version-RFB-320.onnx"
face_detector = ort.InferenceSession(face_detector_onnx)
def faceDetector(orig_image,threshold = 0.7):
image = cv2.cvtColor(orig_image,cv2.COLOR_BGR2RGB)
image = cv2.resize(image,(320,240))
image_mean = np.array([127,127,127])
image = (image - image_mean) / 128
image = np.transpose(image,[2,1])
image = np.expand_dims(image,axis=0)
image = image.astype(np.float32)
input_name = face_detector.get_inputs()[0].name
confidences,Boxes = face_detector.run(None,{input_name: image})
Boxes,labels,probs = predict(orig_image.shape[1],orig_image.shape[0],confidences,Boxes,threshold)
return Boxes,probs
我的tract_onnx 翻译基于onnx-mobilenet-v2 example。我的版本目前看起来像这样:
let model = onnx()
.model_for_path("version-RFB-320.onnx")?
.with_input_fact(
0,InferenceFact::dt_shape(f32::datum_type(),tvec!(1,3,240,320)),)?
.into_optimized()?
.into_runnable()?;
let image = image::open("bruce.jpg").unwrap().to_rgb8();
let resized = image::imageops::resize(&image,320,::image::imageops::FilterType::Triangle);
let image: Tensor = tract_ndarray::Array4::from_shape_fn((1,320),|(_,c,y,x)| {
resized[(x as _,y as _)][c] as f32 / 255.0
}).into();
let result = model.run(tvec!(image))?;
我在将调整大小的图像转换为张量时遇到问题:
thread 'main' panicked at 'Image index (240,0) out of bounds (240,320)'.
这是没有正确维度或每个维度的正确顺序的问题吗?我错过了什么吗?
我知道我还没有实现其他翻译,这是我的下一个问题:如何使用 image_mean
正确规范化、转置和扩展维度?
解决方法
检查项目 https://github.com/recoai/visual-search。它使用 ONNX 在 Rust 中实现了几个深度学习模型。您可以配置任何深度学习模型图像转换或使用 4 个预配置模型。
这是如何配置模型的示例。该项目不支持分割模型,但图像转换管道正在运行。
let model_config = ModelConfig {
model_name: "SqueezeNet".into(),model_url: "https://github.com/onnx/models/raw/master/vision/classification/squeezenet/model/squeezenet1.1-7.onnx".into(),image_transformation: TransformationPipeline {
steps: vec![
ResizeRGBImageAspectRatio { image_size: ImageSize { width: 224,height: 224 },scale: 87.5,filter: FilterType::Nearest }.into(),CenterCrop { crop_size: ImageSize {width: 224,height: 224} }.into(),ToArray {}.into(),Normalization { sub: [0.485,0.456,0.406],div: [0.229,0.224,0.225],zeroone: true }.into(),ToTensor {}.into(),]
},image_size: ImageSize { width: 224,layer_name: Some("squeezenet0_pool3_fwd".to_string()),channels: Channels::CWH
}