rtmo test ok

Author: youngday

examples/rtmo/main.rs

@@ -1,17 +1,51 @@
 use anyhow::Result;
-use usls::{models::RTMO, Annotator, DataLoader, Options, Style, SKELETON_COCO_19};
+use usls::{models::RTMO, Annotator, DataLoader, Config, Style, SKELETON_COCO_19};
+#[derive(argh::FromArgs)]
+/// Example
+struct Args {
+    /// model file
+    #[argh(option)]
+    model: Option<String>,
 
+    /// source
+    #[argh(option, default = "String::from(\"./assets/bus.jpg\")")]
+    source: String,
+    /// dtype
+    #[argh(option, default = "String::from(\"fp16\")")]
+    dtype: String,
+
+    /// device
+    #[argh(option, default = "String::from(\"cpu:0\")")]
+    device: String,
+
+    /// scale: t, s, m, l
+    #[argh(option, default = "String::from(\"t\")")]
+    scale: String,
+}
 fn main() -> Result<()> {
     tracing_subscriber::fmt()
         .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
         .with_timer(tracing_subscriber::fmt::time::ChronoLocal::rfc_3339())
         .init();
-
+    let args: Args = argh::from_env();
+          // build model
+    let config = match args.scale.as_str() {
+        "t" => Config::rtmo_t(),
+        "s" => Config::rtmo_s(),
+        "m" => Config::rtmo_m(),
+        "l" => Config::rtmo_l(),
+        _ => unreachable!(),
+    }
+    .with_model_file(&args.model.unwrap_or_default())
+    .with_model_dtype(args.dtype.parse()?)
+    .with_model_device(args.device.parse()?)
+    .commit()?;
     // build model
-    let mut model = RTMO::new(Options::rtmo_s().commit()?)?;
+    let mut model = RTMO::new(config)?;
 
     // load image
-    let xs = DataLoader::try_read_n(&["./assets/bus.jpg"])?;
+        
+    let xs = DataLoader::try_read_n(&[&args.source])?;
 
     // run
     let ys = model.forward(&xs)?;