File size: 1,766 Bytes
32652fd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
use clap::Parser;
use crate::YOLOTask;
#[derive(Parser, Clone)]
#[command(author, version, about, long_about = None)]
pub struct Args {
/// ONNX model path
#[arg(long, required = true)]
pub model: String,
/// input path
#[arg(long, required = true)]
pub source: String,
/// device id
#[arg(long, default_value_t = 0)]
pub device_id: u32,
/// using TensorRT EP
#[arg(long)]
pub trt: bool,
/// using CUDA EP
#[arg(long)]
pub cuda: bool,
/// input batch size
#[arg(long, default_value_t = 1)]
pub batch: u32,
/// trt input min_batch size
#[arg(long, default_value_t = 1)]
pub batch_min: u32,
/// trt input max_batch size
#[arg(long, default_value_t = 32)]
pub batch_max: u32,
/// using TensorRT --fp16
#[arg(long)]
pub fp16: bool,
/// specify YOLO task
#[arg(long, value_enum)]
pub task: Option<YOLOTask>,
/// num_classes
#[arg(long)]
pub nc: Option<u32>,
/// num_keypoints
#[arg(long)]
pub nk: Option<u32>,
/// num_masks
#[arg(long)]
pub nm: Option<u32>,
/// input image width
#[arg(long)]
pub width: Option<u32>,
/// input image height
#[arg(long)]
pub height: Option<u32>,
/// confidence threshold
#[arg(long, required = false, default_value_t = 0.3)]
pub conf: f32,
/// iou threshold in NMS
#[arg(long, required = false, default_value_t = 0.45)]
pub iou: f32,
/// confidence threshold of keypoint
#[arg(long, required = false, default_value_t = 0.55)]
pub kconf: f32,
/// plot inference result and save
#[arg(long)]
pub plot: bool,
/// check time consumed in each stage
#[arg(long)]
pub profile: bool,
}
|