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            height,
29            buf,
30            stride,
31            valid_stride,
32        }
33    }
34}
35
36pub struct RowChunksExact<'a> {
37    height: usize,
38    buf: &'a [u8],
39    stride: usize,
40    valid_stride: usize,
41}
42
43impl std::fmt::Debug for RowChunksExact<'_> {
44    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
45        f.debug_struct("RowChunksExact")
46            .field("stride", &self.stride)
47            .field("valid_stride", &self.valid_stride)
48            .finish_non_exhaustive()
49    }
50}
51
52impl<'a> Iterator for RowChunksExact<'a> {
53    type Item = &'a [u8];
54
55    fn next(&mut self) -> Option<Self::Item> {
56        if self.height == 0 {
57            return None;
58        }
59        debug_assert!(self.buf.len() >= self.valid_stride);
60        let mut data: &'a [u8] = &[];
61        std::mem::swap(&mut self.buf, &mut data);
62        self.height -= 1;
63        if data.len() > self.stride {
64            let (first, rest) = data.split_at(self.stride);
65            self.buf = rest;
66            Some(&first[..self.valid_stride])
67        } else {
68            Some(&data[..self.valid_stride])
69        }
70    }
71}
72
73impl<'a> DoubleEndedIterator for RowChunksExact<'a> {
74    fn next_back(&mut self) -> Option<Self::Item> {
75        if self.height == 0 {
76            return None;
77        }
78        debug_assert!(self.buf.len() >= self.valid_stride);
79        // self.buf begins with first row. We want to return the last row, which
80        // may be partial.
81        let last_row_start = (self.height - 1) * self.stride;
82        let mut data: &[u8] = &[];
83        std::mem::swap(&mut self.buf, &mut data);
84        let (first, rest) = data.split_at(last_row_start);
85        self.buf = first;
86        self.height -= 1;
87        Some(&rest[..self.valid_stride])
88    }
89}
90
91/// An image whose mutable rows can be iterated over.
92// In a semver-breaking change, we could eliminate this trait and make its
93// method part of ImageMutStride.
94pub trait HasRowChunksExactMut<F>: ImageMutStride<F> {
95    fn rowchunks_exact_mut(&mut self) -> RowChunksExactMut<'_>;
96}
97impl<S, F> HasRowChunksExactMut<F> for S
98where
99    S: ImageMutStride<F>,
100    F: PixelFormat,
101{
102    fn rowchunks_exact_mut(&mut self) -> RowChunksExactMut<'_> {
103        let fmt = pixel_format::pixfmt::<F>().unwrap();
104        let valid_stride = fmt.bits_per_pixel() as usize * self.width() as usize / 8;
105
106        let stride = self.stride();
107        let height = self.height() as usize;
108        let buf = self.buffer_mut_ref().data;
109        let max_len = buf.len().min(stride * height);
110        let buf = &mut buf[..max_len];
111        RowChunksExactMut {
112            height,
113            buf,
114            stride,
115            valid_stride,
116        }
117    }
118}
119
120pub struct RowChunksExactMut<'a> {
121    height: usize,
122    buf: &'a mut [u8],
123    stride: usize,
124    valid_stride: usize,
125}
126
127impl std::fmt::Debug for RowChunksExactMut<'_> {
128    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
129        f.debug_struct("RowChunksExactMut")
130            .field("stride", &self.stride)
131            .field("valid_stride", &self.valid_stride)
132            .finish_non_exhaustive()
133    }
134}
135
136impl<'a> Iterator for RowChunksExactMut<'a> {
137    type Item = &'a mut [u8];
138
139    fn next(&mut self) -> Option<Self::Item> {
140        if self.height == 0 {
141            return None;
142        }
143        debug_assert!(self.buf.len() >= self.valid_stride);
144        let mut data: &'a mut [u8] = &mut [];
145        std::mem::swap(&mut self.buf, &mut data);
146        self.height -= 1;
147        if data.len() > self.stride {
148            let (first, rest) = data.split_at_mut(self.stride);
149            self.buf = rest;
150            Some(&mut first[..self.valid_stride])
151        } else {
152            Some(&mut data[..self.valid_stride])
153        }
154    }
155}
156
157impl<'a> DoubleEndedIterator for RowChunksExactMut<'a> {
158    fn next_back(&mut self) -> Option<Self::Item> {
159        if self.height == 0 {
160            return None;
161        }
162        debug_assert!(self.buf.len() >= self.valid_stride);
163        // self.buf begins with first row. We want to return the last row, which
164        // may be partial.
165        let last_row_start = (self.height - 1) * self.stride;
166        let mut data: &'a mut [u8] = &mut [];
167        std::mem::swap(&mut self.buf, &mut data);
168        let (first, rest) = data.split_at_mut(last_row_start);
169        self.buf = first;
170        self.height -= 1;
171        Some(&mut rest[..self.valid_stride])
172    }
173}
174
175#[cfg(test)]
176mod test {
177    use crate::{
178        iter::{HasRowChunksExact, HasRowChunksExactMut},
179        pixel_format::Mono8,
180        ImageBuffer, ImageBufferMutRef, ImageBufferRef, ImageData, ImageMutData, Stride,
181    };
182
183    struct RoiIm<'a> {
184        width: u32,
185        height: u32,
186        stride: usize,
187        buf: &'a [u8],
188    }
189
190    impl Stride for RoiIm<'_> {
191        fn stride(&self) -> usize {
192            self.stride
193        }
194    }
195
196    impl ImageData<Mono8> for RoiIm<'_> {
197        fn width(&self) -> u32 {
198            self.width
199        }
200        fn height(&self) -> u32 {
201            self.height
202        }
203        fn buffer_ref(&self) -> ImageBufferRef<'_, Mono8> {
204            ImageBufferRef {
205                data: self.buf,
206                pixel_format: std::marker::PhantomData,
207            }
208        }
209        fn buffer(self) -> ImageBuffer<Mono8> {
210            // copy the data
211            self.buffer_ref().to_buffer()
212        }
213    }
214
215    struct RoiImMut<'a> {
216        width: u32,
217        height: u32,
218        stride: usize,
219        buf: &'a mut [u8],
220    }
221
222    impl Stride for RoiImMut<'_> {
223        fn stride(&self) -> usize {
224            self.stride
225        }
226    }
227
228    impl ImageData<Mono8> for RoiImMut<'_> {
229        fn width(&self) -> u32 {
230            self.width
231        }
232        fn height(&self) -> u32 {
233            self.height
234        }
235        fn buffer_ref(&self) -> ImageBufferRef<'_, Mono8> {
236            ImageBufferRef {
237                data: self.buf,
238                pixel_format: std::marker::PhantomData,
239            }
240        }
241        fn buffer(self) -> ImageBuffer<Mono8> {
242            // copy the data
243            self.buffer_ref().to_buffer()
244        }
245    }
246
247    impl ImageMutData<Mono8> for RoiImMut<'_> {
248        fn buffer_mut_ref(&mut self) -> ImageBufferMutRef<'_, Mono8> {
249            ImageBufferMutRef {
250                data: self.buf,
251                pixel_format: std::marker::PhantomData,
252            }
253        }
254    }
255
256    #[test]
257    fn test_roi_at_start() {
258        const STRIDE: usize = 10;
259        const ORIG_W: usize = 10;
260        const ORIG_H: usize = 10;
261        let mut image_data = [0u8; STRIDE * ORIG_H];
262
263        // fill with useful pattern
264        for row in 0..ORIG_H {
265            for col in 0..ORIG_W {
266                image_data[row * STRIDE + col] = (row * 10_usize + col) as u8;
267            }
268        }
269
270        // generate an ROI
271        let width = 2;
272        let height = 2;
273        let (row, col) = (2, 2);
274
275        // create image of this ROI
276        let im = RoiIm {
277            width,
278            height,
279            stride: STRIDE,
280            buf: &image_data[(row * STRIDE + col)..],
281        };
282
283        let mut rowchunk_iter = im.rowchunks_exact();
284        assert_eq!(rowchunk_iter.next(), Some(&[22, 23][..]));
285        assert_eq!(rowchunk_iter.next(), Some(&[32, 33][..]));
286        assert_eq!(rowchunk_iter.next(), None);
287    }
288
289    #[test]
290    fn test_roi_at_end() {
291        const STRIDE: usize = 10;
292        const ORIG_W: usize = 10;
293        const ORIG_H: usize = 10;
294        let mut image_data = [0u8; STRIDE * ORIG_H];
295
296        // fill with useful pattern
297        for row in 0..ORIG_H {
298            for col in 0..ORIG_W {
299                image_data[row * STRIDE + col] = (row * 10_usize + col) as u8;
300            }
301        }
302
303        // generate an ROI
304        let width = 3;
305        let height = 4;
306        let (row, col) = (6, 7);
307
308        // create image of this ROI
309        let im = RoiIm {
310            width,
311            height,
312            stride: STRIDE,
313            buf: &image_data[(row * STRIDE + col)..],
314        };
315
316        let mut rowchunk_iter = im.rowchunks_exact();
317        assert_eq!(rowchunk_iter.next(), Some(&[67, 68, 69][..]));
318        assert_eq!(rowchunk_iter.next(), Some(&[77, 78, 79][..]));
319        assert_eq!(rowchunk_iter.next(), Some(&[87, 88, 89][..]));
320        assert_eq!(rowchunk_iter.next(), Some(&[97, 98, 99][..]));
321        assert_eq!(rowchunk_iter.next(), None);
322    }
323
324    #[test]
325    fn test_mut_roi_at_start() {
326        const STRIDE: usize = 10;
327        const ORIG_W: usize = 10;
328        const ORIG_H: usize = 10;
329        let mut image_data = [0u8; STRIDE * ORIG_H];
330
331        // fill with useful pattern
332        for row in 0..ORIG_H {
333            for col in 0..ORIG_W {
334                image_data[row * STRIDE + col] = (row * 10_usize + col) as u8;
335            }
336        }
337
338        // generate an ROI
339        let width = 2;
340        let height = 2;
341        let (row, col) = (2, 2);
342
343        {
344            // create mutable image of this ROI
345            let mut im = RoiImMut {
346                width,
347                height,
348                stride: STRIDE,
349                buf: &mut image_data[(row * STRIDE + col)..],
350            };
351
352            let mut rowchunk_iter = im.rowchunks_exact_mut();
353            let mut row2 = rowchunk_iter.next();
354            assert_eq!(row2, Some(&mut [22, 23][..]));
355            row2.as_mut().unwrap()[0] += 100;
356            row2.as_mut().unwrap()[1] += 100;
357            let mut row3 = rowchunk_iter.next();
358            assert_eq!(row3, Some(&mut [32, 33][..]));
359            row3.as_mut().unwrap()[0] += 100;
360            row3.as_mut().unwrap()[1] += 100;
361            assert_eq!(rowchunk_iter.next(), None);
362        }
363
364        // create image of this ROI
365        let im = RoiIm {
366            width,
367            height,
368            stride: STRIDE,
369            buf: &image_data[(row * STRIDE + col)..],
370        };
371
372        let mut rowchunk_iter = im.rowchunks_exact();
373        assert_eq!(rowchunk_iter.next(), Some(&[122, 123][..]));
374        assert_eq!(rowchunk_iter.next(), Some(&[132, 133][..]));
375        assert_eq!(rowchunk_iter.next(), None);
376    }
377
378    #[test]
379    fn test_mut_roi_at_end() {
380        const STRIDE: usize = 10;
381        const ORIG_W: usize = 10;
382        const ORIG_H: usize = 10;
383        let mut image_data = [0u8; STRIDE * ORIG_H];
384
385        // fill with useful pattern
386        for row in 0..ORIG_H {
387            for col in 0..ORIG_W {
388                image_data[row * STRIDE + col] = (row * 10_usize + col) as u8;
389            }
390        }
391
392        // generate an ROI
393        let width = 3;
394        let height = 4;
395        let (row, col) = (6, 7);
396
397        {
398            // create mutable image of this ROI
399            let mut im = RoiImMut {
400                width,
401                height,
402                stride: STRIDE,
403                buf: &mut image_data[(row * STRIDE + col)..],
404            };
405
406            let mut rowchunk_iter = im.rowchunks_exact_mut();
407            for row_num in row..(row + height as usize) {
408                let mut this_row = rowchunk_iter.next();
409                assert_eq!(
410                    this_row,
411                    Some(
412                        &mut [
413                            row_num as u8 * 10 + col as u8,
414                            row_num as u8 * 10 + col as u8 + 1,
415                            row_num as u8 * 10 + col as u8 + 2
416                        ][..]
417                    )
418                );
419                for col in 0..width as usize {
420                    this_row.as_mut().unwrap()[col] += 100;
421                }
422            }
423            assert_eq!(rowchunk_iter.next(), None);
424        }
425
426        // create image of this ROI
427        let im = RoiIm {
428            width,
429            height,
430            stride: STRIDE,
431            buf: &image_data[(row * STRIDE + col)..],
432        };
433
434        let mut rowchunk_iter = im.rowchunks_exact();
435        assert_eq!(rowchunk_iter.next(), Some(&[167, 168, 169][..]));
436        assert_eq!(rowchunk_iter.next(), Some(&[177, 178, 179][..]));
437        assert_eq!(rowchunk_iter.next(), Some(&[187, 188, 189][..]));
438        assert_eq!(rowchunk_iter.next(), Some(&[197, 198, 199][..]));
439        assert_eq!(rowchunk_iter.next(), None);
440    }
441
442    #[test]
443    fn test_mut_roi_reverse_iterator() {
444        const STRIDE: usize = 10;
445        const ORIG_W: usize = 10;
446        const ORIG_H: usize = 10;
447        let mut image_data = [0u8; STRIDE * ORIG_H];
448
449        // fill with useful pattern
450        for row in 0..ORIG_H {
451            for col in 0..ORIG_W {
452                image_data[row * STRIDE + col] = (row * 10_usize + col) as u8;
453            }
454        }
455
456        // generate an ROI
457        let width = 2;
458        let height = 4;
459        let (row, col) = (2, 2);
460
461        {
462            // create mutable image of this ROI
463            let mut im = RoiImMut {
464                width,
465                height,
466                stride: STRIDE,
467                buf: &mut image_data[(row * STRIDE + col)..],
468            };
469
470            // test length of backwards iterator
471            {
472                let mut rowchunk_iter = im.rowchunks_exact_mut();
473                assert!(rowchunk_iter.next_back().is_some());
474                assert!(rowchunk_iter.next_back().is_some());
475                assert!(rowchunk_iter.next_back().is_some());
476                assert!(rowchunk_iter.next_back().is_some());
477                assert!(rowchunk_iter.next_back().is_none());
478            }
479
480            // test interleaved forward and backward iteration
481            {
482                let mut rowchunk_iter = im.rowchunks_exact_mut();
483                let mut row5 = rowchunk_iter.next_back();
484                assert_eq!(row5, Some(&mut [52, 53][..]));
485                row5.as_mut().unwrap()[0] += 100;
486                row5.as_mut().unwrap()[1] += 100;
487                let mut row2 = rowchunk_iter.next();
488                assert_eq!(row2, Some(&mut [22, 23][..]));
489                row2.as_mut().unwrap()[0] += 100;
490                row2.as_mut().unwrap()[1] += 100;
491
492                let mut row4 = rowchunk_iter.next_back();
493                assert_eq!(row4, Some(&mut [42, 43][..]));
494                row4.as_mut().unwrap()[0] += 100;
495                row4.as_mut().unwrap()[1] += 100;
496
497                let mut row3 = rowchunk_iter.next();
498                assert_eq!(row3, Some(&mut [32, 33][..]));
499                row3.as_mut().unwrap()[0] += 100;
500                row3.as_mut().unwrap()[1] += 100;
501
502                assert_eq!(rowchunk_iter.next_back(), None);
503            }
504        }
505
506        // create image of this ROI
507        let im = RoiIm {
508            width,
509            height,
510            stride: STRIDE,
511            buf: &image_data[(row * STRIDE + col)..],
512        };
513
514        let mut rowchunk_iter = im.rowchunks_exact();
515        assert_eq!(rowchunk_iter.next(), Some(&[122, 123][..]));
516        assert_eq!(rowchunk_iter.next(), Some(&[132, 133][..]));
517        assert_eq!(rowchunk_iter.next(), Some(&[142, 143][..]));
518        assert_eq!(rowchunk_iter.next(), Some(&[152, 153][..]));
519        assert_eq!(rowchunk_iter.next(), None);
520    }
521
522    #[test]
523    fn test_roi_reverse_iterator() {
524        const STRIDE: usize = 10;
525        const ORIG_W: usize = 10;
526        const ORIG_H: usize = 10;
527        let mut image_data = [0u8; STRIDE * ORIG_H];
528
529        // fill with useful pattern
530        for row in 0..ORIG_H {
531            for col in 0..ORIG_W {
532                image_data[row * STRIDE + col] = (row * 10_usize + col) as u8;
533            }
534        }
535
536        // generate an ROI
537        let width = 2;
538        let height = 4;
539        let (row, col) = (2, 2);
540
541        {
542            // create image of this ROI
543            let im = RoiIm {
544                width,
545                height,
546                stride: STRIDE,
547                buf: &image_data[(row * STRIDE + col)..],
548            };
549
550            // test length of backwards iterator
551            {
552                let mut rowchunk_iter = im.rowchunks_exact();
553                assert!(rowchunk_iter.next_back().is_some());
554                assert!(rowchunk_iter.next_back().is_some());
555                assert!(rowchunk_iter.next_back().is_some());
556                assert!(rowchunk_iter.next_back().is_some());
557                assert!(rowchunk_iter.next_back().is_none());
558            }
559
560            // test interleaved forward and backward iteration
561            {
562                let mut rowchunk_iter = im.rowchunks_exact();
563                let row5 = rowchunk_iter.next_back();
564                assert_eq!(row5, Some(&[52, 53][..]));
565                let row2 = rowchunk_iter.next();
566                assert_eq!(row2, Some(&[22, 23][..]));
567                let row4 = rowchunk_iter.next_back();
568                assert_eq!(row4, Some(&[42, 43][..]));
569                let row3 = rowchunk_iter.next();
570                assert_eq!(row3, Some(&[32, 33][..]));
571                assert_eq!(rowchunk_iter.next_back(), None);
572            }
573        }
574    }
575}
machine_vision_formats/iter.rs

machine_vision_formats/
iter.rs
