Add SAM2 and ONNX (#28)

This commit is contained in:
Jamjamjon
2024-08-01 17:26:06 +08:00
committed by GitHub
parent 451aa8cc7b
commit 46a4456a38
2 changed files with 49 additions and 3 deletions

View File

@@ -38,6 +38,20 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
.with_model("sam-vit-b-decoder-u8.onnx")?;
(options_encoder, options_decoder, "SAM")
}
SamKind::Sam2 => {
let options_encoder = Options::default()
// .with_model("sam2-hiera-tiny-encoder.onnx")?;
// .with_model("sam2-hiera-small-encoder.onnx")?;
.with_model("sam2-hiera-base-plus-encoder.onnx")?;
let options_decoder = Options::default()
.with_i31((1, 1, 1).into())
.with_i41((1, 1, 1).into())
.with_sam_kind(SamKind::Sam2)
// .with_model("sam2-hiera-tiny-decoder.onnx")?;
// .with_model("sam2-hiera-small-decoder.onnx")?;
.with_model("sam2-hiera-base-plus-decoder.onnx")?;
(options_encoder, options_decoder, "SAM2")
}
SamKind::MobileSam => {
let options_encoder = Options::default().with_model("mobile-sam-vit-t-encoder.onnx")?;