mistralrs_vision/
ops.rs

1use candle_core::{Result, Tensor};
2
3/// Pad an image of shape (c, h, w) to (c, max_h, max_w) by padding with zeros on the right and bottom.
4pub fn pad(image: &Tensor, max_h: usize, max_w: usize) -> Result<Tensor> {
5    let (c, h, w) = image.dims3()?;
6    let new_image = Tensor::zeros((c, max_h, max_w), image.dtype(), image.device())?;
7    new_image.slice_assign(&[&(..c), &(..h), &(..w)], image)
8}
9
10/// Generate pixel mask of shape (c, max_h, max_w). 1 indicates valid pixel, 0 indicates padding.
11///
12/// The input tensor is of shape (c, max_h, max_w) and the output mask is the same shape and
13/// represents where pixels are. The mask shape is in the top left corner is passed as `h` and `w`.
14pub fn make_pixel_mask(image: &Tensor, h: usize, w: usize) -> Result<Tensor> {
15    let (_c, max_h, max_w) = image.dims3()?;
16    let mask = Tensor::ones((h, w), image.dtype(), image.device())?;
17    let zeros = Tensor::zeros((max_h, max_w), image.dtype(), image.device())?;
18    // TODO(EricLBuehler): https://github.com/huggingface/candle/pull/2223 will make this nicer
19    zeros.slice_assign(&[&(..h), &(..w)], &mask)
20}
21
22/// Given the image sizes (h, w) and the minimum and maximum lengths, calculate the image dimensions
23/// which will preserve aspect ration while respecing the minimum and maximum lengths.
24pub fn get_resize_image_size(
25    (h, w): (usize, usize),
26    (min_len, max_len): (usize, usize),
27) -> (usize, usize) {
28    let aspect_ratio = w as f64 / h as f64;
29
30    let (new_h, new_w) = if w >= h && w > max_len {
31        ((max_len as f64 / aspect_ratio) as usize, max_len)
32    } else if h > w && h > max_len {
33        (max_len, (max_len as f64 * aspect_ratio) as usize)
34    } else {
35        (h, w)
36    };
37    (new_h.max(min_len), new_w.max(min_len))
38}
mistralrs_vision/ops.rs

mistralrs_vision/
ops.rs
