avenger_wgpu/marks/
image.rs

1use crate::error::AvengerWgpuError;
2use etagere::Size;
3use image::DynamicImage;
4use wgpu::Extent3d;
5
6pub struct ImageAtlasBuilder {
7    extent: Extent3d,
8    next_image: image::RgbaImage,
9    images: Vec<DynamicImage>,
10    initialized: bool,
11    allocator: etagere::AtlasAllocator,
12}
13
14#[derive(Copy, Clone)]
15pub struct ImageAtlasCoords {
16    pub x0: f32,
17    pub y0: f32,
18    pub x1: f32,
19    pub y1: f32,
20}
21
22impl Default for ImageAtlasBuilder {
23    fn default() -> Self {
24        Self::new()
25    }
26}
27
28impl ImageAtlasBuilder {
29    pub fn new() -> Self {
30        Self {
31            extent: Extent3d {
32                width: 1,
33                height: 1,
34                depth_or_array_layers: 1,
35            },
36            next_image: image::RgbaImage::new(1, 1),
37            images: vec![],
38            initialized: false,
39            allocator: etagere::AtlasAllocator::new(etagere::Size::new(1, 1)),
40        }
41    }
42
43    pub fn register_image(
44        &mut self,
45        img: &image::RgbaImage,
46    ) -> Result<(usize, ImageAtlasCoords), AvengerWgpuError> {
47        if !self.initialized {
48            let limits = wgpu::Limits::downlevel_webgl2_defaults();
49
50            // Update extent
51            self.extent = Extent3d {
52                width: limits.max_texture_dimension_1d,
53                height: limits.max_texture_dimension_2d,
54                depth_or_array_layers: 1,
55            };
56
57            // Create backing image
58            self.next_image = image::RgbaImage::new(self.extent.width, self.extent.height);
59
60            // Create allocator
61            self.allocator = etagere::AtlasAllocator::new(etagere::Size::new(
62                self.extent.width as i32,
63                self.extent.height as i32,
64            ));
65
66            // Set initialized
67            self.initialized = true;
68        }
69
70        // Attempt to allocate into the current image
71        let allocation = match self
72            .allocator
73            .allocate(Size::new(img.width() as i32, img.height() as i32))
74        {
75            Some(allocation) => allocation,
76            None => {
77                // Allocation failed, create new image
78                let full_image = std::mem::take(&mut self.next_image);
79                self.next_image = image::RgbaImage::new(self.extent.width, self.extent.height);
80                self.images
81                    .push(image::DynamicImage::ImageRgba8(full_image));
82                self.allocator = etagere::AtlasAllocator::new(etagere::Size::new(
83                    self.extent.width as i32,
84                    self.extent.height as i32,
85                ));
86
87                // Try allocating again with new allocator
88                match self
89                    .allocator
90                    .allocate(Size::new(img.width() as i32, img.height() as i32))
91                {
92                    Some(allocation) => allocation,
93                    None => {
94                        if img.width() > self.extent.width || img.height() > self.extent.height {
95                            return Err(AvengerWgpuError::ImageAllocationError(format!(
96                                "Image dimensions ({}, {}) exceed the maximum size of ({}, {})",
97                                img.width(),
98                                img.height(),
99                                self.extent.width,
100                                self.extent.height
101                            )));
102                        } else {
103                            return Err(AvengerWgpuError::ImageAllocationError(
104                                "Unknown error".to_string(),
105                            ));
106                        }
107                    }
108                }
109            }
110        };
111
112        // Write image to allocated portion of final texture image
113        let p0 = allocation.rectangle.min;
114        let p1 = allocation.rectangle.max;
115
116        let x0 = p0.x;
117        let x1 = p1.x.min(x0 + img.width() as i32);
118        let y0 = p0.y;
119        let y1 = p1.y.min(y0 + img.height() as i32);
120
121        for (src_x, dest_x) in (x0..x1).enumerate() {
122            for (src_y, dest_y) in (y0..y1).enumerate() {
123                self.next_image.put_pixel(
124                    dest_x as u32,
125                    dest_y as u32,
126                    *img.get_pixel(src_x as u32, src_y as u32),
127                );
128            }
129        }
130
131        // Compute texture coordinates
132        let coords = ImageAtlasCoords {
133            x0: x0 as f32 / self.extent.width as f32,
134            x1: x1 as f32 / self.extent.width as f32,
135            y0: y0 as f32 / self.extent.height as f32,
136            y1: y1 as f32 / self.extent.height as f32,
137        };
138
139        // Compute image atlas index
140        let atlas_index = self.images.len();
141
142        Ok((atlas_index, coords))
143    }
144
145    pub fn build(&self) -> (Extent3d, Vec<DynamicImage>) {
146        let mut images = self.images.clone();
147        images.push(image::DynamicImage::ImageRgba8(self.next_image.clone()));
148        (self.extent, images)
149    }
150}