mistralrs_vision/
pad.rs

1use image::{
2    imageops::FilterType, DynamicImage, GenericImage, GenericImageView, ImageBuffer, Rgb, Rgba,
3};
4
5fn resize_image_to_max_edge(img: &DynamicImage, max_edge: u32) -> DynamicImage {
6    // Get the original dimensions of the image
7    let (width, height) = img.dimensions();
8
9    // Calculate the scaling factor
10    let scale = if width > height {
11        max_edge as f32 / width as f32
12    } else {
13        max_edge as f32 / height as f32
14    };
15
16    // New dimensions
17    let new_width = (width as f32 * scale) as u32;
18    let new_height = (height as f32 * scale) as u32;
19
20    // Resize the image
21    img.resize_exact(new_width, new_height, FilterType::Lanczos3)
22}
23
24/// 1) Resize the images to the maximum edge length - preserving aspect ratio
25/// 2) Pad all the images with black padding.
26pub fn pad_to_max_edge(images: &[DynamicImage], max_edge: u32) -> Vec<DynamicImage> {
27    let mut new_images = Vec::new();
28    for image in images {
29        new_images.push(resize_image_to_max_edge(image, max_edge));
30    }
31
32    let mut max_height = 0;
33    let mut max_width = 0;
34    for image in &new_images {
35        let (w, h) = image.dimensions();
36        if w > max_width {
37            max_width = w;
38        }
39        if h > max_height {
40            max_height = h;
41        }
42    }
43
44    for image in &mut new_images {
45        match image {
46            DynamicImage::ImageRgb8(rgb8) => {
47                let mut padded_image = ImageBuffer::from_pixel(max_width, max_height, Rgb([0; 3]));
48
49                padded_image
50                    .copy_from(rgb8, 0, 0)
51                    .expect("Failed to copy image");
52
53                *rgb8 = padded_image;
54            }
55            DynamicImage::ImageRgba8(rgba8) => {
56                let mut padded_image = ImageBuffer::from_pixel(max_width, max_height, Rgba([0; 4]));
57
58                padded_image
59                    .copy_from(rgba8, 0, 0)
60                    .expect("Failed to copy image");
61
62                *rgba8 = padded_image;
63            }
64            _ => panic!("rgb8 or rgba8 are the only supported image types"),
65        }
66    }
67
68    new_images
69}
70
71pub fn pad_to_max_image_size(mut images: Vec<DynamicImage>) -> Vec<DynamicImage> {
72    let mut max_height = 0;
73    let mut max_width = 0;
74    for image in &images {
75        let (w, h) = image.dimensions();
76        if w > max_width {
77            max_width = w;
78        }
79        if h > max_height {
80            max_height = h;
81        }
82    }
83
84    for image in &mut images {
85        match image {
86            DynamicImage::ImageRgb8(rgb8) => {
87                let mut padded_image = ImageBuffer::from_pixel(max_width, max_height, Rgb([0; 3]));
88
89                padded_image
90                    .copy_from(rgb8, 0, 0)
91                    .expect("Failed to copy image");
92
93                *rgb8 = padded_image;
94            }
95            DynamicImage::ImageRgba8(rgba8) => {
96                let mut padded_image = ImageBuffer::from_pixel(max_width, max_height, Rgba([0; 4]));
97
98                padded_image
99                    .copy_from(rgba8, 0, 0)
100                    .expect("Failed to copy image");
101
102                *rgba8 = padded_image;
103            }
104            _ => panic!("rgb8 or rgba8 are the only supported image types"),
105        }
106    }
107
108    images
109}
mistralrs_vision/pad.rs

mistralrs_vision/
pad.rs
