1use crate::{
5 Crop, Error, Flip, FunctionTimer, ImageProcessorTrait, Rect, ResolvedCrop, Result, Rotation,
6};
7use edgefirst_tensor::{
8 DType, PixelFormat, Tensor, TensorDyn, TensorMapTrait, TensorMemory, TensorTrait,
9};
10use edgefirst_tensor::{DetectBox, ProtoData, Segmentation};
11
12mod convert;
13mod masks;
14mod resize;
15mod simd;
16mod tests;
17
18#[derive(Debug, Clone, Copy)]
30pub(crate) struct ColorParams {
31 pub matrix: yuv::YuvStandardMatrix,
33 pub range: yuv::YuvRange,
35 pub encoding: edgefirst_tensor::ColorEncoding,
40 pub range_kind: edgefirst_tensor::ColorRange,
41 pub src_full_range: bool,
44 pub dst_full_range: bool,
47}
48
49#[derive(Debug)]
52pub struct CPUProcessor {
53 resizer: fast_image_resize::Resizer,
54 options: fast_image_resize::ResizeOptions,
55 colors: [[u8; 4]; 20],
56 widen_scratch: Option<TensorDyn>,
62 resize_destride_scratch: Vec<u8>,
70 resize_dst_destride_scratch: Vec<u8>,
81 nv_strip_scratch: Vec<u8>,
87 nv_strip_y_pack: Vec<u8>,
103 nv_strip_uv_pack: Vec<u8>,
105 #[cfg(test)]
112 fused_hits: u64,
113 #[cfg(test)]
120 last_tmp_dims: Option<(usize, usize)>,
121 convert_tmp: Option<Tensor<u8>>,
128 convert_tmp2: Option<Tensor<u8>>,
129 convert_src_sub: Option<Tensor<u8>>,
134}
135
136impl Clone for CPUProcessor {
142 fn clone(&self) -> Self {
143 Self {
144 resizer: self.resizer.clone(),
145 options: self.options,
146 colors: self.colors,
147 widen_scratch: None,
148 resize_destride_scratch: Vec::new(),
149 resize_dst_destride_scratch: Vec::new(),
150 nv_strip_scratch: Vec::new(),
151 nv_strip_y_pack: Vec::new(),
152 nv_strip_uv_pack: Vec::new(),
153 #[cfg(test)]
154 fused_hits: 0,
155 #[cfg(test)]
156 last_tmp_dims: None,
157 convert_tmp: None,
158 convert_tmp2: None,
159 convert_src_sub: None,
160 }
161 }
162}
163
164unsafe impl Send for CPUProcessor {}
165unsafe impl Sync for CPUProcessor {}
166
167impl Default for CPUProcessor {
168 fn default() -> Self {
169 Self::new_bilinear()
170 }
171}
172
173fn prepare_dst_base_cpu(dst: &mut TensorDyn, background: Option<&TensorDyn>) -> Result<()> {
184 match background {
185 Some(bg) => {
186 if bg.shape() != dst.shape() {
187 return Err(Error::InvalidShape(
188 "background shape does not match dst".into(),
189 ));
190 }
191 if bg.format() != dst.format() {
192 return Err(Error::InvalidShape(
193 "background pixel format does not match dst".into(),
194 ));
195 }
196 let bg_u8 = bg.as_typed::<u8>().ok_or(Error::NotAnImage)?;
197 let dst_u8 = dst.as_typed_mut::<u8>().ok_or(Error::NotAnImage)?;
198 CPUProcessor::copy_image(bg_u8, dst_u8)?;
205 }
206 None => {
207 let dst_u8 = dst.as_typed_mut::<u8>().ok_or(Error::NotAnImage)?;
208 let mut dst_map = dst_u8.map_mut()?;
209 dst_map.as_mut_slice().fill(0);
210 }
211 }
212 Ok(())
213}
214
215fn chroma_alignment_ok(fmt: PixelFormat, r: Rect) -> bool {
223 match fmt {
224 PixelFormat::Nv12 => r.left.is_multiple_of(2) && r.top.is_multiple_of(2),
225 PixelFormat::Nv16 => r.left.is_multiple_of(2),
226 PixelFormat::Nv24 => true,
227 _ => false,
228 }
229}
230
231fn row_stride_for(width: usize, fmt: PixelFormat) -> usize {
233 use edgefirst_tensor::PixelLayout;
234 match fmt.layout() {
235 PixelLayout::Packed => width * fmt.channels(),
236 PixelLayout::Planar | PixelLayout::SemiPlanar => width,
237 _ => width, }
239}
240
241fn tensor_row_stride(tensor: &Tensor<u8>) -> usize {
246 tensor.effective_row_stride().unwrap_or_else(|| {
247 let w = tensor.width().unwrap_or(0);
248 let fmt = tensor.format().unwrap_or(PixelFormat::Rgb);
249 row_stride_for(w, fmt)
250 })
251}
252
253fn logical_surface(tensor: &Tensor<u8>) -> Result<(usize, usize)> {
263 use edgefirst_tensor::PixelLayout;
264 let fmt = tensor.format().ok_or(Error::NotAnImage)?;
265 let w = tensor.width().ok_or(Error::NotAnImage)?;
266 let h = tensor.height().ok_or(Error::NotAnImage)?;
267 Ok(match fmt.layout() {
268 PixelLayout::Packed => (h, w * fmt.channels()),
269 PixelLayout::Planar => (fmt.channels() * h, w),
270 PixelLayout::SemiPlanar => (fmt.combined_plane_height(h).unwrap_or(h), w),
271 _ => (h, row_stride_for(w, fmt)),
274 })
275}
276
277fn packed_row_pairs<'s, 'd>(
283 src: &'s [u8],
284 src_stride: usize,
285 src_row_bytes: usize,
286 dst: &'d mut [u8],
287 dst_stride: usize,
288 dst_row_bytes: usize,
289 rows: usize,
290) -> impl Iterator<Item = (&'s [u8], &'d mut [u8])> {
291 src.chunks(src_stride)
292 .zip(dst.chunks_mut(dst_stride))
293 .take(rows)
294 .map(move |(s, d)| (&s[..src_row_bytes], &mut d[..dst_row_bytes]))
295}
296
297fn split_semi_planar(
303 bytes: &[u8],
304 stride: usize,
305 src_h: usize,
306 fmt: PixelFormat,
307) -> Result<(&[u8], &[u8])> {
308 let total_h = fmt.combined_plane_height(src_h).unwrap_or(src_h);
309 let need = stride.checked_mul(total_h).ok_or_else(|| {
310 Error::InvalidShape(format!(
311 "{fmt:?} plane size overflow (stride={stride}, h={src_h})"
312 ))
313 })?;
314 if bytes.len() < need {
315 return Err(Error::InvalidShape(format!(
316 "{fmt:?} source has {} bytes but needs {need} (stride={stride}, h={src_h})",
317 bytes.len()
318 )));
319 }
320 Ok(bytes.split_at(stride * src_h))
321}
322
323fn split_semi_planar_mut(
330 bytes: &mut [u8],
331 stride: usize,
332 dst_h: usize,
333 fmt: PixelFormat,
334) -> Result<(&mut [u8], &mut [u8])> {
335 let total_h = fmt.combined_plane_height(dst_h).unwrap_or(dst_h);
336 let need = stride.checked_mul(total_h).ok_or_else(|| {
337 Error::InvalidShape(format!(
338 "{fmt:?} plane size overflow (stride={stride}, combined_h={total_h})"
339 ))
340 })?;
341 if bytes.len() < need {
342 return Err(Error::InvalidShape(format!(
343 "{fmt:?} destination has {} bytes but needs {need} (stride={stride}, combined_h={total_h})",
344 bytes.len()
345 )));
346 }
347 Ok(bytes.split_at_mut(stride * dst_h))
348}
349
350fn guard_plane(
356 buf_len: usize,
357 stride: usize,
358 rows: usize,
359 row_bytes: usize,
360 what: &str,
361) -> Result<()> {
362 let need = stride.checked_mul(rows).ok_or_else(|| {
363 Error::InvalidShape(format!(
364 "{what} plane size overflow (stride={stride}, rows={rows})"
365 ))
366 })?;
367 if row_bytes > stride || buf_len < need {
368 return Err(Error::InvalidShape(format!(
369 "{what} buffer too small: {buf_len} bytes, need {need} (stride={stride}, rows={rows}, row_bytes={row_bytes})"
370 )));
371 }
372 Ok(())
373}
374
375pub(crate) fn apply_int8_xor_bias(data: &mut [u8], fmt: PixelFormat) {
380 use edgefirst_tensor::PixelLayout;
381 if !fmt.has_alpha() {
382 for b in data.iter_mut() {
383 *b ^= 0x80;
384 }
385 } else if fmt.layout() == PixelLayout::Planar {
386 let channels = fmt.channels();
388 let plane_size = data.len() / channels;
389 for b in data[..plane_size * (channels - 1)].iter_mut() {
390 *b ^= 0x80;
391 }
392 } else {
393 let channels = fmt.channels();
395 for pixel in data.chunks_exact_mut(channels) {
396 for b in &mut pixel[..channels - 1] {
397 *b ^= 0x80;
398 }
399 }
400 }
401}
402
403pub(crate) fn apply_int8_xor_bias_rows(tensor: &mut Tensor<u8>, fmt: PixelFormat) -> Result<()> {
409 use edgefirst_tensor::PixelLayout;
410 let (rows, row_bytes) = logical_surface(tensor)?;
411 let stride = tensor_row_stride(tensor);
412 let color_rows = if fmt.has_alpha() && fmt.layout() == PixelLayout::Planar {
415 rows / fmt.channels() * (fmt.channels() - 1)
416 } else {
417 rows
418 };
419 let row_fmt = if fmt.layout() == PixelLayout::Planar {
422 PixelFormat::Grey
423 } else {
424 fmt
425 };
426 let mut map = tensor.map_mut()?;
427 let buf = map.as_mut_slice();
428 guard_plane(buf.len(), stride, rows, row_bytes, "int8 bias dst")?;
429 for row in buf.chunks_mut(stride).take(color_rows) {
430 apply_int8_xor_bias(&mut row[..row_bytes], row_fmt);
431 }
432 Ok(())
433}
434
435impl CPUProcessor {
436 pub fn new() -> Self {
438 Self::new_bilinear()
439 }
440
441 fn new_bilinear() -> Self {
443 let resizer = fast_image_resize::Resizer::new();
444 let options = fast_image_resize::ResizeOptions::new()
445 .resize_alg(fast_image_resize::ResizeAlg::Convolution(
446 fast_image_resize::FilterType::Bilinear,
447 ))
448 .use_alpha(false);
449
450 log::debug!("CPUConverter created");
451 Self {
452 resizer,
453 options,
454 colors: crate::DEFAULT_COLORS_U8,
455 widen_scratch: None,
456 resize_destride_scratch: Vec::new(),
457 resize_dst_destride_scratch: Vec::new(),
458 nv_strip_scratch: Vec::new(),
459 nv_strip_y_pack: Vec::new(),
460 nv_strip_uv_pack: Vec::new(),
461 #[cfg(test)]
462 fused_hits: 0,
463 #[cfg(test)]
464 last_tmp_dims: None,
465 convert_tmp: None,
466 convert_tmp2: None,
467 convert_src_sub: None,
468 }
469 }
470
471 pub fn set_interpolation(&mut self, mode: crate::Interpolation) {
474 let alg = match mode {
475 crate::Interpolation::Nearest => fast_image_resize::ResizeAlg::Nearest,
476 crate::Interpolation::Bilinear => {
477 fast_image_resize::ResizeAlg::Convolution(fast_image_resize::FilterType::Bilinear)
478 }
479 };
480 self.options = self.options.clone().resize_alg(alg);
481 }
482
483 pub fn new_nearest() -> Self {
485 let resizer = fast_image_resize::Resizer::new();
486 let options = fast_image_resize::ResizeOptions::new()
487 .resize_alg(fast_image_resize::ResizeAlg::Nearest)
488 .use_alpha(false);
489 log::debug!("CPUConverter created");
490 Self {
491 resizer,
492 options,
493 colors: crate::DEFAULT_COLORS_U8,
494 widen_scratch: None,
495 resize_destride_scratch: Vec::new(),
496 resize_dst_destride_scratch: Vec::new(),
497 nv_strip_scratch: Vec::new(),
498 nv_strip_y_pack: Vec::new(),
499 nv_strip_uv_pack: Vec::new(),
500 #[cfg(test)]
501 fused_hits: 0,
502 #[cfg(test)]
503 last_tmp_dims: None,
504 convert_tmp: None,
505 convert_tmp2: None,
506 convert_src_sub: None,
507 }
508 }
509
510 #[cfg(test)]
512 pub(super) fn fused_hits(&self) -> u64 {
513 self.fused_hits
514 }
515
516 #[cfg(test)]
519 pub(super) fn last_tmp_dims(&self) -> Option<(usize, usize)> {
520 self.last_tmp_dims
521 }
522
523 pub(crate) fn support_conversion_pf(src: PixelFormat, dst: PixelFormat) -> bool {
524 use PixelFormat::*;
525 matches!(
526 (src, dst),
527 (Nv12, Rgb)
528 | (Nv12, Rgba)
529 | (Nv12, Grey)
530 | (Nv16, Rgb)
531 | (Nv16, Rgba)
532 | (Nv16, Bgra)
533 | (Nv24, Rgb)
534 | (Nv24, Rgba)
535 | (Nv24, Grey)
536 | (Nv24, Bgra)
537 | (Yuyv, Rgb)
538 | (Yuyv, Rgba)
539 | (Yuyv, Grey)
540 | (Yuyv, Yuyv)
541 | (Yuyv, PlanarRgb)
542 | (Yuyv, PlanarRgba)
543 | (Yuyv, Nv16)
544 | (Vyuy, Rgb)
545 | (Vyuy, Rgba)
546 | (Vyuy, Grey)
547 | (Vyuy, Vyuy)
548 | (Vyuy, PlanarRgb)
549 | (Vyuy, PlanarRgba)
550 | (Vyuy, Nv16)
551 | (Rgba, Rgb)
552 | (Rgba, Rgba)
553 | (Rgba, Grey)
554 | (Rgba, Yuyv)
555 | (Rgba, PlanarRgb)
556 | (Rgba, PlanarRgba)
557 | (Rgba, Nv16)
558 | (Rgb, Rgb)
559 | (Rgb, Rgba)
560 | (Rgb, Grey)
561 | (Rgb, Yuyv)
562 | (Rgb, PlanarRgb)
563 | (Rgb, PlanarRgba)
564 | (Rgb, Nv16)
565 | (Grey, Rgb)
566 | (Grey, Rgba)
567 | (Grey, Grey)
568 | (Grey, Yuyv)
569 | (Grey, PlanarRgb)
570 | (Grey, PlanarRgba)
571 | (Grey, Nv16)
572 | (Nv12, Bgra)
573 | (Yuyv, Bgra)
574 | (Vyuy, Bgra)
575 | (Rgba, Bgra)
576 | (Rgb, Bgra)
577 | (Grey, Bgra)
578 | (Bgra, Bgra)
579 | (PlanarRgb, Rgb)
580 | (PlanarRgb, Rgba)
581 | (PlanarRgba, Rgb)
582 | (PlanarRgba, Rgba)
583 | (PlanarRgb, Bgra)
584 | (PlanarRgba, Bgra)
585 )
586 }
587
588 pub(crate) fn convert_format_pf(
590 src: &Tensor<u8>,
591 dst: &mut Tensor<u8>,
592 src_fmt: PixelFormat,
593 dst_fmt: PixelFormat,
594 cp: ColorParams,
595 ) -> Result<()> {
596 let _timer = FunctionTimer::new(format!(
597 "ImageProcessor::convert_format {} to {}",
598 src_fmt, dst_fmt,
599 ));
600
601 use PixelFormat::*;
602 match (src_fmt, dst_fmt) {
603 (Nv12, Rgb) => Self::convert_nv12_to_rgb(src, dst, cp),
604 (Nv12, Rgba) => Self::convert_nv12_to_rgba(src, dst, cp),
605 (Nv12, Grey) => Self::convert_nv12_to_grey(src, dst, cp),
606 (Yuyv, Rgb) => Self::convert_yuyv_to_rgb(src, dst, cp),
607 (Yuyv, Rgba) => Self::convert_yuyv_to_rgba(src, dst, cp),
608 (Yuyv, Grey) => Self::convert_yuyv_to_grey(src, dst, cp),
609 (Yuyv, Yuyv) => Self::copy_image(src, dst),
610 (Yuyv, PlanarRgb) => Self::convert_yuyv_to_8bps(src, dst, cp),
611 (Yuyv, PlanarRgba) => Self::convert_yuyv_to_prgba(src, dst, cp),
612 (Yuyv, Nv16) => Self::convert_yuyv_to_nv16(src, dst),
613 (Vyuy, Rgb) => Self::convert_vyuy_to_rgb(src, dst, cp),
614 (Vyuy, Rgba) => Self::convert_vyuy_to_rgba(src, dst, cp),
615 (Vyuy, Grey) => Self::convert_vyuy_to_grey(src, dst, cp),
616 (Vyuy, Vyuy) => Self::copy_image(src, dst),
617 (Vyuy, PlanarRgb) => Self::convert_vyuy_to_8bps(src, dst, cp),
618 (Vyuy, PlanarRgba) => Self::convert_vyuy_to_prgba(src, dst, cp),
619 (Vyuy, Nv16) => Self::convert_vyuy_to_nv16(src, dst),
620 (Rgba, Rgb) => Self::convert_rgba_to_rgb(src, dst),
621 (Rgba, Rgba) => Self::copy_image(src, dst),
622 (Rgba, Grey) => Self::convert_rgba_to_grey(src, dst),
623 (Rgba, Yuyv) => Self::convert_rgba_to_yuyv(src, dst, cp),
624 (Rgba, PlanarRgb) => Self::convert_rgba_to_8bps(src, dst),
625 (Rgba, PlanarRgba) => Self::convert_rgba_to_prgba(src, dst),
626 (Rgba, Nv16) => Self::convert_rgba_to_nv16(src, dst, cp),
627 (Rgb, Rgb) => Self::copy_image(src, dst),
628 (Rgb, Rgba) => Self::convert_rgb_to_rgba(src, dst),
629 (Rgb, Grey) => Self::convert_rgb_to_grey(src, dst),
630 (Rgb, Yuyv) => Self::convert_rgb_to_yuyv(src, dst, cp),
631 (Rgb, PlanarRgb) => Self::convert_rgb_to_8bps(src, dst),
632 (Rgb, PlanarRgba) => Self::convert_rgb_to_prgba(src, dst),
633 (Rgb, Nv16) => Self::convert_rgb_to_nv16(src, dst, cp),
634 (Grey, Rgb) => Self::convert_grey_to_rgb(src, dst),
635 (Grey, Rgba) => Self::convert_grey_to_rgba(src, dst),
636 (Grey, Grey) => Self::copy_image(src, dst),
637 (Grey, Yuyv) => Self::convert_grey_to_yuyv(src, dst, cp),
638 (Grey, PlanarRgb) => Self::convert_grey_to_8bps(src, dst),
639 (Grey, PlanarRgba) => Self::convert_grey_to_prgba(src, dst),
640 (Grey, Nv16) => Self::convert_grey_to_nv16(src, dst, cp),
641
642 (Nv16, Rgb) => Self::convert_nv16_to_rgb(src, dst, cp),
644 (Nv16, Rgba) => Self::convert_nv16_to_rgba(src, dst, cp),
645 (Nv24, Rgb) => Self::convert_nv24_to_rgb(src, dst, cp),
646 (Nv24, Rgba) => Self::convert_nv24_to_rgba(src, dst, cp),
647 (Nv24, Grey) => Self::convert_nv24_to_grey(src, dst, cp),
648 (PlanarRgb, Rgb) => Self::convert_8bps_to_rgb(src, dst),
649 (PlanarRgb, Rgba) => Self::convert_8bps_to_rgba(src, dst),
650 (PlanarRgba, Rgb) => Self::convert_prgba_to_rgb(src, dst),
651 (PlanarRgba, Rgba) => Self::convert_prgba_to_rgba(src, dst),
652
653 (Bgra, Bgra) => Self::copy_image(src, dst),
655 (Nv12, Bgra) => {
656 Self::convert_nv12_to_rgba(src, dst, cp)?;
657 Self::swizzle_rb_4chan(dst)
658 }
659 (Nv16, Bgra) => {
660 Self::convert_nv16_to_rgba(src, dst, cp)?;
661 Self::swizzle_rb_4chan(dst)
662 }
663 (Nv24, Bgra) => {
664 Self::convert_nv24_to_rgba(src, dst, cp)?;
665 Self::swizzle_rb_4chan(dst)
666 }
667 (Yuyv, Bgra) => {
668 Self::convert_yuyv_to_rgba(src, dst, cp)?;
669 Self::swizzle_rb_4chan(dst)
670 }
671 (Vyuy, Bgra) => {
672 Self::convert_vyuy_to_rgba(src, dst, cp)?;
673 Self::swizzle_rb_4chan(dst)
674 }
675 (Rgba, Bgra) => {
676 Self::copy_image(src, dst)?;
677 Self::swizzle_rb_4chan(dst)
678 }
679 (Rgb, Bgra) => {
680 Self::convert_rgb_to_rgba(src, dst)?;
681 Self::swizzle_rb_4chan(dst)
682 }
683 (Grey, Bgra) => {
684 Self::convert_grey_to_rgba(src, dst)?;
685 Self::swizzle_rb_4chan(dst)
686 }
687 (PlanarRgb, Bgra) => {
688 Self::convert_8bps_to_rgba(src, dst)?;
689 Self::swizzle_rb_4chan(dst)
690 }
691 (PlanarRgba, Bgra) => {
692 Self::convert_prgba_to_rgba(src, dst)?;
693 Self::swizzle_rb_4chan(dst)
694 }
695
696 (s, d) => Err(Error::NotSupported(format!("Conversion from {s} to {d}",))),
697 }
698 }
699
700 pub(crate) fn fill_image_outside_crop_u8(
702 dst: &mut Tensor<u8>,
703 rgba: [u8; 4],
704 crop: Rect,
705 ) -> Result<()> {
706 let dst_fmt = dst.format().unwrap();
707 let dst_w = dst.width().unwrap();
708 let dst_h = dst.height().unwrap();
709 let cm = crate::colorimetry::resolve_colorimetry(dst.colorimetry(), dst.height());
713 let cp = ColorParams {
714 matrix: crate::colorimetry::yuv_matrix(cm.encoding.unwrap()),
715 range: crate::colorimetry::yuv_range(cm.range.unwrap()),
716 encoding: cm.encoding.unwrap(),
717 range_kind: cm.range.unwrap(),
718 src_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
719 dst_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
720 };
721 let dst_stride = tensor_row_stride(dst);
722 let mut dst_map = dst.map_mut()?;
723 let dst_tup = (dst_map.as_mut_slice(), dst_w, dst_h, dst_stride);
724 Self::fill_outside_crop_dispatch(dst_tup, dst_fmt, rgba, crop, cp)
725 }
726
727 fn fill_outside_crop_dispatch(
732 dst: (&mut [u8], usize, usize, usize),
733 fmt: PixelFormat,
734 rgba: [u8; 4],
735 crop: Rect,
736 cp: ColorParams,
737 ) -> Result<()> {
738 use PixelFormat::*;
739 match fmt {
740 Rgba | Bgra => Self::fill_image_outside_crop_(dst, rgba, crop),
741 Rgb => Self::fill_image_outside_crop_(dst, Self::rgba_to_rgb(rgba), crop),
742 Grey => Self::fill_image_outside_crop_(dst, Self::rgba_to_grey(rgba), crop),
743 Yuyv => {
744 let (bytes, w, h, stride) = dst;
745 let yuyv = Self::rgba_to_yuyv(rgba, cp);
746 Self::fill_image_outside_crop_(
749 (&mut *bytes, w / 2, h, stride),
750 yuyv,
751 Rect::new(crop.left / 2, crop.top, crop.width.div_ceil(2), crop.height),
752 )?;
753 if w % 2 == 1 {
761 let last = w - 1;
762 let col_outside = last < crop.left || last >= crop.left + crop.width;
763 for y in 0..h {
764 let row_outside = y < crop.top || y >= crop.top + crop.height;
765 if row_outside || col_outside {
766 let off = y * stride + last * 2;
767 bytes[off] = yuyv[0];
768 bytes[off + 1] = yuyv[1];
769 }
770 }
771 }
772 Ok(())
773 }
774 PlanarRgb => Self::fill_image_outside_crop_planar(dst, Self::rgba_to_rgb(rgba), crop),
775 PlanarRgba => Self::fill_image_outside_crop_planar(dst, rgba, crop),
776 Nv16 => {
777 let yuyv = Self::rgba_to_yuyv(rgba, cp);
778 Self::fill_image_outside_crop_yuv_semiplanar(dst, yuyv[0], [yuyv[1], yuyv[3]], crop)
779 }
780 _ => Err(Error::Internal(format!(
781 "Found unexpected destination {fmt}",
782 ))),
783 }
784 }
785}
786
787impl ImageProcessorTrait for CPUProcessor {
788 fn convert(
789 &mut self,
790 src: &TensorDyn,
791 dst: &mut TensorDyn,
792 rotation: Rotation,
793 flip: Flip,
794 crop: Crop,
795 ) -> Result<()> {
796 let crop = crop.resolve(
797 src.width().unwrap_or(0),
798 src.height().unwrap_or(0),
799 dst.width().unwrap_or(0),
800 dst.height().unwrap_or(0),
801 )?;
802 self.convert_impl(src, dst, rotation, flip, crop)
803 }
804
805 fn draw_decoded_masks(
806 &mut self,
807 dst: &mut TensorDyn,
808 detect: &[DetectBox],
809 segmentation: &[Segmentation],
810 overlay: crate::MaskOverlay<'_>,
811 ) -> Result<()> {
812 prepare_dst_base_cpu(dst, overlay.background)?;
816 let dst = dst.as_typed_mut::<u8>().ok_or(Error::NotAnImage)?;
817 self.draw_decoded_masks_impl(
818 dst,
819 detect,
820 segmentation,
821 overlay.opacity,
822 overlay.color_mode,
823 )
824 }
825
826 fn draw_proto_masks(
827 &mut self,
828 dst: &mut TensorDyn,
829 detect: &[DetectBox],
830 proto_data: &ProtoData,
831 overlay: crate::MaskOverlay<'_>,
832 ) -> Result<()> {
833 prepare_dst_base_cpu(dst, overlay.background)?;
834 let dst = dst.as_typed_mut::<u8>().ok_or(Error::NotAnImage)?;
835 self.draw_proto_masks_impl(
836 dst,
837 detect,
838 proto_data,
839 overlay.opacity,
840 overlay.letterbox,
841 overlay.color_mode,
842 )
843 }
844
845 fn set_class_colors(&mut self, colors: &[[u8; 4]]) -> Result<()> {
846 for (c, new_c) in self.colors.iter_mut().zip(colors.iter()) {
847 *c = *new_c;
848 }
849 Ok(())
850 }
851}
852
853impl CPUProcessor {
855 pub(crate) fn convert_impl(
857 &mut self,
858 src: &TensorDyn,
859 dst: &mut TensorDyn,
860 rotation: Rotation,
861 flip: Flip,
862 crop: ResolvedCrop,
863 ) -> Result<()> {
864 let src_fmt = src.format().ok_or(Error::NotAnImage)?;
865 let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
866
867 let src_cm = crate::colorimetry::effective_colorimetry(src);
872 let dst_cm = crate::colorimetry::effective_colorimetry(dst);
873 let src_full = src_cm.range == Some(edgefirst_tensor::ColorRange::Full);
874 let dst_full = dst_cm.range == Some(edgefirst_tensor::ColorRange::Full);
875 let src_params = ColorParams {
876 matrix: crate::colorimetry::yuv_matrix(src_cm.encoding.unwrap()),
877 range: crate::colorimetry::yuv_range(src_cm.range.unwrap()),
878 encoding: src_cm.encoding.unwrap(),
879 range_kind: src_cm.range.unwrap(),
880 src_full_range: src_full,
881 dst_full_range: dst_full,
882 };
883 let dst_params = ColorParams {
884 matrix: crate::colorimetry::yuv_matrix(dst_cm.encoding.unwrap()),
885 range: crate::colorimetry::yuv_range(dst_cm.range.unwrap()),
886 encoding: dst_cm.encoding.unwrap(),
887 range_kind: dst_cm.range.unwrap(),
888 src_full_range: src_full,
889 dst_full_range: dst_full,
890 };
891 match (src.dtype(), dst.dtype()) {
892 (DType::U8, DType::U8) => {
893 let src = src.as_typed::<u8>().unwrap();
894 let dst = dst.as_typed_mut::<u8>().unwrap();
895 self.convert_u8(
896 src, dst, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
897 )
898 }
899 (DType::U8, DType::I8) => {
900 let src_u8 = src.as_typed::<u8>().unwrap();
903 let dst_i8 = dst.as_typed_mut::<i8>().unwrap();
904 let dst_u8 = unsafe { &mut *(dst_i8 as *mut Tensor<i8> as *mut Tensor<u8>) };
908 self.convert_u8(
909 src_u8, dst_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
910 )?;
911 apply_int8_xor_bias_rows(dst_u8, dst_fmt)
913 }
914 (DType::U8, d @ (DType::F32 | DType::F16)) => {
915 let src_u8 = src.as_typed::<u8>().unwrap();
916 let dw = dst.width().ok_or(Error::NotAnImage)?;
917 let dh = dst.height().ok_or(Error::NotAnImage)?;
918 let scratch_matches = self.widen_scratch.as_ref().is_some_and(|t| {
923 t.width() == Some(dw) && t.height() == Some(dh) && t.format() == Some(dst_fmt)
924 });
925 let mut tmp = if scratch_matches {
926 self.widen_scratch.take().unwrap()
927 } else {
928 TensorDyn::image(
929 dw,
930 dh,
931 dst_fmt,
932 DType::U8,
933 Some(TensorMemory::Mem),
934 edgefirst_tensor::CpuAccess::ReadWrite,
935 )?
936 };
937 {
938 let tmp_u8 = tmp.as_typed_mut::<u8>().unwrap();
939 self.convert_u8(
940 src_u8, tmp_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params,
941 dst_params,
942 )?;
943 }
944 {
954 let tmp_u8 = tmp.as_typed::<u8>().unwrap();
955 let (rows, row_len) = logical_surface(tmp_u8)?;
956 let src_stride = tensor_row_stride(tmp_u8);
957 let dst_stride_bytes = dst.effective_row_stride().ok_or(Error::NotAnImage)?;
958 let src_map = tmp_u8.map_read()?;
959 guard_plane(
960 src_map.as_slice().len(),
961 src_stride,
962 rows,
963 row_len,
964 "widen src",
965 )?;
966 let src_rows = src_map.as_slice().chunks(src_stride).take(rows);
967 let elem = d.size();
968 if !dst_stride_bytes.is_multiple_of(elem) {
969 return Err(Error::InvalidShape(format!(
970 "{d} destination row stride {dst_stride_bytes} is not a multiple of \
971 the element size {elem}"
972 )));
973 }
974 let dst_stride = dst_stride_bytes / elem;
975 match d {
976 DType::F32 => {
977 let dst_t = dst.as_typed_mut::<f32>().unwrap();
978 let mut dst_map = dst_t.map_mut()?;
979 guard_plane(
981 dst_map.as_slice().len(),
982 dst_stride,
983 rows,
984 row_len,
985 "widen f32 dst",
986 )?;
987 for (s, dr) in
991 src_rows.zip(dst_map.as_mut_slice().chunks_mut(dst_stride))
992 {
993 simd::widen_u8_to_f32_norm(s, &mut dr[..row_len]);
994 }
995 }
996 DType::F16 => {
997 let dst_t = dst.as_typed_mut::<half::f16>().unwrap();
998 let mut dst_map = dst_t.map_mut()?;
999 guard_plane(
1000 dst_map.as_slice().len(),
1001 dst_stride,
1002 rows,
1003 row_len,
1004 "widen f16 dst",
1005 )?;
1006 for (s, dr) in
1011 src_rows.zip(dst_map.as_mut_slice().chunks_mut(dst_stride))
1012 {
1013 simd::widen_u8_to_f16_norm(s, &mut dr[..row_len]);
1014 }
1015 }
1016 _ => unreachable!(),
1017 }
1018 }
1019 self.widen_scratch = Some(tmp);
1020 Ok(())
1021 }
1022 (s, d) => Err(Error::NotSupported(format!("dtype {s} -> {d}",))),
1023 }
1024 }
1025
1026 fn reuse_or_alloc_image(
1032 cached: Option<Tensor<u8>>,
1033 w: usize,
1034 h: usize,
1035 fmt: PixelFormat,
1036 ) -> Result<Tensor<u8>> {
1037 if let Some(t) = cached {
1038 if t.width() == Some(w) && t.height() == Some(h) && t.format() == Some(fmt) {
1039 return Ok(t);
1040 }
1041 }
1042 Ok(Tensor::<u8>::image(
1043 w,
1044 h,
1045 fmt,
1046 Some(TensorMemory::Mem),
1047 edgefirst_tensor::CpuAccess::ReadWrite,
1048 )?)
1049 }
1050
1051 fn pre_resize_region(
1072 &self,
1073 src_fmt: PixelFormat,
1074 (src_w, src_h): (usize, usize),
1075 (dst_w, dst_h): (usize, usize),
1076 rotation: Rotation,
1077 crop: ResolvedCrop,
1078 ) -> Option<Rect> {
1079 use PixelFormat::{Nv12, Nv16, Nv24};
1080
1081 if !matches!(src_fmt, Nv12 | Nv16 | Nv24) {
1082 return None;
1083 }
1084 let r = crop.src_rect?;
1085 let full_src = Rect {
1086 left: 0,
1087 top: 0,
1088 width: src_w,
1089 height: src_h,
1090 };
1091 if r == full_src {
1092 return None;
1093 }
1094
1095 let d = crop.dst_rect.unwrap_or(Rect {
1100 left: 0,
1101 top: 0,
1102 width: dst_w,
1103 height: dst_h,
1104 });
1105 let (dst_x, dst_y) = match rotation {
1106 Rotation::None | Rotation::Rotate180 => (d.width, d.height),
1107 Rotation::Clockwise90 | Rotation::CounterClockwise90 => (d.height, d.width),
1108 };
1109 let halo_x = self.filter_halo(r.width, dst_x)?;
1110 let halo_y = self.filter_halo(r.height, dst_y)?;
1111
1112 let mut left = r.left.saturating_sub(halo_x);
1113 let mut top = r.top.saturating_sub(halo_y);
1114 let mut right = (r.left + r.width + halo_x).min(src_w);
1115 let mut bottom = (r.top + r.height + halo_y).min(src_h);
1116
1117 let (align_x, align_y) = match src_fmt {
1123 Nv12 => (2, 2),
1124 Nv16 => (2, 1),
1125 _ => (1, 1),
1126 };
1127 left -= left % align_x;
1128 top -= top % align_y;
1129 right = right.next_multiple_of(align_x).min(src_w);
1130 bottom = bottom.next_multiple_of(align_y).min(src_h);
1131
1132 let grown = Rect {
1133 left,
1134 top,
1135 width: right - left,
1136 height: bottom - top,
1137 };
1138 (grown != full_src).then_some(grown)
1141 }
1142
1143 #[allow(clippy::too_many_arguments)]
1145 fn convert_u8(
1146 &mut self,
1147 src: &Tensor<u8>,
1148 dst: &mut Tensor<u8>,
1149 src_fmt: PixelFormat,
1150 dst_fmt: PixelFormat,
1151 rotation: Rotation,
1152 flip: Flip,
1153 crop: ResolvedCrop,
1154 src_params: ColorParams,
1155 dst_params: ColorParams,
1156 ) -> Result<()> {
1157 use PixelFormat::*;
1158
1159 #[cfg(test)]
1160 {
1161 self.last_tmp_dims = None;
1162 }
1163
1164 let src_w = src.width().unwrap();
1165 let src_h = src.height().unwrap();
1166 let dst_w = dst.width().unwrap();
1167 let dst_h = dst.height().unwrap();
1168
1169 crop.check_crop_dims(src_w, src_h, dst_w, dst_h)?;
1170
1171 let intermediate = match (src_fmt, dst_fmt) {
1173 (Nv12, Rgb) => Rgb,
1174 (Nv12, Rgba) => Rgba,
1175 (Nv12, Grey) => Grey,
1176 (Nv12, Yuyv) => Rgba,
1177 (Nv12, Nv16) => Rgba,
1178 (Nv12, PlanarRgb) => Rgb,
1179 (Nv12, PlanarRgba) => Rgba,
1180 (Nv16, PlanarRgb) => Rgb,
1181 (Nv16, PlanarRgba) => Rgba,
1182 (Nv24, PlanarRgb) => Rgb,
1183 (Nv24, PlanarRgba) => Rgba,
1184 (Yuyv, Rgb) => Rgb,
1185 (Yuyv, Rgba) => Rgba,
1186 (Yuyv, Grey) => Grey,
1187 (Yuyv, Yuyv) => Rgba,
1188 (Yuyv, PlanarRgb) => Rgb,
1189 (Yuyv, PlanarRgba) => Rgba,
1190 (Yuyv, Nv16) => Rgba,
1191 (Vyuy, Rgb) => Rgb,
1192 (Vyuy, Rgba) => Rgba,
1193 (Vyuy, Grey) => Grey,
1194 (Vyuy, Vyuy) => Rgba,
1195 (Vyuy, PlanarRgb) => Rgb,
1196 (Vyuy, PlanarRgba) => Rgba,
1197 (Vyuy, Nv16) => Rgba,
1198 (Rgba, Rgb) => Rgba,
1199 (Rgba, Rgba) => Rgba,
1200 (Rgba, Grey) => Grey,
1201 (Rgba, Yuyv) => Rgba,
1202 (Rgba, PlanarRgb) => Rgba,
1203 (Rgba, PlanarRgba) => Rgba,
1204 (Rgba, Nv16) => Rgba,
1205 (Rgb, Rgb) => Rgb,
1206 (Rgb, Rgba) => Rgb,
1207 (Rgb, Grey) => Grey,
1208 (Rgb, Yuyv) => Rgb,
1209 (Rgb, PlanarRgb) => Rgb,
1210 (Rgb, PlanarRgba) => Rgb,
1211 (Rgb, Nv16) => Rgb,
1212 (Grey, Rgb) => Rgb,
1213 (Grey, Rgba) => Rgba,
1214 (Grey, Grey) => Grey,
1215 (Grey, Yuyv) => Grey,
1216 (Grey, PlanarRgb) => Grey,
1217 (Grey, PlanarRgba) => Grey,
1218 (Grey, Nv16) => Grey,
1219 (Nv12, Bgra) => Rgba,
1220 (Yuyv, Bgra) => Rgba,
1221 (Vyuy, Bgra) => Rgba,
1222 (Rgba, Bgra) => Rgba,
1223 (Rgb, Bgra) => Rgb,
1224 (Grey, Bgra) => Grey,
1225 (Bgra, Bgra) => Bgra,
1226 (Nv16, Rgb) => Rgb,
1227 (Nv16, Rgba) => Rgba,
1228 (Nv16, Bgra) => Rgba,
1229 (Nv24, Rgb) => Rgb,
1230 (Nv24, Rgba) => Rgba,
1231 (Nv24, Grey) => Grey,
1232 (Nv24, Bgra) => Rgba,
1233 (PlanarRgb, Rgb) => Rgb,
1234 (PlanarRgb, Rgba) => Rgb,
1235 (PlanarRgb, Bgra) => Rgb,
1236 (PlanarRgba, Rgb) => Rgba,
1237 (PlanarRgba, Rgba) => Rgba,
1238 (PlanarRgba, Bgra) => Rgba,
1239 (s, d) => {
1240 return Err(Error::NotSupported(format!("Conversion from {s} to {d}",)));
1241 }
1242 };
1243
1244 let need_resize_flip_rotation = rotation != Rotation::None
1245 || flip != Flip::None
1246 || src_w != dst_w
1247 || src_h != dst_h
1248 || crop.src_rect.is_some_and(|c| {
1249 c != Rect {
1250 left: 0,
1251 top: 0,
1252 width: src_w,
1253 height: src_h,
1254 }
1255 })
1256 || crop.dst_rect.is_some_and(|c| {
1257 c != Rect {
1258 left: 0,
1259 top: 0,
1260 width: dst_w,
1261 height: dst_h,
1262 }
1263 });
1264
1265 let direct_is_yuv_src = matches!(src_fmt, Nv12 | Nv16 | Nv24 | Yuyv | Vyuy);
1268 let direct_params = if direct_is_yuv_src {
1269 src_params
1270 } else {
1271 dst_params
1272 };
1273
1274 let full_dst_rect = Rect {
1291 left: 0,
1292 top: 0,
1293 width: dst_w,
1294 height: dst_h,
1295 };
1296 let fused_region = if rotation != Rotation::None || flip != Flip::None {
1297 None
1298 } else {
1299 match crop.src_rect {
1300 None if src_w == dst_w && src_h == dst_h => Some(None),
1301 None => None,
1302 Some(r)
1303 if r.width == dst_w
1304 && r.height == dst_h
1305 && crop.dst_rect.is_none_or(|d| d == full_dst_rect)
1306 && chroma_alignment_ok(src_fmt, r) =>
1307 {
1308 Some(Some(r))
1309 }
1310 Some(_) => None,
1311 }
1312 };
1313 if let Some(region) = fused_region {
1314 if matches!(src_fmt, Nv12 | Nv16 | Nv24) && matches!(dst_fmt, PlanarRgb | PlanarRgba) {
1315 #[cfg(test)]
1316 {
1317 self.fused_hits += 1;
1318 }
1319 return self.convert_nv_to_planar_fused(
1320 src,
1321 dst,
1322 src_fmt,
1323 dst_fmt,
1324 direct_params,
1325 region,
1326 );
1327 }
1328 }
1329
1330 if !need_resize_flip_rotation && Self::support_conversion_pf(src_fmt, dst_fmt) {
1332 return Self::convert_format_pf(src, dst, src_fmt, dst_fmt, direct_params);
1333 }
1334
1335 if dst_fmt == Yuyv && !dst_w.is_multiple_of(2) {
1337 return Err(Error::NotSupported(format!(
1338 "{} destination must have width divisible by 2",
1339 dst_fmt,
1340 )));
1341 }
1342
1343 let mut cached_tmp = self.convert_tmp.take();
1351 let mut cached_tmp2 = self.convert_tmp2.take();
1352
1353 let pre_region = if intermediate != src_fmt {
1358 self.pre_resize_region(src_fmt, (src_w, src_h), (dst_w, dst_h), rotation, crop)
1359 } else {
1360 None
1361 };
1362
1363 let tmp_holder: Option<Tensor<u8>> = if intermediate != src_fmt {
1365 let _s = tracing::trace_span!(
1366 "image.convert.cpu.format_convert",
1367 from = ?src_fmt,
1368 to = ?intermediate,
1369 pass = "pre_resize",
1370 )
1371 .entered();
1372 let (tmp_w, tmp_h) = pre_region.map_or((src_w, src_h), |g| (g.width, g.height));
1373 let mut t = Self::reuse_or_alloc_image(cached_tmp.take(), tmp_w, tmp_h, intermediate)?;
1374 #[cfg(test)]
1375 {
1376 self.last_tmp_dims = Some((tmp_w, tmp_h));
1377 }
1378 match pre_region {
1379 Some(g) => {
1380 let mut sub = Self::reuse_or_alloc_image(
1381 self.convert_src_sub.take(),
1382 g.width,
1383 g.height,
1384 src_fmt,
1385 )?;
1386 {
1387 let _s = tracing::trace_span!(
1388 "image.convert.cpu.extract_region",
1389 region_w = g.width,
1390 region_h = g.height,
1391 )
1392 .entered();
1393 Self::extract_nv_region(src, &mut sub, src_fmt, g)?;
1394 }
1395 Self::convert_format_pf(&sub, &mut t, src_fmt, intermediate, src_params)?;
1396 self.convert_src_sub = Some(sub);
1397 }
1398 None => Self::convert_format_pf(src, &mut t, src_fmt, intermediate, src_params)?,
1399 }
1400 Some(t)
1401 } else {
1402 None
1403 };
1404
1405 let crop = match (pre_region, crop.src_rect) {
1410 (Some(g), Some(r)) => ResolvedCrop {
1411 src_rect: Some(Rect {
1412 left: r.left - g.left,
1413 top: r.top - g.top,
1414 ..r
1415 }),
1416 ..crop
1417 },
1418 _ => crop,
1419 };
1420 let (tmp, tmp_fmt): (&Tensor<u8>, PixelFormat) = match &tmp_holder {
1421 Some(t) => (t, intermediate),
1422 None => (src, src_fmt),
1423 };
1424
1425 debug_assert!(matches!(tmp_fmt, Rgb | Rgba | Grey));
1427 if tmp_fmt == dst_fmt {
1428 let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
1429 self.resize_flip_rotate_pf(tmp, dst, dst_fmt, rotation, flip, crop)?;
1430 } else if !need_resize_flip_rotation {
1431 let _s = tracing::trace_span!(
1432 "image.convert.cpu.format_convert",
1433 from = ?tmp_fmt,
1434 to = ?dst_fmt,
1435 pass = "direct",
1436 )
1437 .entered();
1438 Self::convert_format_pf(tmp, dst, tmp_fmt, dst_fmt, dst_params)?;
1439 } else {
1440 let mut tmp2 = Self::reuse_or_alloc_image(cached_tmp2.take(), dst_w, dst_h, tmp_fmt)?;
1441 if crop.dst_rect.is_some_and(|c| {
1442 c != Rect {
1443 left: 0,
1444 top: 0,
1445 width: dst_w,
1446 height: dst_h,
1447 }
1448 }) && crop.dst_color.is_none()
1449 {
1450 Self::convert_format_pf(dst, &mut tmp2, dst_fmt, tmp_fmt, dst_params)?;
1451 }
1452 {
1453 let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
1454 self.resize_flip_rotate_pf(tmp, &mut tmp2, tmp_fmt, rotation, flip, crop)?;
1455 }
1456 {
1457 let _s = tracing::trace_span!(
1458 "image.convert.cpu.format_convert",
1459 from = ?tmp_fmt,
1460 to = ?dst_fmt,
1461 pass = "post_resize",
1462 )
1463 .entered();
1464 Self::convert_format_pf(&tmp2, dst, tmp_fmt, dst_fmt, dst_params)?;
1465 }
1466 cached_tmp2 = Some(tmp2);
1467 }
1468 if let Some(t) = tmp_holder {
1471 cached_tmp = Some(t);
1472 }
1473 self.convert_tmp = cached_tmp;
1474 self.convert_tmp2 = cached_tmp2;
1475
1476 if let (Some(dst_rect), Some(dst_color)) = (crop.dst_rect, crop.dst_color) {
1477 let full_rect = Rect {
1478 left: 0,
1479 top: 0,
1480 width: dst_w,
1481 height: dst_h,
1482 };
1483 if dst_rect != full_rect {
1484 Self::fill_image_outside_crop_u8(dst, dst_color, dst_rect)?;
1485 }
1486 }
1487
1488 Ok(())
1489 }
1490
1491 fn draw_decoded_masks_impl(
1492 &mut self,
1493 dst: &mut Tensor<u8>,
1494 detect: &[DetectBox],
1495 segmentation: &[Segmentation],
1496 opacity: f32,
1497 color_mode: crate::ColorMode,
1498 ) -> Result<()> {
1499 let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
1500 if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
1501 return Err(crate::Error::NotSupported(
1502 "CPU image rendering only supports RGBA or RGB images".to_string(),
1503 ));
1504 }
1505
1506 let _timer = FunctionTimer::new("CPUProcessor::draw_decoded_masks");
1507
1508 let dst_w = dst.width().unwrap();
1509 let dst_h = dst.height().unwrap();
1510 let dst_rs = tensor_row_stride(dst);
1511 let dst_c = dst_fmt.channels();
1512
1513 let mut map = dst.map_mut()?;
1514 let dst_slice = map.as_mut_slice();
1515
1516 self.render_box(dst_w, dst_h, dst_rs, dst_c, dst_slice, detect, color_mode)?;
1517
1518 if segmentation.is_empty() {
1519 return Ok(());
1520 }
1521
1522 if let Some(bad) = segmentation
1529 .iter()
1530 .find(|s| s.segmentation.shape().len() != 3)
1531 {
1532 return Err(crate::Error::InvalidShape(format!(
1533 "segmentation must be [H, W, C], got {:?}",
1534 bad.segmentation.shape()
1535 )));
1536 }
1537
1538 let is_semantic = segmentation[0].segmentation.shape()[2] > 1;
1541
1542 if is_semantic {
1543 self.render_modelpack_segmentation(
1544 dst_w,
1545 dst_h,
1546 dst_rs,
1547 dst_c,
1548 dst_slice,
1549 &segmentation[0],
1550 opacity,
1551 )?;
1552 } else {
1553 for (idx, (seg, det)) in segmentation.iter().zip(detect).enumerate() {
1554 let color_index = color_mode.index(idx, det.label);
1555 self.render_yolo_segmentation(
1556 dst_w,
1557 dst_h,
1558 dst_rs,
1559 dst_c,
1560 dst_slice,
1561 seg,
1562 color_index,
1563 opacity,
1564 )?;
1565 }
1566 }
1567
1568 Ok(())
1569 }
1570
1571 fn draw_proto_masks_impl(
1572 &mut self,
1573 dst: &mut Tensor<u8>,
1574 detect: &[DetectBox],
1575 proto_data: &ProtoData,
1576 opacity: f32,
1577 letterbox: Option<[f32; 4]>,
1578 color_mode: crate::ColorMode,
1579 ) -> Result<()> {
1580 let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
1581 if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
1582 return Err(crate::Error::NotSupported(
1583 "CPU image rendering only supports RGBA or RGB images".to_string(),
1584 ));
1585 }
1586
1587 let _timer = FunctionTimer::new("CPUProcessor::draw_proto_masks");
1588
1589 let dst_w = dst.width().unwrap();
1590 let dst_h = dst.height().unwrap();
1591 let dst_rs = tensor_row_stride(dst);
1592 let channels = dst_fmt.channels();
1593
1594 let mut map = dst.map_mut()?;
1595 let dst_slice = map.as_mut_slice();
1596
1597 self.render_box(
1598 dst_w, dst_h, dst_rs, channels, dst_slice, detect, color_mode,
1599 )?;
1600
1601 if detect.is_empty() {
1602 return Ok(());
1603 }
1604 let proto_shape = proto_data.protos.shape();
1605 if proto_shape.len() != 3 {
1606 return Err(Error::InvalidShape(format!(
1607 "protos tensor must be rank-3, got {proto_shape:?}"
1608 )));
1609 }
1610 let proto_h = proto_shape[0];
1611 let proto_w = proto_shape[1];
1612 let num_protos = proto_shape[2];
1613 let coeff_shape = proto_data.mask_coefficients.shape();
1614 if coeff_shape.len() != 2 {
1615 return Err(Error::InvalidShape(format!(
1616 "mask_coefficients tensor must be rank-2, got {coeff_shape:?}"
1617 )));
1618 }
1619 if coeff_shape[0] == 0 {
1621 return Ok(());
1622 }
1623 if coeff_shape[1] != num_protos {
1624 return Err(Error::InvalidShape(format!(
1625 "mask_coefficients second dimension must match num_protos \
1626 ({num_protos}), got {coeff_shape:?}"
1627 )));
1628 }
1629
1630 let coeff_f32: Vec<f32> = match proto_data.mask_coefficients.dtype() {
1632 DType::F32 => {
1633 let t = proto_data.mask_coefficients.as_typed::<f32>().expect("F32");
1634 let m = t.map_read()?;
1635 m.as_slice().to_vec()
1636 }
1637 DType::F16 => {
1638 let t = proto_data
1639 .mask_coefficients
1640 .as_typed::<half::f16>()
1641 .expect("F16");
1642 let m = t.map_read()?;
1643 m.as_slice().iter().map(|v| v.to_f32()).collect()
1644 }
1645 DType::I8 => {
1646 let t = proto_data.mask_coefficients.as_typed::<i8>().expect("I8");
1647 let m = t.map_read()?;
1648 if let Some(q) = t.quantization() {
1649 use edgefirst_tensor::QuantMode;
1650 let (scale, zp) = match q.mode() {
1651 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1652 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1653 other => {
1654 return Err(Error::NotSupported(format!(
1655 "I8 mask_coefficients quantization mode {other:?} not supported"
1656 )));
1657 }
1658 };
1659 m.as_slice()
1660 .iter()
1661 .map(|&v| (v as f32 - zp) * scale)
1662 .collect()
1663 } else {
1664 m.as_slice().iter().map(|&v| v as f32).collect()
1665 }
1666 }
1667 DType::I16 => {
1668 let t = proto_data.mask_coefficients.as_typed::<i16>().expect("I16");
1669 let m = t.map_read()?;
1670 if let Some(q) = t.quantization() {
1671 use edgefirst_tensor::QuantMode;
1672 let (scale, zp) = match q.mode() {
1673 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1674 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1675 other => {
1676 return Err(Error::NotSupported(format!(
1677 "I16 mask_coefficients quantization mode {other:?} not supported"
1678 )));
1679 }
1680 };
1681 m.as_slice()
1682 .iter()
1683 .map(|&v| (v as f32 - zp) * scale)
1684 .collect()
1685 } else {
1686 m.as_slice().iter().map(|&v| v as f32).collect()
1687 }
1688 }
1689 other => {
1690 return Err(Error::InvalidShape(format!(
1691 "mask_coefficients dtype {other:?} not supported"
1692 )));
1693 }
1694 };
1695
1696 let (lx0, lx_range, ly0, ly_range) = match letterbox {
1698 Some([lx0, ly0, lx1, ly1]) => (lx0, lx1 - lx0, ly0, ly1 - ly0),
1699 None => (0.0_f32, 1.0_f32, 0.0_f32, 1.0_f32),
1700 };
1701
1702 match proto_data.protos.dtype() {
1705 DType::F32 => {
1706 let t = proto_data.protos.as_typed::<f32>().expect("F32");
1707 let m = t.map_read()?;
1708 self.draw_proto_masks_inner(
1709 dst_slice,
1710 dst_w,
1711 dst_h,
1712 dst_rs,
1713 channels,
1714 detect,
1715 m.as_slice(),
1716 &coeff_f32,
1717 proto_h,
1718 proto_w,
1719 num_protos,
1720 opacity,
1721 (lx0, lx_range, ly0, ly_range),
1722 color_mode,
1723 0.0_f32,
1724 |p: &f32, _| *p,
1725 );
1726 }
1727 DType::F16 => {
1728 let t = proto_data.protos.as_typed::<half::f16>().expect("F16");
1729 let m = t.map_read()?;
1730 self.draw_proto_masks_inner(
1731 dst_slice,
1732 dst_w,
1733 dst_h,
1734 dst_rs,
1735 channels,
1736 detect,
1737 m.as_slice(),
1738 &coeff_f32,
1739 proto_h,
1740 proto_w,
1741 num_protos,
1742 opacity,
1743 (lx0, lx_range, ly0, ly_range),
1744 color_mode,
1745 0.0_f32,
1746 |p: &half::f16, _| p.to_f32(),
1747 );
1748 }
1749 DType::I8 => {
1750 use edgefirst_tensor::QuantMode;
1751 let t = proto_data.protos.as_typed::<i8>().expect("I8");
1752 let m = t.map_read()?;
1753 let quant = t.quantization().ok_or_else(|| {
1754 Error::InvalidShape("I8 protos require quantization metadata".into())
1755 })?;
1756 let (scale, zp) = match quant.mode() {
1757 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1758 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1759 QuantMode::PerChannel { axis, .. }
1760 | QuantMode::PerChannelSymmetric { axis, .. } => {
1761 return Err(Error::NotSupported(format!(
1762 "per-channel quantization (axis={axis}) in draw_proto_masks \
1763 CPU path not yet supported"
1764 )));
1765 }
1766 };
1767 self.draw_proto_masks_inner(
1768 dst_slice,
1769 dst_w,
1770 dst_h,
1771 dst_rs,
1772 channels,
1773 detect,
1774 m.as_slice(),
1775 &coeff_f32,
1776 proto_h,
1777 proto_w,
1778 num_protos,
1779 opacity,
1780 (lx0, lx_range, ly0, ly_range),
1781 color_mode,
1782 scale,
1783 move |p: &i8, _| (*p as f32) - zp,
1784 );
1785 }
1786 other => {
1787 return Err(Error::InvalidShape(format!(
1788 "proto tensor dtype {other:?} not supported"
1789 )));
1790 }
1791 }
1792
1793 Ok(())
1794 }
1795
1796 #[allow(clippy::too_many_arguments)]
1797 fn draw_proto_masks_inner<P: Copy>(
1798 &self,
1799 dst_slice: &mut [u8],
1800 dst_w: usize,
1801 dst_h: usize,
1802 dst_rs: usize,
1803 channels: usize,
1804 detect: &[DetectBox],
1805 protos: &[P],
1806 coeff_all_f32: &[f32],
1807 proto_h: usize,
1808 proto_w: usize,
1809 num_protos: usize,
1810 opacity: f32,
1811 letterbox_xy: (f32, f32, f32, f32),
1812 color_mode: crate::ColorMode,
1813 acc_scale: f32,
1814 load_f32: impl Fn(&P, f32) -> f32 + Copy,
1815 ) {
1816 let (lx0, lx_range, ly0, ly_range) = letterbox_xy;
1817 let stride_y = proto_w * num_protos;
1818 for (idx, det) in detect.iter().enumerate() {
1819 let coeff = &coeff_all_f32[idx * num_protos..(idx + 1) * num_protos];
1820 let color_index = color_mode.index(idx, det.label);
1821 let color = self.colors[color_index % self.colors.len()];
1822 let alpha = if opacity == 1.0 {
1823 color[3] as u16
1824 } else {
1825 (color[3] as f32 * opacity).round() as u16
1826 };
1827
1828 let start_x = (dst_w as f32 * det.bbox.xmin).round() as usize;
1829 let start_y = (dst_h as f32 * det.bbox.ymin).round() as usize;
1830 let end_x = ((dst_w as f32 * det.bbox.xmax).round() as usize).min(dst_w);
1831 let end_y = ((dst_h as f32 * det.bbox.ymax).round() as usize).min(dst_h);
1832
1833 for y in start_y..end_y {
1834 for x in start_x..end_x {
1835 let px = (lx0 + (x as f32 / dst_w as f32) * lx_range) * proto_w as f32 - 0.5;
1836 let py = (ly0 + (y as f32 / dst_h as f32) * ly_range) * proto_h as f32 - 0.5;
1837
1838 let x0 = (px.floor() as isize).clamp(0, proto_w as isize - 1) as usize;
1842 let y0 = (py.floor() as isize).clamp(0, proto_h as isize - 1) as usize;
1843 let x1 = (x0 + 1).min(proto_w - 1);
1844 let y1 = (y0 + 1).min(proto_h - 1);
1845 let fx = px - px.floor();
1846 let fy = py - py.floor();
1847 let w00 = (1.0 - fx) * (1.0 - fy);
1848 let w10 = fx * (1.0 - fy);
1849 let w01 = (1.0 - fx) * fy;
1850 let w11 = fx * fy;
1851 let b00 = y0 * stride_y + x0 * num_protos;
1852 let b10 = y0 * stride_y + x1 * num_protos;
1853 let b01 = y1 * stride_y + x0 * num_protos;
1854 let b11 = y1 * stride_y + x1 * num_protos;
1855 let mut acc = 0.0_f32;
1856 for p in 0..num_protos {
1857 let v00 = load_f32(&protos[b00 + p], 0.0);
1858 let v10 = load_f32(&protos[b10 + p], 0.0);
1859 let v01 = load_f32(&protos[b01 + p], 0.0);
1860 let v11 = load_f32(&protos[b11 + p], 0.0);
1861 let val = w00 * v00 + w10 * v10 + w01 * v01 + w11 * v11;
1862 acc += coeff[p] * val;
1863 }
1864 let final_acc = if acc_scale == 0.0 {
1865 acc
1866 } else {
1867 acc_scale * acc
1868 };
1869 let mask = 1.0 / (1.0 + (-final_acc).exp());
1873 if mask < 0.5 {
1874 continue;
1875 }
1876 let dst_index = y * dst_rs + x * channels;
1877 for c in 0..3 {
1878 dst_slice[dst_index + c] = ((color[c] as u16 * alpha
1879 + dst_slice[dst_index + c] as u16 * (255 - alpha))
1880 / 255) as u8;
1881 }
1882 }
1883 }
1884 }
1885 }
1886}