mistralrs_vision/
lib.rs

1//! This crate provides vision utilities for mistral.rs inspired by torchvision.
2//! In particular, it represents transformations on some `Self` type which are applied
3//! sequentially.
4//!
5//! ## Example
6//! ```rust
7//! use candle_core::Device;
8//! use image::{ColorType, DynamicImage};
9//! use mistralrs_vision::{ApplyTransforms, Normalize, ToTensor, Transforms};
10//!
11//! let image = DynamicImage::new(3, 4, ColorType::Rgb8);
12//! let transforms = Transforms {
13//!     input: &ToTensor,
14//!     inner_transforms: &[&Normalize {
15//!         mean: vec![0.5, 0.5, 0.5],
16//!         std: vec![0.5, 0.5, 0.5],
17//!     }],
18//! };
19//! let transformed = image.apply(transforms, &Device::Cpu).unwrap();
20//! assert_eq!(transformed.dims(), &[3, 4, 3]);
21//! ```
22
23use candle_core::{Device, Result, Tensor};
24use image::DynamicImage;
25mod ops;
26mod pad;
27mod transforms;
28pub(crate) mod utils;
29
30pub use ops::{get_resize_image_size, make_pixel_mask, pad};
31pub use pad::{pad_to_max_edge, pad_to_max_image_size};
32pub use transforms::{InterpolateResize, Normalize, Rescale, ToTensor, ToTensorNoNorm};
33
34/// A transform over an image. The input may vary but the output is always a Tensor.
35pub trait ImageTransform {
36    type Input;
37    type Output;
38
39    fn map(&self, x: &Self::Input, device: &Device) -> Result<Self::Output>;
40}
41
42/// Transforms to apply, starting with the `input` and then with each transform in
43/// `inner_transforms` applied sequentially
44pub struct Transforms<'a> {
45    pub input: &'a dyn ImageTransform<Input = DynamicImage, Output = Tensor>,
46    pub inner_transforms: &'a [&'a dyn ImageTransform<Input = Tensor, Output = Tensor>],
47}
48
49/// Transforms, with each of `inner_transforms` applied sequentially
50pub struct TensorTransforms<'a> {
51    pub inner_transforms: &'a [&'a dyn ImageTransform<Input = Tensor, Output = Tensor>],
52}
53
54/// Application of transforms to the Self type.
55pub trait ApplyTransforms<'a> {
56    fn apply(&self, transforms: Transforms<'a>, device: &Device) -> Result<Tensor>;
57}
58
59impl<'a> ApplyTransforms<'a> for DynamicImage {
60    fn apply(&self, transforms: Transforms<'a>, device: &Device) -> Result<Tensor> {
61        let mut res = transforms.input.map(self, device)?;
62        for transform in transforms.inner_transforms {
63            res = transform.map(&res, device)?;
64        }
65        Ok(res)
66    }
67}
68
69/// Application of transforms to the Self type.
70pub trait ApplyTensorTransforms<'a> {
71    fn apply(&self, transforms: TensorTransforms<'a>, device: &Device) -> Result<Tensor>;
72}
73
74impl<'a> ApplyTensorTransforms<'a> for Tensor {
75    fn apply(&self, transforms: TensorTransforms<'a>, device: &Device) -> Result<Tensor> {
76        let mut res = self.clone();
77        for transform in transforms.inner_transforms {
78            res = transform.map(&res, device)?;
79        }
80        Ok(res)
81    }
82}
mistralrs_vision/lib.rs

mistralrs_vision/
lib.rs
