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
44#[derive(Clone, Copy)]
45pub struct Transforms<'a> {
46    pub input: &'a dyn ImageTransform<Input = DynamicImage, Output = Tensor>,
47    pub inner_transforms: &'a [&'a dyn ImageTransform<Input = Tensor, Output = Tensor>],
48}
49
50/// Transforms, with each of `inner_transforms` applied sequentially
51#[derive(Clone, Copy)]
52pub struct TensorTransforms<'a> {
53    pub inner_transforms: &'a [&'a dyn ImageTransform<Input = Tensor, Output = Tensor>],
54}
55
56/// Application of transforms to the Self type.
57pub trait ApplyTransforms<'a> {
58    fn apply(&self, transforms: Transforms<'a>, device: &Device) -> Result<Tensor>;
59}
60
61impl<'a> ApplyTransforms<'a> for DynamicImage {
62    fn apply(&self, transforms: Transforms<'a>, device: &Device) -> Result<Tensor> {
63        let mut res = transforms.input.map(self, device)?;
64        for transform in transforms.inner_transforms {
65            res = transform.map(&res, device)?;
66        }
67        Ok(res)
68    }
69}
70
71/// Application of transforms to the Self type.
72pub trait ApplyTensorTransforms<'a> {
73    fn apply(&self, transforms: TensorTransforms<'a>, device: &Device) -> Result<Tensor>;
74}
75
76impl<'a> ApplyTensorTransforms<'a> for Tensor {
77    fn apply(&self, transforms: TensorTransforms<'a>, device: &Device) -> Result<Tensor> {
78        let mut res = self.clone();
79        for transform in transforms.inner_transforms {
80            res = transform.map(&res, device)?;
81        }
82        Ok(res)
83    }
84}
mistralrs_vision/lib.rs

mistralrs_vision/
lib.rs
