huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-17 11:08:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

dinov2 - read images from disk and compute the class probabilities (#503)

Browse Source 
* Load the image from disk and convert it to a tensor.

* Tweak the function name.
This commit is contained in:
Laurent Mazare
2023-08-18 15:50:33 +01:00
committed by GitHub
parent 95462c6a2e
commit 4f1541526c
4 changed files with 48 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
candle-examples/examples/stable-diffusion/utils.rs
View File

@ -12,25 +12,6 @@ pub fn linspace(start: f64, stop: f64, steps: usize) -> Result<Tensor> {
Tensor::from_vec(vs, steps, &Device::Cpu)
}
/// Saves an image to disk using the image crate, this expects an input with shape
/// (c, width, height).
pub fn save_image<P: AsRef<std::path::Path>>(img: &Tensor, p: P) -> Result<()> {
let p = p.as_ref();
let (channel, width, height) = img.dims3()?;
if channel != 3 {
candle::bail!("save_image expects an input of shape (3, width, height)")
}
let img = img.transpose(0, 1)?.t()?.flatten_all()?;
let pixels = img.to_vec1::<u8>()?;
let image: image::ImageBuffer<image::Rgb<u8>, Vec<u8>> =
match image::ImageBuffer::from_raw(width as u32, height as u32, pixels) {
Some(image) => image,
None => candle::bail!("error saving image {p:?}"),
};
image.save(p).map_err(candle::Error::wrap)?;
Ok(())
}
// Wrap the conv2d op to provide some tracing.
#[derive(Debug)]
pub struct Conv2d {