mirror of
https://github.com/mii443/usls.git
synced 2025-08-22 15:45:41 +00:00
Add BEN2 model (#91)
This commit is contained in:
5
examples/ben2/README.md
Normal file
5
examples/ben2/README.md
Normal file
@ -0,0 +1,5 @@
|
|||||||
|
## Quick Start
|
||||||
|
|
||||||
|
```shell
|
||||||
|
cargo run -r --example ben2
|
||||||
|
```
|
48
examples/ben2/main.rs
Normal file
48
examples/ben2/main.rs
Normal file
@ -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(())
|
||||||
|
}
|
9
src/models/ben2/README.md
Normal file
9
src/models/ben2/README.md
Normal file
@ -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)
|
6
src/models/ben2/config.rs
Normal file
6
src/models/ben2/config.rs
Normal file
@ -0,0 +1,6 @@
|
|||||||
|
/// Model configuration for `BEN2`
|
||||||
|
impl crate::Options {
|
||||||
|
pub fn ben2_base() -> Self {
|
||||||
|
Self::rmbg().with_model_file("ben2-base.onnx")
|
||||||
|
}
|
||||||
|
}
|
1
src/models/ben2/mod.rs
Normal file
1
src/models/ben2/mod.rs
Normal file
@ -0,0 +1 @@
|
|||||||
|
mod config;
|
@ -1,4 +1,5 @@
|
|||||||
mod beit;
|
mod beit;
|
||||||
|
mod ben2;
|
||||||
mod blip;
|
mod blip;
|
||||||
mod clip;
|
mod clip;
|
||||||
mod convnext;
|
mod convnext;
|
||||||
|
@ -1,4 +1,4 @@
|
|||||||
/// Model configuration for `RMBG-2.0`
|
/// Model configuration for `RMBG`
|
||||||
impl crate::Options {
|
impl crate::Options {
|
||||||
pub fn rmbg() -> Self {
|
pub fn rmbg() -> Self {
|
||||||
Self::default()
|
Self::default()
|
||||||
|
Reference in New Issue
Block a user