76 lines
1.8 KiB
Rust
76 lines
1.8 KiB
Rust
use std::path::PathBuf;
|
|
|
|
use mnn::ForwardType;
|
|
#[derive(Debug, clap::Parser)]
|
|
pub struct Cli {
|
|
#[clap(subcommand)]
|
|
pub cmd: SubCommand,
|
|
}
|
|
|
|
#[derive(Debug, clap::Subcommand)]
|
|
pub enum SubCommand {
|
|
#[clap(name = "detect")]
|
|
Detect(Detect),
|
|
#[clap(name = "list")]
|
|
List(List),
|
|
#[clap(name = "completions")]
|
|
Completions { shell: clap_complete::Shell },
|
|
}
|
|
|
|
#[derive(Debug, clap::ValueEnum, Clone, Copy)]
|
|
pub enum Models {
|
|
RetinaFace,
|
|
Yolo,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub enum Executor {
|
|
Mnn(mnn::ForwardType),
|
|
Ort(Vec<detector::ort_ep::ExecutionProvider>),
|
|
}
|
|
|
|
#[derive(Debug, clap::Args)]
|
|
pub struct Detect {
|
|
#[clap(short, long)]
|
|
pub model: Option<PathBuf>,
|
|
#[clap(short = 'M', long, default_value = "retina-face")]
|
|
pub model_type: Models,
|
|
#[clap(short, long)]
|
|
pub output: Option<PathBuf>,
|
|
#[clap(
|
|
short = 'p',
|
|
long,
|
|
default_value = "cpu",
|
|
group = "execution_provider",
|
|
required_unless_present = "mnn_forward_type"
|
|
)]
|
|
pub ort_execution_provider: Vec<detector::ort_ep::ExecutionProvider>,
|
|
#[clap(
|
|
short = 'f',
|
|
long,
|
|
group = "execution_provider",
|
|
required_unless_present = "ort_execution_provider"
|
|
)]
|
|
pub mnn_forward_type: Option<mnn::ForwardType>,
|
|
#[clap(short, long, default_value_t = 0.8)]
|
|
pub threshold: f32,
|
|
#[clap(short, long, default_value_t = 0.3)]
|
|
pub nms_threshold: f32,
|
|
pub image: PathBuf,
|
|
}
|
|
|
|
#[derive(Debug, clap::Args)]
|
|
pub struct List {}
|
|
|
|
impl Cli {
|
|
pub fn completions(shell: clap_complete::Shell) {
|
|
let mut command = <Cli as clap::CommandFactory>::command();
|
|
clap_complete::generate(
|
|
shell,
|
|
&mut command,
|
|
env!("CARGO_BIN_NAME"),
|
|
&mut std::io::stdout(),
|
|
);
|
|
}
|
|
}
|