machine_vision_formats/
iter.rs

1//! Types to facilitate iterating over images
2
3use crate::{pixel_format, ImageMutStride, ImageStride, PixelFormat};
4
5/// An image whose rows can be iterated over.
6// In a semver-breaking change, we could eliminate this trait and make its
7// method part of ImageStride.
8pub trait HasRowChunksExact<F>: ImageStride<F> {
9    fn rowchunks_exact(&self) -> RowChunksExact<'_>;
10}
11
12impl<S, F> HasRowChunksExact<F> for S
13where
14    S: ImageStride<F>,
15    F: PixelFormat,
16{
17    fn rowchunks_exact(&self) -> RowChunksExact<'_> {
18        let fmt = pixel_format::pixfmt::<F>().unwrap();
19        let valid_stride = fmt.bits_per_pixel() as usize * self.width() as usize / 8;
20
21        let stride = self.stride();
22        let height = self.height() as usize;
23        let buf = self.buffer_ref().data;
24        let max_len = buf.len().min(stride * height);
25        let buf = &buf[..max_len];
26
27        RowChunksExact {
28            buf,
29            stride,
30            valid_stride,
31        }
32    }
33}
34
35pub struct RowChunksExact<'a> {
36    buf: &'a [u8],
37    stride: usize,
38    valid_stride: usize,
39}
40
41impl std::fmt::Debug for RowChunksExact<'_> {
42    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
43        f.debug_struct("RowChunksExact")
44            .field("stride", &self.stride)
45            .field("valid_stride", &self.valid_stride)
46            .finish_non_exhaustive()
47    }
48}
49
50impl<'a> Iterator for RowChunksExact<'a> {
51    type Item = &'a [u8];
52
53    fn next(&mut self) -> Option<Self::Item> {
54        if self.buf.len() >= self.valid_stride {
55            let mut data: &'a [u8] = &[];
56            std::mem::swap(&mut self.buf, &mut data);
57            if data.len() > self.stride {
58                let (first, rest) = data.split_at(self.stride);
59                self.buf = rest;
60                Some(&first[..self.valid_stride])
61            } else {
62                Some(&data[..self.valid_stride])
63            }
64        } else {
65            None
66        }
67    }
68}
69
70/// An image whose mutable rows can be iterated over.
71// In a semver-breaking change, we could eliminate this trait and make its
72// method part of ImageMutStride.
73pub trait HasRowChunksExactMut<F>: ImageMutStride<F> {
74    fn rowchunks_exact_mut(&mut self) -> RowChunksExactMut<'_>;
75}
76impl<S, F> HasRowChunksExactMut<F> for S
77where
78    S: ImageMutStride<F>,
79    F: PixelFormat,
80{
81    fn rowchunks_exact_mut(&mut self) -> RowChunksExactMut<'_> {
82        let fmt = pixel_format::pixfmt::<F>().unwrap();
83        let valid_stride = fmt.bits_per_pixel() as usize * self.width() as usize / 8;
84
85        let stride = self.stride();
86        let height = self.height() as usize;
87        let buf = self.buffer_mut_ref().data;
88        let max_len = buf.len().min(stride * height);
89        let buf = &mut buf[..max_len];
90        RowChunksExactMut {
91            buf,
92            stride,
93            valid_stride,
94        }
95    }
96}
97
98pub struct RowChunksExactMut<'a> {
99    buf: &'a mut [u8],
100    stride: usize,
101    valid_stride: usize,
102}
103
104impl std::fmt::Debug for RowChunksExactMut<'_> {
105    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
106        f.debug_struct("RowChunksExactMut")
107            .field("stride", &self.stride)
108            .field("valid_stride", &self.valid_stride)
109            .finish_non_exhaustive()
110    }
111}
112
113impl<'a> Iterator for RowChunksExactMut<'a> {
114    type Item = &'a mut [u8];
115
116    fn next(&mut self) -> Option<Self::Item> {
117        if self.buf.len() >= self.valid_stride {
118            let mut data: &'a mut [u8] = &mut [];
119            std::mem::swap(&mut self.buf, &mut data);
120            if data.len() > self.stride {
121                let (first, rest) = data.split_at_mut(self.stride);
122                self.buf = rest;
123                Some(&mut first[..self.valid_stride])
124            } else {
125                Some(&mut data[..self.valid_stride])
126            }
127        } else {
128            None
129        }
130    }
131}
132
133#[cfg(test)]
134mod test {
135    use crate::{
136        iter::{HasRowChunksExact, HasRowChunksExactMut},
137        pixel_format::Mono8,
138        ImageBuffer, ImageBufferMutRef, ImageBufferRef, ImageData, ImageMutData, Stride,
139    };
140
141    struct RoiIm<'a> {
142        width: u32,
143        height: u32,
144        stride: usize,
145        buf: &'a [u8],
146    }
147
148    impl Stride for RoiIm<'_> {
149        fn stride(&self) -> usize {
150            self.stride
151        }
152    }
153
154    impl ImageData<Mono8> for RoiIm<'_> {
155        fn width(&self) -> u32 {
156            self.width
157        }
158        fn height(&self) -> u32 {
159            self.height
160        }
161        fn buffer_ref(&self) -> ImageBufferRef<'_, Mono8> {
162            ImageBufferRef {
163                data: self.buf,
164                pixel_format: std::marker::PhantomData,
165            }
166        }
167        fn buffer(self) -> ImageBuffer<Mono8> {
168            // copy the data
169            self.buffer_ref().to_buffer()
170        }
171    }
172
173    struct RoiImMut<'a> {
174        width: u32,
175        height: u32,
176        stride: usize,
177        buf: &'a mut [u8],
178    }
179
180    impl Stride for RoiImMut<'_> {
181        fn stride(&self) -> usize {
182            self.stride
183        }
184    }
185
186    impl ImageData<Mono8> for RoiImMut<'_> {
187        fn width(&self) -> u32 {
188            self.width
189        }
190        fn height(&self) -> u32 {
191            self.height
192        }
193        fn buffer_ref(&self) -> ImageBufferRef<'_, Mono8> {
194            ImageBufferRef {
195                data: self.buf,
196                pixel_format: std::marker::PhantomData,
197            }
198        }
199        fn buffer(self) -> ImageBuffer<Mono8> {
200            // copy the data
201            self.buffer_ref().to_buffer()
202        }
203    }
204
205    impl ImageMutData<Mono8> for RoiImMut<'_> {
206        fn buffer_mut_ref(&mut self) -> ImageBufferMutRef<'_, Mono8> {
207            ImageBufferMutRef {
208                data: self.buf,
209                pixel_format: std::marker::PhantomData,
210            }
211        }
212    }
213
214    #[test]
215    fn test_roi_at_start() {
216        const STRIDE: usize = 10;
217        const ORIG_W: usize = 10;
218        const ORIG_H: usize = 10;
219        let mut image_data = [0u8; STRIDE * ORIG_H];
220
221        // fill with useful pattern
222        for row in 0..ORIG_H {
223            for col in 0..ORIG_W {
224                image_data[row * STRIDE + col] = (row * 10_usize + col) as u8;
225            }
226        }
227
228        // generate an ROI
229        let width = 2;
230        let height = 2;
231        let (row, col) = (2, 2);
232
233        // create image of this ROI
234        let im = RoiIm {
235            width,
236            height,
237            stride: STRIDE,
238            buf: &image_data[(row * STRIDE + col)..],
239        };
240
241        let mut rowchunk_iter = im.rowchunks_exact();
242        assert_eq!(rowchunk_iter.next(), Some(&[22, 23][..]));
243        assert_eq!(rowchunk_iter.next(), Some(&[32, 33][..]));
244        assert_eq!(rowchunk_iter.next(), None);
245    }
246
247    #[test]
248    fn test_roi_at_end() {
249        const STRIDE: usize = 10;
250        const ORIG_W: usize = 10;
251        const ORIG_H: usize = 10;
252        let mut image_data = [0u8; STRIDE * ORIG_H];
253
254        // fill with useful pattern
255        for row in 0..ORIG_H {
256            for col in 0..ORIG_W {
257                image_data[row * STRIDE + col] = (row * 10_usize + col) as u8;
258            }
259        }
260
261        // generate an ROI
262        let width = 3;
263        let height = 4;
264        let (row, col) = (6, 7);
265
266        // create image of this ROI
267        let im = RoiIm {
268            width,
269            height,
270            stride: STRIDE,
271            buf: &image_data[(row * STRIDE + col)..],
272        };
273
274        let mut rowchunk_iter = im.rowchunks_exact();
275        assert_eq!(rowchunk_iter.next(), Some(&[67, 68, 69][..]));
276        assert_eq!(rowchunk_iter.next(), Some(&[77, 78, 79][..]));
277        assert_eq!(rowchunk_iter.next(), Some(&[87, 88, 89][..]));
278        assert_eq!(rowchunk_iter.next(), Some(&[97, 98, 99][..]));
279        assert_eq!(rowchunk_iter.next(), None);
280    }
281
282    #[test]
283    fn test_mut_roi_at_start() {
284        const STRIDE: usize = 10;
285        const ORIG_W: usize = 10;
286        const ORIG_H: usize = 10;
287        let mut image_data = [0u8; STRIDE * ORIG_H];
288
289        // fill with useful pattern
290        for row in 0..ORIG_H {
291            for col in 0..ORIG_W {
292                image_data[row * STRIDE + col] = (row * 10_usize + col) as u8;
293            }
294        }
295
296        // generate an ROI
297        let width = 2;
298        let height = 2;
299        let (row, col) = (2, 2);
300
301        {
302            // create image of this ROI
303            let mut im = RoiImMut {
304                width,
305                height,
306                stride: STRIDE,
307                buf: &mut image_data[(row * STRIDE + col)..],
308            };
309
310            let mut rowchunk_iter = im.rowchunks_exact_mut();
311            let mut row2 = rowchunk_iter.next();
312            assert_eq!(row2, Some(&mut [22, 23][..]));
313            row2.as_mut().unwrap()[0] += 100;
314            row2.as_mut().unwrap()[1] += 100;
315            let mut row3 = rowchunk_iter.next();
316            assert_eq!(row3, Some(&mut [32, 33][..]));
317            row3.as_mut().unwrap()[0] += 100;
318            row3.as_mut().unwrap()[1] += 100;
319            assert_eq!(rowchunk_iter.next(), None);
320        }
321
322        // create image of this ROI
323        let im = RoiIm {
324            width,
325            height,
326            stride: STRIDE,
327            buf: &image_data[(row * STRIDE + col)..],
328        };
329
330        let mut rowchunk_iter = im.rowchunks_exact();
331        assert_eq!(rowchunk_iter.next(), Some(&[122, 123][..]));
332        assert_eq!(rowchunk_iter.next(), Some(&[132, 133][..]));
333        assert_eq!(rowchunk_iter.next(), None);
334    }
335
336    #[test]
337    fn test_mut_roi_at_end() {
338        const STRIDE: usize = 10;
339        const ORIG_W: usize = 10;
340        const ORIG_H: usize = 10;
341        let mut image_data = [0u8; STRIDE * ORIG_H];
342
343        // fill with useful pattern
344        for row in 0..ORIG_H {
345            for col in 0..ORIG_W {
346                image_data[row * STRIDE + col] = (row * 10_usize + col) as u8;
347            }
348        }
349
350        // generate an ROI
351        let width = 3;
352        let height = 4;
353        let (row, col) = (6, 7);
354
355        {
356            // create image of this ROI
357            let mut im = RoiImMut {
358                width,
359                height,
360                stride: STRIDE,
361                buf: &mut image_data[(row * STRIDE + col)..],
362            };
363
364            let mut rowchunk_iter = im.rowchunks_exact_mut();
365            for row_num in row..(row + height as usize) {
366                let mut this_row = rowchunk_iter.next();
367                assert_eq!(
368                    this_row,
369                    Some(
370                        &mut [
371                            row_num as u8 * 10 + col as u8,
372                            row_num as u8 * 10 + col as u8 + 1,
373                            row_num as u8 * 10 + col as u8 + 2
374                        ][..]
375                    )
376                );
377                for col in 0..width as usize {
378                    this_row.as_mut().unwrap()[col] += 100;
379                }
380            }
381            assert_eq!(rowchunk_iter.next(), None);
382        }
383
384        // create image of this ROI
385        let im = RoiIm {
386            width,
387            height,
388            stride: STRIDE,
389            buf: &image_data[(row * STRIDE + col)..],
390        };
391
392        let mut rowchunk_iter = im.rowchunks_exact();
393        assert_eq!(rowchunk_iter.next(), Some(&[167, 168, 169][..]));
394        assert_eq!(rowchunk_iter.next(), Some(&[177, 178, 179][..]));
395        assert_eq!(rowchunk_iter.next(), Some(&[187, 188, 189][..]));
396        assert_eq!(rowchunk_iter.next(), Some(&[197, 198, 199][..]));
397        assert_eq!(rowchunk_iter.next(), None);
398    }
399}
machine_vision_formats/iter.rs

machine_vision_formats/
iter.rs
