From 3898e500debd632b520054ebfa42f8333323a20e Mon Sep 17 00:00:00 2001 From: Laurent Mazare Date: Fri, 8 Sep 2023 05:53:08 +0100 Subject: Generate a mask image + the scaled input image. (#769) * Also round-trip the original image. * Make it possible to use a safetensors input. --- candle-examples/examples/segment-anything/model_sam.rs | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) (limited to 'candle-examples/examples/segment-anything/model_sam.rs') diff --git a/candle-examples/examples/segment-anything/model_sam.rs b/candle-examples/examples/segment-anything/model_sam.rs index acba7ef4..237163a3 100644 --- a/candle-examples/examples/segment-anything/model_sam.rs +++ b/candle-examples/examples/segment-anything/model_sam.rs @@ -87,7 +87,15 @@ impl Sam { Ok((low_res_mask, iou_predictions)) } - fn preprocess(&self, img: &Tensor) -> Result { + pub fn unpreprocess(&self, img: &Tensor) -> Result { + let img = img + .broadcast_mul(&self.pixel_std)? + .broadcast_add(&self.pixel_mean)?; + img.maximum(&img.zeros_like()?)? + .minimum(&(img.ones_like()? * 255.)?) + } + + pub fn preprocess(&self, img: &Tensor) -> Result { let (c, h, w) = img.dims3()?; let img = img .to_dtype(DType::F32)? -- cgit v1.2.3