diff --git a/examples/ben2/README.md b/examples/ben2/README.md new file mode 100644 index 0000000..7947565 --- /dev/null +++ b/examples/ben2/README.md @@ -0,0 +1,5 @@ +## Quick Start + +```shell +cargo run -r --example ben2 +``` diff --git a/examples/ben2/main.rs b/examples/ben2/main.rs new file mode 100644 index 0000000..9096318 --- /dev/null +++ b/examples/ben2/main.rs @@ -0,0 +1,48 @@ +use usls::{models::RMBG, Annotator, DataLoader, Options}; + +#[derive(argh::FromArgs)] +/// Example +struct Args { + /// dtype + #[argh(option, default = "String::from(\"auto\")")] + dtype: String, + + /// device + #[argh(option, default = "String::from(\"cpu:0\")")] + device: String, +} + +fn main() -> anyhow::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 options = Options::ben2_base() + .with_model_dtype(args.dtype.as_str().try_into()?) + .with_model_device(args.device.as_str().try_into()?) + .commit()?; + let mut model = RMBG::new(options)?; + + // load image + let xs = DataLoader::try_read_n(&["./assets/cat.png"])?; + + // run + let ys = model.forward(&xs)?; + + // annotate + let annotator = Annotator::default(); + for (x, y) in xs.iter().zip(ys.iter()) { + annotator.annotate(x, y)?.save(format!( + "{}.jpg", + usls::Dir::Current + .base_dir_with_subs(&["runs", model.spec()])? + .join(usls::timestamp(None)) + .display(), + ))?; + } + + Ok(()) +} diff --git a/src/models/ben2/README.md b/src/models/ben2/README.md new file mode 100644 index 0000000..2adf389 --- /dev/null +++ b/src/models/ben2/README.md @@ -0,0 +1,9 @@ +# BEN2: Background Erase Network + +## Official Repository + +The official repository can be found on: [HuggingFace](https://huggingface.co/PramaLLC/BEN2) + +## Example + +Refer to the [example](../../../examples/ben2) diff --git a/src/models/ben2/config.rs b/src/models/ben2/config.rs new file mode 100644 index 0000000..b6b5632 --- /dev/null +++ b/src/models/ben2/config.rs @@ -0,0 +1,6 @@ +/// Model configuration for `BEN2` +impl crate::Options { + pub fn ben2_base() -> Self { + Self::rmbg().with_model_file("ben2-base.onnx") + } +} diff --git a/src/models/ben2/mod.rs b/src/models/ben2/mod.rs new file mode 100644 index 0000000..1bf79df --- /dev/null +++ b/src/models/ben2/mod.rs @@ -0,0 +1 @@ +mod config; diff --git a/src/models/mod.rs b/src/models/mod.rs index 13303dc..b0768cf 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,4 +1,5 @@ mod beit; +mod ben2; mod blip; mod clip; mod convnext; diff --git a/src/models/rmbg/config.rs b/src/models/rmbg/config.rs index e44ce8f..65ae2bb 100644 --- a/src/models/rmbg/config.rs +++ b/src/models/rmbg/config.rs @@ -1,4 +1,4 @@ -/// Model configuration for `RMBG-2.0` +/// Model configuration for `RMBG` impl crate::Options { pub fn rmbg() -> Self { Self::default()