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