huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-15 18:28:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

Segment-anything fixes: avoid normalizing twice. (#767)

Browse Source 
* Segment-anything fixes: avoid normalizing twice.

* More fixes for the image aspect ratio.
This commit is contained in:
Laurent Mazare
2023-09-07 21:45:16 +01:00
committed by GitHub
parent 7396b8ed1a
commit 79c27fc489
3 changed files with 33 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
candle-examples/examples/segment-anything/main.rs
View File

@ -108,7 +108,8 @@ pub fn main() -> anyhow::Result<()> {
let device = candle_examples::device(args.cpu)?;
let image = candle_examples::imagenet::load_image224(args.image)?.to_device(&device)?;
let image =
candle_examples::load_image(args.image, Some(model_sam::IMAGE_SIZE))?.to_device(&device)?;
println!("loaded image {image:?}");
let model = match args.model {
@ -125,7 +126,7 @@ pub fn main() -> anyhow::Result<()> {
let sam = model_sam::Sam::new(768, 12, 12, &[2, 5, 8, 11], vb)?; // sam_vit_b
let (mask, iou_predictions) = sam.forward(&image, false)?;
println!("mask: {mask:?}");
println!("mask:\n{mask}");
println!("iou_predictions: {iou_predictions:?}");
Ok(())
}

3
candle-examples/examples/segment-anything/model_sam.rs
View File

@ -6,7 +6,7 @@ use crate::model_mask_decoder::MaskDecoder;
use crate::model_prompt_encoder::PromptEncoder;
const PROMPT_EMBED_DIM: usize = 256;
const IMAGE_SIZE: usize = 1024;
pub const IMAGE_SIZE: usize = 1024;
const VIT_PATCH_SIZE: usize = 16;
#[derive(Debug)]
@ -90,6 +90,7 @@ impl Sam {
fn preprocess(&self, img: &Tensor) -> Result<Tensor> {
let (c, h, w) = img.dims3()?;
let img = img
.to_dtype(DType::F32)?
.broadcast_sub(&self.pixel_mean)?
.broadcast_div(&self.pixel_std)?;
if h > IMAGE_SIZE || w > IMAGE_SIZE {