334 lines
7.7 KiB
Rust
334 lines
7.7 KiB
Rust
use crate::{auto_load, Device, MinOptMax};
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct Options {
|
|
pub onnx_path: String,
|
|
pub device: Device,
|
|
pub profile: bool,
|
|
pub num_dry_run: usize,
|
|
pub i00: Option<MinOptMax>, // 1st input, axis 0, batch usually
|
|
pub i01: Option<MinOptMax>, // 1st input, axis 1
|
|
pub i02: Option<MinOptMax>,
|
|
pub i03: Option<MinOptMax>,
|
|
pub i04: Option<MinOptMax>,
|
|
pub i05: Option<MinOptMax>,
|
|
pub i10: Option<MinOptMax>, // 2nd input, axis 0
|
|
pub i11: Option<MinOptMax>, // 2nd input, axis 1
|
|
pub i12: Option<MinOptMax>,
|
|
pub i13: Option<MinOptMax>,
|
|
pub i14: Option<MinOptMax>,
|
|
pub i15: Option<MinOptMax>,
|
|
pub i20: Option<MinOptMax>,
|
|
pub i21: Option<MinOptMax>,
|
|
pub i22: Option<MinOptMax>,
|
|
pub i23: Option<MinOptMax>,
|
|
pub i24: Option<MinOptMax>,
|
|
pub i25: Option<MinOptMax>,
|
|
pub i30: Option<MinOptMax>,
|
|
pub i31: Option<MinOptMax>,
|
|
pub i32_: Option<MinOptMax>,
|
|
pub i33: Option<MinOptMax>,
|
|
pub i34: Option<MinOptMax>,
|
|
pub i35: Option<MinOptMax>,
|
|
|
|
// trt related
|
|
pub trt_engine_cache_enable: bool,
|
|
pub trt_int8_enable: bool,
|
|
pub trt_fp16_enable: bool,
|
|
|
|
// options for Vision and Language models
|
|
pub nc: Option<usize>,
|
|
pub nk: Option<usize>,
|
|
pub nm: Option<usize>,
|
|
pub confs: Vec<f32>,
|
|
pub kconfs: Vec<f32>,
|
|
pub iou: f32,
|
|
pub apply_nms: bool,
|
|
pub tokenizer: Option<String>,
|
|
pub vocab: Option<String>,
|
|
pub names: Option<Vec<String>>, // class names
|
|
pub anchors_first: bool, // otuput format: [bs, anchors/na, pos+nc+nm]
|
|
pub min_width: Option<f32>,
|
|
pub min_height: Option<f32>,
|
|
pub unclip_ratio: f32, // DB
|
|
}
|
|
|
|
impl Default for Options {
|
|
fn default() -> Self {
|
|
Self {
|
|
onnx_path: String::new(),
|
|
device: Device::Cuda(0),
|
|
profile: false,
|
|
num_dry_run: 3,
|
|
i00: None,
|
|
i01: None,
|
|
i02: None,
|
|
i03: None,
|
|
i04: None,
|
|
i05: None,
|
|
i10: None,
|
|
i11: None,
|
|
i12: None,
|
|
i13: None,
|
|
i14: None,
|
|
i15: None,
|
|
i20: None,
|
|
i21: None,
|
|
i22: None,
|
|
i23: None,
|
|
i24: None,
|
|
i25: None,
|
|
i30: None,
|
|
i31: None,
|
|
i32_: None,
|
|
i33: None,
|
|
i34: None,
|
|
i35: None,
|
|
trt_engine_cache_enable: true,
|
|
trt_int8_enable: false,
|
|
trt_fp16_enable: false,
|
|
nc: None,
|
|
nk: None,
|
|
nm: None,
|
|
confs: vec![0.4f32],
|
|
kconfs: vec![0.5f32],
|
|
iou: 0.45f32,
|
|
apply_nms: true,
|
|
tokenizer: None,
|
|
vocab: None,
|
|
names: None,
|
|
anchors_first: false,
|
|
min_width: None,
|
|
min_height: None,
|
|
unclip_ratio: 1.5,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Options {
|
|
pub fn with_model(mut self, onnx_path: &str) -> Self {
|
|
self.onnx_path = auto_load(onnx_path).unwrap();
|
|
self
|
|
}
|
|
|
|
pub fn with_dry_run(mut self, n: usize) -> Self {
|
|
self.num_dry_run = n;
|
|
self
|
|
}
|
|
|
|
pub fn with_cuda(mut self, id: usize) -> Self {
|
|
self.device = Device::Cuda(id);
|
|
self
|
|
}
|
|
|
|
pub fn with_trt(mut self, id: usize) -> Self {
|
|
self.device = Device::Trt(id);
|
|
self
|
|
}
|
|
|
|
pub fn with_cpu(mut self) -> Self {
|
|
self.device = Device::Cpu(0);
|
|
self
|
|
}
|
|
|
|
pub fn with_coreml(mut self, id: usize) -> Self {
|
|
self.device = Device::CoreML(id);
|
|
self
|
|
}
|
|
|
|
pub fn with_fp16(mut self, x: bool) -> Self {
|
|
self.trt_fp16_enable = x;
|
|
self
|
|
}
|
|
|
|
pub fn with_profile(mut self, profile: bool) -> Self {
|
|
self.profile = profile;
|
|
self
|
|
}
|
|
|
|
pub fn with_names(mut self, names: &[&str]) -> Self {
|
|
self.names = Some(names.iter().map(|x| x.to_string()).collect::<Vec<String>>());
|
|
self
|
|
}
|
|
|
|
pub fn with_vocab(mut self, vocab: &str) -> Self {
|
|
self.vocab = Some(auto_load(vocab).unwrap());
|
|
self
|
|
}
|
|
|
|
pub fn with_tokenizer(mut self, tokenizer: &str) -> Self {
|
|
self.tokenizer = Some(auto_load(tokenizer).unwrap());
|
|
self
|
|
}
|
|
|
|
pub fn with_unclip_ratio(mut self, x: f32) -> Self {
|
|
self.unclip_ratio = x;
|
|
self
|
|
}
|
|
|
|
pub fn with_min_width(mut self, x: f32) -> Self {
|
|
self.min_width = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_min_height(mut self, x: f32) -> Self {
|
|
self.min_height = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_anchors_first(mut self) -> Self {
|
|
self.anchors_first = true;
|
|
self
|
|
}
|
|
|
|
pub fn with_nms(mut self, apply_nms: bool) -> Self {
|
|
self.apply_nms = apply_nms;
|
|
self
|
|
}
|
|
|
|
pub fn with_nc(mut self, nc: usize) -> Self {
|
|
self.nc = Some(nc);
|
|
self
|
|
}
|
|
|
|
pub fn with_nk(mut self, nk: usize) -> Self {
|
|
self.nk = Some(nk);
|
|
self
|
|
}
|
|
|
|
pub fn with_iou(mut self, x: f32) -> Self {
|
|
self.iou = x;
|
|
self
|
|
}
|
|
|
|
pub fn with_confs(mut self, confs: &[f32]) -> Self {
|
|
self.confs = confs.to_vec();
|
|
self
|
|
}
|
|
|
|
pub fn with_kconfs(mut self, kconfs: &[f32]) -> Self {
|
|
self.kconfs = kconfs.to_vec();
|
|
self
|
|
}
|
|
|
|
pub fn with_i00(mut self, x: MinOptMax) -> Self {
|
|
self.i00 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i01(mut self, x: MinOptMax) -> Self {
|
|
self.i01 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i02(mut self, x: MinOptMax) -> Self {
|
|
self.i02 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i03(mut self, x: MinOptMax) -> Self {
|
|
self.i03 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i04(mut self, x: MinOptMax) -> Self {
|
|
self.i04 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i05(mut self, x: MinOptMax) -> Self {
|
|
self.i05 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i10(mut self, x: MinOptMax) -> Self {
|
|
self.i10 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i11(mut self, x: MinOptMax) -> Self {
|
|
self.i11 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i12(mut self, x: MinOptMax) -> Self {
|
|
self.i12 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i13(mut self, x: MinOptMax) -> Self {
|
|
self.i13 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i14(mut self, x: MinOptMax) -> Self {
|
|
self.i14 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i15(mut self, x: MinOptMax) -> Self {
|
|
self.i15 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i20(mut self, x: MinOptMax) -> Self {
|
|
self.i20 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i21(mut self, x: MinOptMax) -> Self {
|
|
self.i21 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i22(mut self, x: MinOptMax) -> Self {
|
|
self.i22 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i23(mut self, x: MinOptMax) -> Self {
|
|
self.i23 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i24(mut self, x: MinOptMax) -> Self {
|
|
self.i24 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i25(mut self, x: MinOptMax) -> Self {
|
|
self.i25 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i30(mut self, x: MinOptMax) -> Self {
|
|
self.i30 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i31(mut self, x: MinOptMax) -> Self {
|
|
self.i31 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i32_(mut self, x: MinOptMax) -> Self {
|
|
self.i32_ = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i33(mut self, x: MinOptMax) -> Self {
|
|
self.i33 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i34(mut self, x: MinOptMax) -> Self {
|
|
self.i34 = Some(x);
|
|
self
|
|
}
|
|
|
|
pub fn with_i35(mut self, x: MinOptMax) -> Self {
|
|
self.i35 = Some(x);
|
|
self
|
|
}
|
|
}
|