diff options
Diffstat (limited to 'candle-examples/src/imagenet.rs')
| -rw-r--r-- | candle-examples/src/imagenet.rs | 18 |
1 files changed, 18 insertions, 0 deletions
diff --git a/candle-examples/src/imagenet.rs b/candle-examples/src/imagenet.rs index cefbd71b..781dcd4f 100644 --- a/candle-examples/src/imagenet.rs +++ b/candle-examples/src/imagenet.rs @@ -17,6 +17,24 @@ pub fn load_image224<P: AsRef<std::path::Path>>(p: P) -> Result<Tensor> { .broadcast_div(&std) } +/// Loads an image from disk using the image crate, this returns a tensor with shape +/// (3, 518, 518). imagenet normalization is applied. +/// The model dinov2 reg4 analyzes images with dimensions 3x518x518 (resulting in 37x37 transformer tokens). +pub fn load_image518<P: AsRef<std::path::Path>>(p: P) -> Result<Tensor> { + let img = image::io::Reader::open(p)? + .decode() + .map_err(candle::Error::wrap)? + .resize_to_fill(518, 518, image::imageops::FilterType::Triangle); + let img = img.to_rgb8(); + let data = img.into_raw(); + let data = Tensor::from_vec(data, (518, 518, 3), &Device::Cpu)?.permute((2, 0, 1))?; + let mean = Tensor::new(&[0.485f32, 0.456, 0.406], &Device::Cpu)?.reshape((3, 1, 1))?; + let std = Tensor::new(&[0.229f32, 0.224, 0.225], &Device::Cpu)?.reshape((3, 1, 1))?; + (data.to_dtype(candle::DType::F32)? / 255.)? + .broadcast_sub(&mean)? + .broadcast_div(&std) +} + pub const CLASS_COUNT: i64 = 1000; pub const CLASSES: [&str; 1000] = [ |