1use bayer as wang_debayer;
10use machine_vision_formats as formats;
11
12use formats::{
13 cow::CowImage,
14 image_ref::{ImageRef, ImageRefMut},
15 iter::{HasRowChunksExact, HasRowChunksExactMut},
16 owned::OImage,
17 pixel_format::{Mono8, NV12, RGB8},
18 ImageBufferMutRef, OwnedImageStride, PixFmt, PixelFormat,
19};
20
21type Result<T> = std::result::Result<T, Error>;
22
23#[derive(thiserror::Error, Debug)]
25pub enum Error {
26 #[error("unimplemented pixel_format: {0:?}")]
27 UnimplementedPixelFormat(PixFmt),
28 #[error("unimplemented ROI width conversion")]
29 UnimplementedRoiWidthConversion,
30 #[error("ROI size exceeds original image")]
31 RoiExceedsOriginal,
32 #[error("invalid allocated buffer size")]
33 InvalidAllocatedBufferSize,
34 #[error("invalid allocated buffer stride")]
35 InvalidAllocatedBufferStride,
36 #[error("{0}")]
37 Bayer(#[from] wang_debayer::BayerError),
38 #[error("io error: {0}")]
39 Io(#[from] std::io::Error),
40 #[error("{0}")]
41 Image(#[from] image::ImageError),
42 #[error("unimplemented conversion {0} -> {1}")]
43 UnimplementedConversion(PixFmt, PixFmt),
44}
45
46#[inline]
47const fn calc_min_stride(w: u32, pixfmt: PixFmt) -> usize {
48 w as usize * pixfmt.bits_per_pixel() as usize / 8
49}
50
51#[inline]
52const fn calc_min_buf_size(w: u32, h: u32, stride: usize, pixfmt: PixFmt) -> usize {
53 if h == 0 {
54 return 0;
55 }
56 let all_but_last = (h - 1) as usize * stride;
57 let last = calc_min_stride(w, pixfmt);
58 debug_assert!(stride >= last);
59 all_but_last + last
60}
61
62#[allow(non_camel_case_types)]
63#[allow(non_snake_case)]
64#[derive(PartialEq, Eq, Debug)]
65struct RGB888 {
66 R: u8,
67 G: u8,
68 B: u8,
69}
70
71#[cfg(test)]
72impl RGB888 {
73 fn max_channel_distance(&self, other: &RGB888) -> i32 {
74 let dr = (self.R as i32 - other.R as i32).abs();
75 let dg = (self.G as i32 - other.G as i32).abs();
76 let db = (self.B as i32 - other.B as i32).abs();
77
78 let m1 = if dr > dg { dr } else { dg };
79
80 if m1 > db {
81 m1
82 } else {
83 db
84 }
85 }
86
87 fn distance(&self, other: &RGB888) -> i32 {
88 let dr = (self.R as i32 - other.R as i32).abs();
89 let dg = (self.G as i32 - other.G as i32).abs();
90 let db = (self.B as i32 - other.B as i32).abs();
91 dr + dg + db
92 }
93}
94
95#[allow(non_camel_case_types)]
96#[allow(non_snake_case)]
97#[derive(PartialEq, Eq, Debug)]
98struct YUV444 {
99 Y: u8,
100 U: u8,
101 V: u8,
102}
103
104#[test]
105fn test_f32_to_u8() {
106 assert_eq!(-1.0f32 as u8, 0u8);
108 assert_eq!(-2.0f32 as u8, 0u8);
109 assert_eq!(-100.0f32 as u8, 0u8);
110 assert_eq!(255.0f32 as u8, 255u8);
111 assert_eq!(255.1f32 as u8, 255u8);
112 assert_eq!(255.8f32 as u8, 255u8);
113 assert_eq!(5000.0f32 as u8, 255u8);
114}
115
116#[allow(non_snake_case)]
117fn YUV444_bt601_toRGB(Y: u8, U: u8, V: u8) -> RGB888 {
118 let Y = Y as f32;
120 let U = U as f32 - 128.0;
121 let V = V as f32 - 128.0;
122
123 let R = 1.0 * Y + 1.402 * V;
124 let G = 1.0 * Y + -0.344136 * U + -0.714136 * V;
125 let B = 1.0 * Y + 1.772 * U;
126
127 RGB888 {
128 R: R as u8,
129 G: G as u8,
130 B: B as u8,
131 }
132}
133
134#[allow(non_snake_case)]
135#[inline]
136fn RGB888toYUV444_bt601_full_swing(R: u8, G: u8, B: u8) -> YUV444 {
137 let Y = RGB888toY4_bt601_full_swing(R, G, B);
141 let R = R as i32;
142 let G = G as i32;
143 let B = B as i32;
144 let U = ((-43 * R - 84 * G + 127 * B + 128) >> 8) + 128;
145 let V = ((127 * R - 106 * G - 21 * B + 128) >> 8) + 128;
146 YUV444 {
147 Y,
148 U: U as u8,
149 V: V as u8,
150 }
151}
152
153#[allow(non_snake_case)]
154#[inline]
155fn RGB888toY4_bt601_full_swing(R: u8, G: u8, B: u8) -> u8 {
156 let R = R as i32;
160 let G = G as i32;
161 let B = B as i32;
162 let Y = (77 * R + 150 * G + 29 * B + 128) >> 8;
163 Y as u8
164}
165
166pub fn image_to_rgb8(
168 input: image::DynamicImage,
169) -> Result<impl OwnedImageStride<formats::pixel_format::RGB8>> {
170 let rgb = input.to_rgb8();
171 let (width, height) = rgb.dimensions();
172 let stride = width as usize * 3;
173 let data = rgb.into_vec();
174
175 Ok(OImage::new(width, height, stride, data).unwrap())
176}
177
178fn yuv422_into_rgb(
180 src_yuv422: &dyn HasRowChunksExact<formats::pixel_format::YUV422>,
181 dest_rgb: &mut dyn HasRowChunksExactMut<RGB8>,
182) -> Result<()> {
183 if src_yuv422.width() != dest_rgb.width() || src_yuv422.height() != dest_rgb.height() {
184 return Err(Error::InvalidAllocatedBufferSize);
185 }
186
187 let min_stride = src_yuv422.width() as usize * PixFmt::RGB8.bits_per_pixel() as usize / 8;
189 if dest_rgb.stride() < min_stride {
190 return Err(Error::InvalidAllocatedBufferStride);
191 }
192
193 let w = src_yuv422.width() as usize;
194 for (src_row, dest_row) in src_yuv422
195 .rowchunks_exact()
196 .zip(dest_rgb.rowchunks_exact_mut())
197 {
198 for (result_chunk, yuv422_pixpair) in dest_row[..(w * 3)]
199 .chunks_exact_mut(6)
200 .zip(src_row[..w * 2].chunks_exact(4))
201 {
202 let u = yuv422_pixpair[0];
203 let y1 = yuv422_pixpair[1];
204 let v = yuv422_pixpair[2];
205 let y2 = yuv422_pixpair[3];
206
207 let tmp_rgb1 = YUV444_bt601_toRGB(y1, u, v);
208 let tmp_rgb2 = YUV444_bt601_toRGB(y2, u, v);
209
210 result_chunk[0] = tmp_rgb1.R;
211 result_chunk[1] = tmp_rgb1.G;
212 result_chunk[2] = tmp_rgb1.B;
213
214 result_chunk[3] = tmp_rgb2.R;
215 result_chunk[4] = tmp_rgb2.G;
216 result_chunk[5] = tmp_rgb2.B;
217 }
218 }
219 Ok(())
220}
221
222fn into_yuv444<FMT>(
223 frame: &dyn HasRowChunksExact<FMT>,
224 dest: &mut dyn HasRowChunksExactMut<machine_vision_formats::pixel_format::YUV444>,
225) -> Result<()>
226where
227 FMT: PixelFormat,
228{
229 if frame.width() != dest.width() || frame.height() != dest.height() {
230 return Err(Error::InvalidAllocatedBufferSize);
231 }
232 let w = frame.width() as usize;
233
234 let frame = to_rgb8_or_mono8(frame)?;
240
241 match &frame {
242 SupportedEncoding::Mono(mono) => {
243 for (dest_row, src_row) in dest.rowchunks_exact_mut().zip(mono.rowchunks_exact()) {
244 for (dest_pixel, src_pixel) in
245 dest_row[..(w * 3)].chunks_exact_mut(3).zip(&src_row[..w])
246 {
247 let yuv = YUV444 {
248 Y: *src_pixel,
249 U: 128,
250 V: 128,
251 };
252 dest_pixel[0] = yuv.Y;
253 dest_pixel[1] = yuv.U;
254 dest_pixel[2] = yuv.V;
255 }
256 }
257 }
258 SupportedEncoding::Rgb(rgb) => {
259 for (dest_row, src_row) in dest.rowchunks_exact_mut().zip(rgb.rowchunks_exact()) {
260 for (dest_pixel, src_pixel) in dest_row[..(w * 3)]
261 .chunks_exact_mut(3)
262 .zip(src_row[..(w * 3)].chunks_exact(3))
263 {
264 let yuv =
265 RGB888toYUV444_bt601_full_swing(src_pixel[0], src_pixel[1], src_pixel[2]);
266 dest_pixel[0] = yuv.Y;
267 dest_pixel[1] = yuv.U;
268 dest_pixel[2] = yuv.V;
269 }
270 }
271 }
272 };
273 Ok(())
274}
275
276fn bayer_into_rgb<FMT>(
278 frame: &dyn HasRowChunksExact<FMT>,
279 dest_rgb: &mut dyn HasRowChunksExactMut<RGB8>,
280) -> Result<()>
281where
282 FMT: formats::PixelFormat,
283{
284 let dest_stride = dest_rgb.stride();
285
286 if frame.stride() != frame.width() as usize {
287 return Err(Error::UnimplementedRoiWidthConversion);
288 }
289
290 let expected_stride = frame.width() as usize * PixFmt::RGB8.bits_per_pixel() as usize / 8;
292 if dest_stride != expected_stride {
293 return Err(Error::InvalidAllocatedBufferStride);
294 }
295
296 let src_fmt = machine_vision_formats::pixel_format::pixfmt::<FMT>().unwrap();
297
298 let cfa = match src_fmt {
299 formats::pixel_format::PixFmt::BayerRG8 => wang_debayer::CFA::RGGB,
300 formats::pixel_format::PixFmt::BayerGB8 => wang_debayer::CFA::GBRG,
301 formats::pixel_format::PixFmt::BayerGR8 => wang_debayer::CFA::GRBG,
302 formats::pixel_format::PixFmt::BayerBG8 => wang_debayer::CFA::BGGR,
303 _ => {
304 return Err(Error::UnimplementedPixelFormat(src_fmt));
305 }
306 };
307
308 use std::io::Cursor;
309
310 {
311 let mut dst = wang_debayer::RasterMut::new(
312 frame.width() as usize,
313 frame.height() as usize,
314 wang_debayer::RasterDepth::Depth8,
315 dest_rgb.buffer_mut_ref().data,
316 );
317
318 wang_debayer::run_demosaic(
319 &mut Cursor::new(&frame.image_data()),
320 wang_debayer::BayerDepth::Depth8,
321 cfa,
322 wang_debayer::Demosaic::Cubic,
323 &mut dst,
324 )?;
325 }
326 Ok(())
327}
328
329fn mono8_into_rgb8(
333 src: &dyn HasRowChunksExact<formats::pixel_format::Mono8>,
334 dest_rgb: &mut dyn HasRowChunksExactMut<RGB8>,
335) -> Result<()> {
336 let dest_stride = dest_rgb.stride();
337 let min_stride = src.width() as usize * PixFmt::RGB8.bits_per_pixel() as usize / 8;
339 if dest_stride < min_stride {
340 return Err(Error::InvalidAllocatedBufferStride);
341 }
342
343 let w = src.width() as usize;
344 for (src_row, dest_row) in src.rowchunks_exact().zip(dest_rgb.rowchunks_exact_mut()) {
345 for (dest_pix, src_pix) in dest_row[..(w * 3)].chunks_exact_mut(3).zip(&src_row[..w]) {
346 dest_pix[0] = *src_pix;
347 dest_pix[1] = *src_pix;
348 dest_pix[2] = *src_pix;
349 }
350 }
351 Ok(())
352}
353
354fn rgba_into_rgb(
356 frame: &dyn HasRowChunksExact<formats::pixel_format::RGBA8>,
357 dest: &mut dyn HasRowChunksExactMut<RGB8>,
358) -> Result<()> {
359 let dest_stride = dest.stride();
360
361 let min_stride = frame.width() as usize * PixFmt::RGB8.bits_per_pixel() as usize / 8;
363 if dest_stride < min_stride {
364 return Err(Error::InvalidAllocatedBufferStride);
365 }
366
367 let w = frame.width() as usize;
368 for (src_row, dest_row) in frame.rowchunks_exact().zip(dest.rowchunks_exact_mut()) {
369 for (dest_pix, src_pix) in dest_row[..(w * 3)]
370 .chunks_exact_mut(3)
371 .zip(src_row[..(w * 4)].chunks_exact(4))
372 {
373 dest_pix[0] = src_pix[0];
374 dest_pix[1] = src_pix[1];
375 dest_pix[2] = src_pix[2];
376 }
378 }
379 Ok(())
380}
381
382fn rgb8_into_mono8(
384 frame: &dyn HasRowChunksExact<formats::pixel_format::RGB8>,
385 dest: &mut dyn HasRowChunksExactMut<Mono8>,
386) -> Result<()> {
387 if !(dest.height() == frame.height() && dest.width() == frame.width()) {
388 return Err(Error::InvalidAllocatedBufferSize);
389 }
390
391 let w = frame.width() as usize;
392 for (src_row, dest_row) in frame.rowchunks_exact().zip(dest.rowchunks_exact_mut()) {
393 let y_iter = src_row[..w * 3]
394 .chunks_exact(3)
395 .map(|rgb| RGB888toY4_bt601_full_swing(rgb[0], rgb[1], rgb[2]));
396
397 let dest_iter = dest_row[0..w].iter_mut();
398
399 for (ydest, y) in dest_iter.zip(y_iter) {
400 *ydest = y;
401 }
402 }
403
404 Ok(())
405}
406
407fn yuv444_into_mono8(
409 frame: &dyn HasRowChunksExact<formats::pixel_format::YUV444>,
410 dest: &mut dyn HasRowChunksExactMut<Mono8>,
411) -> Result<()> {
412 if !(dest.height() == frame.height() && dest.width() == frame.width()) {
413 return Err(Error::InvalidAllocatedBufferSize);
414 }
415
416 let w = frame.width() as usize;
417 for (src_row, dest_row) in frame.rowchunks_exact().zip(dest.rowchunks_exact_mut()) {
418 let y_iter = src_row[..w * 3].chunks_exact(3).map(|yuv444| yuv444[0]);
419
420 let dest_iter = dest_row[0..w].iter_mut();
421
422 for (ydest, y) in dest_iter.zip(y_iter) {
423 *ydest = y;
424 }
425 }
426
427 Ok(())
428}
429
430fn nv12_into_mono8(
432 frame: &dyn HasRowChunksExact<formats::pixel_format::NV12>,
433 dest: &mut dyn HasRowChunksExactMut<Mono8>,
434) -> Result<()> {
435 if !(dest.height() == frame.height() && dest.width() == frame.width()) {
436 return Err(Error::InvalidAllocatedBufferSize);
437 }
438
439 for (src_row, dest_row) in frame.rowchunks_exact().zip(dest.rowchunks_exact_mut()) {
440 dest_row[..frame.width() as usize].copy_from_slice(&src_row[..frame.width() as usize]);
441 }
442
443 Ok(())
444}
445
446fn remove_padding<FMT>(frame: &dyn HasRowChunksExact<FMT>) -> Result<CowImage<'_, FMT>>
448where
449 FMT: PixelFormat,
450{
451 let fmt = machine_vision_formats::pixel_format::pixfmt::<FMT>().unwrap();
452 let bytes_per_pixel = fmt.bits_per_pixel() as usize / 8;
453 let dest_stride = frame.width() as usize * bytes_per_pixel;
454 if dest_stride == frame.stride() {
455 Ok(CowImage::Borrowed(force_pixel_format_ref(frame)))
456 } else {
457 if frame.stride() < dest_stride {
458 return Err(Error::InvalidAllocatedBufferStride);
459 }
460 let mut dest_buf = vec![0u8; frame.height() as usize * dest_stride];
462 frame
464 .rowchunks_exact()
465 .zip(dest_buf.chunks_exact_mut(dest_stride))
466 .for_each(|(src_row_full, dest_row)| {
467 dest_row[..dest_stride].copy_from_slice(&src_row_full[..dest_stride]);
468 });
469 Ok(CowImage::Owned(
471 OImage::new(frame.width(), frame.height(), dest_stride, dest_buf).unwrap(),
472 ))
473 }
474}
475
476pub fn force_pixel_format<FRAME, FMT1, FMT2>(frame: FRAME) -> impl OwnedImageStride<FMT2>
482where
483 FRAME: OwnedImageStride<FMT1>,
484 FMT2: PixelFormat,
485{
486 let width = frame.width();
487 let height = frame.height();
488 let stride = frame.stride();
489 let image_data = frame.into(); OImage::new(width, height, stride, image_data).unwrap()
492}
493
494pub fn force_pixel_format_ref<FMT1, FMT2>(frame: &dyn HasRowChunksExact<FMT1>) -> ImageRef<'_, FMT2>
503where
504 FMT2: PixelFormat,
505{
506 ImageRef::new(
507 frame.width(),
508 frame.height(),
509 frame.stride(),
510 frame.image_data(),
511 )
512 .unwrap()
513}
514
515fn force_pixel_format_ref_mut<FMT1, FMT2>(
524 frame: &mut dyn HasRowChunksExactMut<FMT1>,
525) -> ImageRefMut<'_, FMT2>
526where
527 FMT1: PixelFormat,
528 FMT2: PixelFormat,
529{
530 ImageRefMut::new(
531 frame.width(),
532 frame.height(),
533 frame.stride(),
534 frame.buffer_mut_ref().data,
535 )
536 .unwrap()
537}
538
539fn force_buffer_pixel_format_ref<FMT1, FMT2>(
541 orig: ImageBufferMutRef<'_, FMT1>,
542) -> ImageBufferMutRef<'_, FMT2> {
543 ImageBufferMutRef::new(orig.data)
544}
545
546pub fn convert_owned<OWNED, SRC, DEST>(source: OWNED) -> Result<impl HasRowChunksExact<DEST>>
560where
561 OWNED: OwnedImageStride<SRC>,
562 SRC: PixelFormat,
563 DEST: PixelFormat,
564{
565 let src_fmt = machine_vision_formats::pixel_format::pixfmt::<SRC>().unwrap();
566 let dest_fmt = machine_vision_formats::pixel_format::pixfmt::<DEST>().unwrap();
567
568 if src_fmt == dest_fmt {
570 let width = source.width();
571 let height = source.height();
572 let stride = source.stride();
573 let buf: Vec<u8> = source.into();
574 let dest = OImage::new(width, height, stride, buf).unwrap();
575 return Ok(CowImage::Owned(dest));
576 }
577
578 let dest_min_stride = dest_fmt.bits_per_pixel() as usize * source.width() as usize / 8;
580 let dest_size = source.height() as usize * dest_min_stride;
581 let image_data = vec![0u8; dest_size];
582 let mut dest =
583 OImage::new(source.width(), source.height(), dest_min_stride, image_data).unwrap();
584
585 convert_into(&source, &mut dest)?;
587
588 Ok(CowImage::Owned(dest))
590}
591
592pub fn convert_ref<SRC, DEST>(
606 source: &dyn HasRowChunksExact<SRC>,
607) -> Result<impl HasRowChunksExact<DEST> + '_>
608where
609 SRC: PixelFormat,
610 DEST: PixelFormat,
611{
612 let src_fmt = machine_vision_formats::pixel_format::pixfmt::<SRC>().unwrap();
613 let dest_fmt = machine_vision_formats::pixel_format::pixfmt::<DEST>().unwrap();
614
615 if src_fmt == dest_fmt {
617 return Ok(CowImage::Borrowed(force_pixel_format_ref(source)));
618 }
619
620 let dest_stride = calc_min_stride(source.width(), dest_fmt);
622 let dest_size = calc_min_buf_size(source.width(), source.height(), dest_stride, dest_fmt);
623 let image_data = vec![0u8; dest_size];
624 let mut dest = OImage::new(source.width(), source.height(), dest_stride, image_data).unwrap();
625
626 convert_into(source, &mut dest)?;
628
629 Ok(CowImage::Owned(dest))
631}
632
633pub fn convert_into<SRC, DEST>(
640 source: &dyn HasRowChunksExact<SRC>,
641 dest: &mut dyn HasRowChunksExactMut<DEST>,
642) -> Result<()>
643where
644 SRC: PixelFormat,
645 DEST: PixelFormat,
646{
647 let src_fmt = machine_vision_formats::pixel_format::pixfmt::<SRC>().unwrap();
648 let dest_fmt = machine_vision_formats::pixel_format::pixfmt::<DEST>().unwrap();
649
650 if dest.width() != source.width() || dest.height() != source.height() {
651 return Err(Error::InvalidAllocatedBufferSize);
652 }
653
654 let dest_stride = dest.stride();
655
656 if src_fmt == dest_fmt {
658 use itertools::izip;
659 let dest_stride = calc_min_stride(source.width(), dest_fmt);
660 for (src_row, dest_row) in izip![source.rowchunks_exact(), dest.rowchunks_exact_mut(),] {
661 debug_assert_eq!(src_row.len(), dest_stride);
662 dest_row[..dest_stride].copy_from_slice(src_row);
663 }
664 return Ok(());
665 }
666
667 match dest_fmt {
668 formats::pixel_format::PixFmt::RGB8 => {
669 let mut dest_rgb = ImageRefMut::new(
670 dest.width(),
671 dest.height(),
672 dest.stride(),
673 dest.buffer_mut_ref().data,
674 )
675 .unwrap();
676 match src_fmt {
678 formats::pixel_format::PixFmt::BayerRG8
679 | formats::pixel_format::PixFmt::BayerGB8
680 | formats::pixel_format::PixFmt::BayerGR8
681 | formats::pixel_format::PixFmt::BayerBG8 => {
682 let exact_stride = remove_padding(source)?;
685 bayer_into_rgb(&exact_stride, &mut dest_rgb)?;
686 Ok(())
687 }
688 formats::pixel_format::PixFmt::Mono8 => {
689 let mono8 = force_pixel_format_ref(source);
691 mono8_into_rgb8(&mono8, &mut dest_rgb)?;
692 Ok(())
693 }
694 formats::pixel_format::PixFmt::RGBA8 => {
695 let rgba8 = force_pixel_format_ref(source);
697 rgba_into_rgb(&rgba8, &mut dest_rgb)?;
698 Ok(())
699 }
700 formats::pixel_format::PixFmt::YUV422 => {
701 let yuv422 = force_pixel_format_ref(source);
703 yuv422_into_rgb(&yuv422, &mut dest_rgb)?;
704 Ok(())
705 }
706 _ => Err(Error::UnimplementedConversion(src_fmt, dest_fmt)),
707 }
708 }
709 formats::pixel_format::PixFmt::Mono8 => {
710 let mut dest_mono8 = ImageRefMut::new(
711 dest.width(),
712 dest.height(),
713 dest.stride(),
714 dest.buffer_mut_ref().data,
715 )
716 .unwrap();
717
718 match src_fmt {
720 formats::pixel_format::PixFmt::RGB8 => {
721 let tmp = force_pixel_format_ref(source);
723 {
724 rgb8_into_mono8(&tmp, &mut dest_mono8)?;
725 }
726 Ok(())
727 }
728 formats::pixel_format::PixFmt::YUV444 => {
729 let yuv444 = force_pixel_format_ref(source);
731 yuv444_into_mono8(&yuv444, &mut dest_mono8)?;
732 Ok(())
733 }
734 formats::pixel_format::PixFmt::NV12 => {
735 let nv12 = force_pixel_format_ref(source);
737 nv12_into_mono8(&nv12, &mut dest_mono8)?;
738 Ok(())
739 }
740 formats::pixel_format::PixFmt::BayerRG8
741 | formats::pixel_format::PixFmt::BayerGB8
742 | formats::pixel_format::PixFmt::BayerGR8
743 | formats::pixel_format::PixFmt::BayerBG8 => {
744 let width: usize = source.width().try_into().unwrap();
747 let height: usize = source.height().try_into().unwrap();
748 let stride = width * 3;
749 let rgb_buf = &mut vec![0u8; stride * height];
750 let mut tmp_rgb =
751 ImageRefMut::new(source.width(), source.height(), stride, rgb_buf).unwrap();
752 let exact_stride = remove_padding(source)?;
754 bayer_into_rgb(&exact_stride, &mut tmp_rgb)?;
755 rgb8_into_mono8(&tmp_rgb, &mut dest_mono8)?;
757 Ok(())
758 }
759 _ => Err(Error::UnimplementedConversion(src_fmt, dest_fmt)),
760 }
761 }
762 formats::pixel_format::PixFmt::YUV444 => {
763 let mut dest2: ImageRefMut<'_, machine_vision_formats::pixel_format::YUV444> =
765 force_pixel_format_ref_mut(dest);
766 into_yuv444(source, &mut dest2)?;
767 Ok(())
768 }
769 formats::pixel_format::PixFmt::NV12 => {
770 let mut dest2 = force_buffer_pixel_format_ref(dest.buffer_mut_ref());
772 encode_into_nv12_inner(source, &mut dest2, dest_stride)?;
773 Ok(())
774 }
775 _ => Err(Error::UnimplementedConversion(src_fmt, dest_fmt)),
776 }
777}
778
779enum SupportedEncoding<'a> {
786 Rgb(Box<dyn HasRowChunksExact<formats::pixel_format::RGB8> + 'a>),
787 Mono(Box<dyn HasRowChunksExact<formats::pixel_format::Mono8> + 'a>),
788}
789
790impl SupportedEncoding<'_> {
791 #[inline]
792 fn width(&self) -> u32 {
793 match self {
794 SupportedEncoding::Rgb(m) => m.width(),
795 SupportedEncoding::Mono(m) => m.width(),
796 }
797 }
798 #[inline]
799 fn height(&self) -> u32 {
800 match self {
801 SupportedEncoding::Rgb(m) => m.height(),
802 SupportedEncoding::Mono(m) => m.height(),
803 }
804 }
805 #[inline]
806 fn stride(&self) -> usize {
807 match self {
808 SupportedEncoding::Rgb(m) => m.stride(),
809 SupportedEncoding::Mono(m) => m.stride(),
810 }
811 }
812 #[inline]
813 fn image_data(&self) -> &[u8] {
814 match self {
815 SupportedEncoding::Rgb(m) => m.image_data(),
816 SupportedEncoding::Mono(m) => m.image_data(),
817 }
818 }
819 #[inline]
820 fn rowchunks_exact(&self) -> machine_vision_formats::iter::RowChunksExact {
821 match self {
822 SupportedEncoding::Rgb(m) => m.rowchunks_exact(),
823 SupportedEncoding::Mono(m) => m.rowchunks_exact(),
824 }
825 }
826 #[inline]
827 fn pixfmt(&self) -> formats::pixel_format::PixFmt {
828 match self {
829 SupportedEncoding::Rgb(_) => {
830 machine_vision_formats::pixel_format::pixfmt::<RGB8>().unwrap()
831 }
832 SupportedEncoding::Mono(_) => {
833 machine_vision_formats::pixel_format::pixfmt::<Mono8>().unwrap()
834 }
835 }
836 }
837}
838
839fn to_rgb8_or_mono8<FMT>(frame: &dyn HasRowChunksExact<FMT>) -> Result<SupportedEncoding<'_>>
841where
842 FMT: PixelFormat,
843{
844 if machine_vision_formats::pixel_format::pixfmt::<FMT>().unwrap()
845 == formats::pixel_format::PixFmt::Mono8
846 {
847 let im = convert_ref::<_, formats::pixel_format::Mono8>(frame)?;
848 Ok(SupportedEncoding::Mono(Box::new(im)))
849 } else {
850 let im = convert_ref::<_, formats::pixel_format::RGB8>(frame)?;
851 Ok(SupportedEncoding::Rgb(Box::new(im)))
852 }
853}
854
855pub fn frame_to_image<FMT>(frame: &dyn HasRowChunksExact<FMT>) -> Result<image::DynamicImage>
857where
858 FMT: PixelFormat,
859{
860 let frame = to_rgb8_or_mono8(frame)?;
861
862 let (coding, pixfmt) = match &frame {
863 SupportedEncoding::Mono(_) => {
864 let pixfmt = machine_vision_formats::pixel_format::pixfmt::<Mono8>().unwrap();
865 debug_assert_eq!(frame.pixfmt(), pixfmt);
866 (image::ColorType::L8, pixfmt)
867 }
868 SupportedEncoding::Rgb(_) => {
869 let pixfmt = machine_vision_formats::pixel_format::pixfmt::<RGB8>().unwrap();
870 debug_assert_eq!(frame.pixfmt(), pixfmt);
871 (image::ColorType::Rgb8, pixfmt)
872 }
873 };
874
875 let packed_stride = calc_min_stride(frame.width(), pixfmt);
876
877 let mut packed = None;
881 if frame.stride() != packed_stride {
882 let dest_sz = calc_min_buf_size(frame.width(), frame.height(), packed_stride, pixfmt);
883 let mut dest = Vec::with_capacity(dest_sz);
884 for src_row in frame.rowchunks_exact() {
885 debug_assert_eq!(src_row.len(), packed_stride);
886 dest.extend_from_slice(&src_row[..packed_stride]);
887 }
888 packed = Some(dest);
889 }
890
891 let packed = match packed {
892 None => frame.image_data().to_vec(),
893 Some(p) => p,
894 };
895
896 match coding {
897 image::ColorType::L8 => {
898 let imbuf: image::ImageBuffer<image::Luma<_>, _> =
899 image::ImageBuffer::from_raw(frame.width(), frame.height(), packed).unwrap();
900 Ok(imbuf.into())
901 }
902 image::ColorType::Rgb8 => {
903 let imbuf: image::ImageBuffer<image::Rgb<_>, _> =
904 image::ImageBuffer::from_raw(frame.width(), frame.height(), packed).unwrap();
905 Ok(imbuf.into())
906 }
907 _ => {
908 unreachable!()
909 }
910 }
911}
912
913#[derive(Copy, Clone, Debug, PartialEq, Eq)]
915pub enum EncoderOptions {
916 Jpeg(u8),
918 Png,
920}
921
922pub fn frame_to_encoded_buffer<FMT>(
925 frame: &dyn HasRowChunksExact<FMT>,
926 opts: EncoderOptions,
927) -> Result<Vec<u8>>
928where
929 FMT: PixelFormat,
930{
931 let mut result = Vec::new();
932
933 let frame = to_rgb8_or_mono8(frame)?;
934
935 let (coding, bytes_per_pixel) = match &frame {
936 SupportedEncoding::Mono(_) => (image::ColorType::L8, 1),
937 SupportedEncoding::Rgb(_) => (image::ColorType::Rgb8, 3),
938 };
939
940 let mut packed = None;
944 let packed_stride = frame.width() as usize * bytes_per_pixel as usize;
945 if frame.stride() != packed_stride {
946 let mut dest = Vec::with_capacity(packed_stride * frame.height() as usize);
947
948 for src_row in frame.rowchunks_exact() {
949 dest.extend_from_slice(&src_row[..]);
950 }
951
952 packed = Some(dest);
953 }
954
955 let use_frame = match &packed {
956 None => frame.image_data(),
957 Some(p) => p.as_slice(),
958 };
959
960 match opts {
961 EncoderOptions::Jpeg(quality) => {
962 let mut encoder =
963 image::codecs::jpeg::JpegEncoder::new_with_quality(&mut result, quality);
964 encoder.encode(use_frame, frame.width(), frame.height(), coding.into())?;
965 }
966 EncoderOptions::Png => {
967 use image::ImageEncoder;
968 let encoder = image::codecs::png::PngEncoder::new(&mut result);
969 encoder.write_image(use_frame, frame.width(), frame.height(), coding.into())?;
970 }
971 }
972 Ok(result)
973}
974
975fn encode_into_nv12_inner<FMT>(
976 frame: &dyn HasRowChunksExact<FMT>,
977 dest: &mut ImageBufferMutRef<NV12>,
978 dest_stride: usize,
979) -> Result<()>
980where
981 FMT: PixelFormat,
982{
983 use itertools::izip;
984
985 let frame = to_rgb8_or_mono8(frame)?;
986
987 let luma_size = frame.height() as usize * dest_stride;
988
989 let (nv12_luma, nv12_chroma) = dest.data.split_at_mut(luma_size);
990
991 match &frame {
992 SupportedEncoding::Mono(frame) => {
993 let w = frame.width() as usize;
995 for y in 0..frame.height() as usize {
996 let start = dest_stride * y;
997 let src = frame.stride() * y;
998 nv12_luma[start..(start + w)].copy_from_slice(&frame.image_data()[src..(src + w)]);
999 }
1000
1001 for y in 0..(frame.height() as usize / 2) {
1002 let start = dest_stride * y;
1003 for x in (0..frame.width() as usize).step_by(2) {
1004 nv12_chroma[start + x] = 128u8;
1005 nv12_chroma[start + (x + 1)] = 128u8;
1006 }
1007 }
1008 }
1009 SupportedEncoding::Rgb(frame) => {
1010 let w = frame.width() as usize;
1011
1012 let mut u_plane_full: Vec<u8> = vec![0; nv12_luma.len()];
1015 let mut v_plane_full: Vec<u8> = vec![0; nv12_luma.len()];
1016
1017 for (src_row, dest_row, udest_row, vdest_row) in izip![
1018 frame.image_data().chunks_exact(frame.stride()),
1019 nv12_luma.chunks_exact_mut(dest_stride),
1020 u_plane_full.chunks_exact_mut(dest_stride),
1021 v_plane_full.chunks_exact_mut(dest_stride),
1022 ] {
1023 let yuv_iter = src_row[..w * 3]
1024 .chunks_exact(3)
1025 .map(|rgb| RGB888toYUV444_bt601_full_swing(rgb[0], rgb[1], rgb[2]));
1026
1027 let dest_iter = dest_row[0..w].iter_mut();
1028
1029 for (ydest, udest, vdest, yuv) in izip![dest_iter, udest_row, vdest_row, yuv_iter] {
1030 *ydest = yuv.Y;
1031 *udest = yuv.U;
1032 *vdest = yuv.V;
1033 }
1034 }
1035
1036 let half_stride = dest_stride; for y in 0..(frame.height() as usize / 2) {
1039 for x in 0..(frame.width() as usize / 2) {
1040 let u_sum: u16 = u_plane_full[dest_stride * 2 * y + 2 * x] as u16
1041 + u_plane_full[dest_stride * 2 * y + 2 * x + 1] as u16
1042 + u_plane_full[dest_stride * (2 * y + 1) + 2 * x] as u16
1043 + u_plane_full[dest_stride * (2 * y + 1) + 2 * x + 1] as u16;
1044 let v_sum: u16 = v_plane_full[dest_stride * 2 * y + 2 * x] as u16
1045 + v_plane_full[dest_stride * 2 * y + 2 * x + 1] as u16
1046 + v_plane_full[dest_stride * (2 * y + 1) + 2 * x] as u16
1047 + v_plane_full[dest_stride * (2 * y + 1) + 2 * x + 1] as u16;
1048
1049 nv12_chroma[(half_stride * y) + 2 * x] = (u_sum / 4) as u8;
1050 nv12_chroma[(half_stride * y) + 2 * x + 1] = (v_sum / 4) as u8;
1051 }
1052 }
1053 }
1054 }
1055 Ok(())
1056}
1057
1058#[cfg(test)]
1059mod tests {
1060 use crate::*;
1061 use formats::{ImageBuffer, ImageBufferRef, ImageData, Stride};
1062
1063 struct RoiImage<'a, F> {
1066 data: &'a [u8],
1067 w: u32,
1068 h: u32,
1069 stride: usize,
1070 fmt: std::marker::PhantomData<F>,
1071 }
1072
1073 impl<'a, F> RoiImage<'a, F>
1074 where
1075 F: PixelFormat,
1076 {
1077 fn new(
1079 frame: &'a dyn HasRowChunksExact<F>,
1080 w: u32,
1081 h: u32,
1082 x: u32,
1083 y: u32,
1084 ) -> Result<RoiImage<'a, F>> {
1085 let stride = frame.stride();
1086 let fmt = machine_vision_formats::pixel_format::pixfmt::<F>().unwrap();
1087 let col_offset = x as usize * fmt.bits_per_pixel() as usize / 8;
1088 if col_offset >= stride {
1089 return Err(Error::RoiExceedsOriginal);
1090 }
1091 let offset = y as usize * stride + col_offset;
1092 Ok(RoiImage {
1093 data: &frame.image_data()[offset..],
1094 w,
1095 h,
1096 stride,
1097 fmt: std::marker::PhantomData,
1098 })
1099 }
1100 }
1101
1102 impl<F> Stride for RoiImage<'_, F> {
1103 fn stride(&self) -> usize {
1104 self.stride
1105 }
1106 }
1107
1108 impl<F> ImageData<F> for RoiImage<'_, F> {
1109 fn width(&self) -> u32 {
1110 self.w
1111 }
1112 fn height(&self) -> u32 {
1113 self.h
1114 }
1115 fn buffer_ref(&self) -> ImageBufferRef<'_, F> {
1116 let image_data = self.data;
1117 ImageBufferRef::new(image_data)
1118 }
1119 fn buffer(self) -> ImageBuffer<F> {
1120 let copied = self.data.to_vec();
1121 ImageBuffer::new(copied)
1122 }
1123 }
1124
1125 fn imstr<F>(frame: &dyn HasRowChunksExact<F>) -> String
1126 where
1127 F: PixelFormat,
1128 {
1129 let fmt = machine_vision_formats::pixel_format::pixfmt::<F>().unwrap();
1130 let bytes_per_pixel = fmt.bits_per_pixel() as usize / 8;
1131
1132 frame
1133 .rowchunks_exact()
1134 .map(|row| {
1135 let image_row = &row[..frame.width() as usize * bytes_per_pixel];
1136 image_row
1137 .chunks_exact(fmt.bits_per_pixel() as usize / 8)
1138 .map(|x| format!("{:?}", x))
1139 .collect::<Vec<_>>()
1140 .join(", ")
1141 })
1142 .collect::<Vec<_>>()
1143 .join("\n")
1144 }
1145
1146 #[test]
1147 fn check_roi_mono8() {
1148 const STRIDE: usize = 10;
1150 const W: u32 = 8;
1151 const H: u32 = 6;
1152 let mut image_data = vec![255; H as usize * STRIDE];
1154 for row in 0..H as usize {
1156 let start_idx = row * STRIDE;
1157 for col in 0..W as usize {
1158 image_data[start_idx + col] = (row * W as usize + col) as u8;
1159 }
1160 }
1161 let frame: OImage<formats::pixel_format::Mono8> =
1162 OImage::new(W, H, STRIDE, image_data).unwrap();
1163 println!("frame: {:?}", frame.image_data());
1164 println!("frame: \n{}", imstr(&frame));
1165 let roi = RoiImage::new(&frame, 6, 2, 1, 1).unwrap();
1166 println!("roi: {:?}", roi.image_data());
1167 println!("roi: \n{}", imstr(&roi));
1168 let small = super::remove_padding(&roi).unwrap();
1169 println!("small: {:?}", small.image_data());
1170 println!("small: \n{}", imstr(&small));
1171 assert_eq!(
1172 small.image_data(),
1173 &[9, 10, 11, 12, 13, 14, 17, 18, 19, 20, 21, 22]
1174 );
1175 }
1176
1177 #[test]
1178 fn check_roi_rgb8() {
1179 const STRIDE: usize = 30;
1181 const W: u32 = 8;
1182 const H: u32 = 6;
1183 let mut image_data = vec![255; H as usize * STRIDE];
1185 for row in 0..H as usize {
1186 let start_idx = row * STRIDE;
1187 for col in 0..W as usize {
1188 let col_offset = col * 3;
1189 image_data[start_idx + col_offset] = ((row * W as usize + col) * 3) as u8;
1190 image_data[start_idx + col_offset + 1] = ((row * W as usize + col) * 3) as u8 + 1;
1191 image_data[start_idx + col_offset + 2] = ((row * W as usize + col) * 3) as u8 + 2;
1192 }
1193 }
1194 let frame: OImage<formats::pixel_format::RGB8> =
1195 OImage::new(W, H, STRIDE, image_data).unwrap();
1196 println!("frame: {:?}", frame.image_data());
1197 println!("frame: \n{}", imstr(&frame));
1198 let roi = RoiImage::new(&frame, 6, 2, 1, 1).unwrap();
1199 println!("roi: {:?}", roi.image_data());
1200 println!("roi: \n{}", imstr(&roi));
1201 let small = super::remove_padding(&roi).unwrap();
1202 println!("small: {:?}", small.image_data());
1203 println!("small: \n{}", imstr(&small));
1204 assert_eq!(
1205 small.image_data(),
1206 &[
1207 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 51, 52, 53,
1208 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68
1209 ]
1210 );
1211 }
1212 #[test]
1213 fn check_stride_conversion_to_image() {
1214 const STRIDE: usize = 6;
1216 const W: u32 = 4;
1217 const H: u32 = 2;
1218 let mut image_data = vec![42; H as usize * STRIDE];
1220 for row in 0..H as usize {
1222 let start_idx = row * STRIDE;
1223 for col in 0..W as usize {
1224 image_data[start_idx + col] = 0;
1225 }
1226 }
1227 let frame: OImage<formats::pixel_format::Mono8> =
1228 OImage::new(W, H, STRIDE, image_data).unwrap();
1229 let buf = frame_to_encoded_buffer(&frame, EncoderOptions::Png).unwrap();
1230
1231 let im2 = image::load_from_memory_with_format(&buf, image::ImageFormat::Png).unwrap();
1233
1234 let im2gray = im2.into_luma8();
1235 assert_eq!(im2gray.width(), W);
1236 assert_eq!(im2gray.height(), H);
1237 for row in 0..H {
1238 for col in 0..W {
1239 let pixel = im2gray.get_pixel(col, row);
1240 assert_eq!(pixel.0[0], 0, "at pixel {},{}", col, row);
1241 }
1242 }
1243 }
1244
1245 #[test]
1246 fn check_bayer_conversion_to_jpg() {
1247 const STRIDE: usize = 6;
1249 const W: u32 = 4;
1250 const H: u32 = 4;
1251 let mut image_data = vec![42; H as usize * STRIDE];
1253 for row in 0..H as usize {
1255 let start_idx = row * STRIDE;
1256 for col in 0..W as usize {
1257 image_data[start_idx + col] = 0;
1258 }
1259 }
1260 let frame: OImage<formats::pixel_format::BayerRG8> =
1261 OImage::new(W, H, STRIDE, image_data).unwrap();
1262 frame_to_encoded_buffer(&frame, EncoderOptions::Jpeg(100)).unwrap();
1263 }
1264
1265 #[test]
1266 fn prevent_unnecessary_copy_mono8() {
1267 let frame: OImage<formats::pixel_format::Mono8> =
1268 OImage::new(10, 10, 10, vec![42; 100]).unwrap();
1269 let im2 = convert_ref::<_, formats::pixel_format::Mono8>(&frame).unwrap();
1271 assert_eq!(im2.image_data(), frame.image_data());
1273
1274 let const_ptr = frame.image_data().as_ptr();
1276 let data_ptr = const_ptr as *mut u8;
1278 unsafe {
1280 *data_ptr = 2;
1281 }
1282 assert_eq!(im2.image_data()[0], 2);
1284 }
1285
1286 #[test]
1287 fn prevent_unnecessary_copy_rgb8() {
1288 let frame: OImage<formats::pixel_format::RGB8> =
1289 OImage::new(10, 10, 30, vec![42; 300]).unwrap();
1290 let im2 = convert_ref::<_, formats::pixel_format::RGB8>(&frame).unwrap();
1292 assert_eq!(im2.image_data(), frame.image_data());
1294
1295 let const_ptr = frame.image_data().as_ptr();
1297 let data_ptr = const_ptr as *mut u8;
1299 unsafe {
1301 *data_ptr = 2;
1302 }
1303 assert_eq!(im2.image_data()[0], 2);
1305 }
1306
1307 #[test]
1308 fn test_rgb_yuv_roundtrip() {
1309 let black_rgb = RGB888 { R: 0, G: 0, B: 0 };
1312 let black_yuv = RGB888toYUV444_bt601_full_swing(black_rgb.R, black_rgb.G, black_rgb.B);
1313 let black_rgb2 = YUV444_bt601_toRGB(black_yuv.Y, black_yuv.U, black_yuv.V);
1314 assert_eq!(black_rgb, black_rgb2);
1315
1316 let white_rgb = RGB888 {
1317 R: 255,
1318 G: 255,
1319 B: 255,
1320 };
1321 let white_yuv = RGB888toYUV444_bt601_full_swing(white_rgb.R, white_rgb.G, white_rgb.B);
1322 let white_rgb2 = YUV444_bt601_toRGB(white_yuv.Y, white_yuv.U, white_yuv.V);
1323 assert_eq!(white_rgb, white_rgb2);
1324
1325 for r in 0..255 {
1326 for g in 0..255 {
1327 for b in 0..255 {
1328 let expected = RGB888 { R: r, G: g, B: b };
1329 let yuv = RGB888toYUV444_bt601_full_swing(expected.R, expected.G, expected.B);
1330 let actual = YUV444_bt601_toRGB(yuv.Y, yuv.U, yuv.V);
1331 assert!(
1332 actual.distance(&expected) <= 7,
1333 "expected: {:?}, actual: {:?}",
1334 expected,
1335 actual
1336 );
1337 assert!(
1338 actual.max_channel_distance(&expected) <= 4,
1339 "expected: {:?}, actual: {:?}",
1340 expected,
1341 actual
1342 );
1343 }
1344 }
1345 }
1346 }
1347
1348 #[test]
1349 fn test_mono8_rgb8() -> Result<()> {
1351 let orig: OImage<formats::pixel_format::Mono8> =
1352 OImage::new(256, 1, 256, (0u8..=255u8).collect()).unwrap();
1353 let rgb = convert_ref::<_, formats::pixel_format::RGB8>(&orig)?;
1354 for (i, rgb_pix) in rgb.image_data().chunks_exact(3).enumerate() {
1355 assert_eq!(i, rgb_pix[0] as usize);
1356 assert_eq!(i, rgb_pix[1] as usize);
1357 assert_eq!(i, rgb_pix[2] as usize);
1358 }
1359 Ok(())
1360 }
1361
1362 #[test]
1363 fn test_mono8_rgb_roundtrip() -> Result<()> {
1364 let orig: OImage<formats::pixel_format::Mono8> =
1365 OImage::new(256, 1, 256, (0u8..=255u8).collect()).unwrap();
1366 let rgb = convert_ref::<_, formats::pixel_format::RGB8>(&orig)?;
1367 let actual = convert_ref::<_, formats::pixel_format::Mono8>(&rgb)?;
1368 assert_eq!(orig.image_data(), actual.image_data());
1369 Ok(())
1370 }
1371
1372 #[test]
1373 fn test_mono8_nv12_roundtrip() -> Result<()> {
1374 let orig: OImage<formats::pixel_format::Mono8> =
1375 OImage::new(256, 1, 256, (0u8..=255u8).collect()).unwrap();
1376 let nv12 = convert_ref::<_, formats::pixel_format::NV12>(&orig)?;
1377 let actual = convert_ref::<_, formats::pixel_format::Mono8>(&nv12)?;
1378 for i in 0..256 {
1379 assert_eq!(orig.image_data()[i], actual.image_data()[i]);
1380 }
1381 assert_eq!(orig.image_data(), actual.image_data());
1382 Ok(())
1383 }
1384
1385 #[test]
1386 fn test_mono8_yuv_roundtrip() -> Result<()> {
1388 let orig: OImage<formats::pixel_format::Mono8> =
1389 OImage::new(256, 1, 256, (0u8..=255u8).collect()).unwrap();
1390 let yuv = convert_ref::<_, formats::pixel_format::YUV444>(&orig)?;
1391 let actual = convert_ref::<_, formats::pixel_format::Mono8>(&yuv)?;
1392 assert_eq!(orig.image_data(), actual.image_data());
1393 Ok(())
1394 }
1395}