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 new_nearest() -> Self {
473 let resizer = fast_image_resize::Resizer::new();
474 let options = fast_image_resize::ResizeOptions::new()
475 .resize_alg(fast_image_resize::ResizeAlg::Nearest)
476 .use_alpha(false);
477 log::debug!("CPUConverter created");
478 Self {
479 resizer,
480 options,
481 colors: crate::DEFAULT_COLORS_U8,
482 widen_scratch: None,
483 resize_destride_scratch: Vec::new(),
484 resize_dst_destride_scratch: Vec::new(),
485 nv_strip_scratch: Vec::new(),
486 nv_strip_y_pack: Vec::new(),
487 nv_strip_uv_pack: Vec::new(),
488 #[cfg(test)]
489 fused_hits: 0,
490 #[cfg(test)]
491 last_tmp_dims: None,
492 convert_tmp: None,
493 convert_tmp2: None,
494 convert_src_sub: None,
495 }
496 }
497
498 #[cfg(test)]
500 pub(super) fn fused_hits(&self) -> u64 {
501 self.fused_hits
502 }
503
504 #[cfg(test)]
507 pub(super) fn last_tmp_dims(&self) -> Option<(usize, usize)> {
508 self.last_tmp_dims
509 }
510
511 pub(crate) fn support_conversion_pf(src: PixelFormat, dst: PixelFormat) -> bool {
512 use PixelFormat::*;
513 matches!(
514 (src, dst),
515 (Nv12, Rgb)
516 | (Nv12, Rgba)
517 | (Nv12, Grey)
518 | (Nv16, Rgb)
519 | (Nv16, Rgba)
520 | (Nv16, Bgra)
521 | (Nv24, Rgb)
522 | (Nv24, Rgba)
523 | (Nv24, Grey)
524 | (Nv24, Bgra)
525 | (Yuyv, Rgb)
526 | (Yuyv, Rgba)
527 | (Yuyv, Grey)
528 | (Yuyv, Yuyv)
529 | (Yuyv, PlanarRgb)
530 | (Yuyv, PlanarRgba)
531 | (Yuyv, Nv16)
532 | (Vyuy, Rgb)
533 | (Vyuy, Rgba)
534 | (Vyuy, Grey)
535 | (Vyuy, Vyuy)
536 | (Vyuy, PlanarRgb)
537 | (Vyuy, PlanarRgba)
538 | (Vyuy, Nv16)
539 | (Rgba, Rgb)
540 | (Rgba, Rgba)
541 | (Rgba, Grey)
542 | (Rgba, Yuyv)
543 | (Rgba, PlanarRgb)
544 | (Rgba, PlanarRgba)
545 | (Rgba, Nv16)
546 | (Rgb, Rgb)
547 | (Rgb, Rgba)
548 | (Rgb, Grey)
549 | (Rgb, Yuyv)
550 | (Rgb, PlanarRgb)
551 | (Rgb, PlanarRgba)
552 | (Rgb, Nv16)
553 | (Grey, Rgb)
554 | (Grey, Rgba)
555 | (Grey, Grey)
556 | (Grey, Yuyv)
557 | (Grey, PlanarRgb)
558 | (Grey, PlanarRgba)
559 | (Grey, Nv16)
560 | (Nv12, Bgra)
561 | (Yuyv, Bgra)
562 | (Vyuy, Bgra)
563 | (Rgba, Bgra)
564 | (Rgb, Bgra)
565 | (Grey, Bgra)
566 | (Bgra, Bgra)
567 | (PlanarRgb, Rgb)
568 | (PlanarRgb, Rgba)
569 | (PlanarRgba, Rgb)
570 | (PlanarRgba, Rgba)
571 | (PlanarRgb, Bgra)
572 | (PlanarRgba, Bgra)
573 )
574 }
575
576 pub(crate) fn convert_format_pf(
578 src: &Tensor<u8>,
579 dst: &mut Tensor<u8>,
580 src_fmt: PixelFormat,
581 dst_fmt: PixelFormat,
582 cp: ColorParams,
583 ) -> Result<()> {
584 let _timer = FunctionTimer::new(format!(
585 "ImageProcessor::convert_format {} to {}",
586 src_fmt, dst_fmt,
587 ));
588
589 use PixelFormat::*;
590 match (src_fmt, dst_fmt) {
591 (Nv12, Rgb) => Self::convert_nv12_to_rgb(src, dst, cp),
592 (Nv12, Rgba) => Self::convert_nv12_to_rgba(src, dst, cp),
593 (Nv12, Grey) => Self::convert_nv12_to_grey(src, dst, cp),
594 (Yuyv, Rgb) => Self::convert_yuyv_to_rgb(src, dst, cp),
595 (Yuyv, Rgba) => Self::convert_yuyv_to_rgba(src, dst, cp),
596 (Yuyv, Grey) => Self::convert_yuyv_to_grey(src, dst, cp),
597 (Yuyv, Yuyv) => Self::copy_image(src, dst),
598 (Yuyv, PlanarRgb) => Self::convert_yuyv_to_8bps(src, dst, cp),
599 (Yuyv, PlanarRgba) => Self::convert_yuyv_to_prgba(src, dst, cp),
600 (Yuyv, Nv16) => Self::convert_yuyv_to_nv16(src, dst),
601 (Vyuy, Rgb) => Self::convert_vyuy_to_rgb(src, dst, cp),
602 (Vyuy, Rgba) => Self::convert_vyuy_to_rgba(src, dst, cp),
603 (Vyuy, Grey) => Self::convert_vyuy_to_grey(src, dst, cp),
604 (Vyuy, Vyuy) => Self::copy_image(src, dst),
605 (Vyuy, PlanarRgb) => Self::convert_vyuy_to_8bps(src, dst, cp),
606 (Vyuy, PlanarRgba) => Self::convert_vyuy_to_prgba(src, dst, cp),
607 (Vyuy, Nv16) => Self::convert_vyuy_to_nv16(src, dst),
608 (Rgba, Rgb) => Self::convert_rgba_to_rgb(src, dst),
609 (Rgba, Rgba) => Self::copy_image(src, dst),
610 (Rgba, Grey) => Self::convert_rgba_to_grey(src, dst),
611 (Rgba, Yuyv) => Self::convert_rgba_to_yuyv(src, dst, cp),
612 (Rgba, PlanarRgb) => Self::convert_rgba_to_8bps(src, dst),
613 (Rgba, PlanarRgba) => Self::convert_rgba_to_prgba(src, dst),
614 (Rgba, Nv16) => Self::convert_rgba_to_nv16(src, dst, cp),
615 (Rgb, Rgb) => Self::copy_image(src, dst),
616 (Rgb, Rgba) => Self::convert_rgb_to_rgba(src, dst),
617 (Rgb, Grey) => Self::convert_rgb_to_grey(src, dst),
618 (Rgb, Yuyv) => Self::convert_rgb_to_yuyv(src, dst, cp),
619 (Rgb, PlanarRgb) => Self::convert_rgb_to_8bps(src, dst),
620 (Rgb, PlanarRgba) => Self::convert_rgb_to_prgba(src, dst),
621 (Rgb, Nv16) => Self::convert_rgb_to_nv16(src, dst, cp),
622 (Grey, Rgb) => Self::convert_grey_to_rgb(src, dst),
623 (Grey, Rgba) => Self::convert_grey_to_rgba(src, dst),
624 (Grey, Grey) => Self::copy_image(src, dst),
625 (Grey, Yuyv) => Self::convert_grey_to_yuyv(src, dst, cp),
626 (Grey, PlanarRgb) => Self::convert_grey_to_8bps(src, dst),
627 (Grey, PlanarRgba) => Self::convert_grey_to_prgba(src, dst),
628 (Grey, Nv16) => Self::convert_grey_to_nv16(src, dst, cp),
629
630 (Nv16, Rgb) => Self::convert_nv16_to_rgb(src, dst, cp),
632 (Nv16, Rgba) => Self::convert_nv16_to_rgba(src, dst, cp),
633 (Nv24, Rgb) => Self::convert_nv24_to_rgb(src, dst, cp),
634 (Nv24, Rgba) => Self::convert_nv24_to_rgba(src, dst, cp),
635 (Nv24, Grey) => Self::convert_nv24_to_grey(src, dst, cp),
636 (PlanarRgb, Rgb) => Self::convert_8bps_to_rgb(src, dst),
637 (PlanarRgb, Rgba) => Self::convert_8bps_to_rgba(src, dst),
638 (PlanarRgba, Rgb) => Self::convert_prgba_to_rgb(src, dst),
639 (PlanarRgba, Rgba) => Self::convert_prgba_to_rgba(src, dst),
640
641 (Bgra, Bgra) => Self::copy_image(src, dst),
643 (Nv12, Bgra) => {
644 Self::convert_nv12_to_rgba(src, dst, cp)?;
645 Self::swizzle_rb_4chan(dst)
646 }
647 (Nv16, Bgra) => {
648 Self::convert_nv16_to_rgba(src, dst, cp)?;
649 Self::swizzle_rb_4chan(dst)
650 }
651 (Nv24, Bgra) => {
652 Self::convert_nv24_to_rgba(src, dst, cp)?;
653 Self::swizzle_rb_4chan(dst)
654 }
655 (Yuyv, Bgra) => {
656 Self::convert_yuyv_to_rgba(src, dst, cp)?;
657 Self::swizzle_rb_4chan(dst)
658 }
659 (Vyuy, Bgra) => {
660 Self::convert_vyuy_to_rgba(src, dst, cp)?;
661 Self::swizzle_rb_4chan(dst)
662 }
663 (Rgba, Bgra) => {
664 Self::copy_image(src, dst)?;
665 Self::swizzle_rb_4chan(dst)
666 }
667 (Rgb, Bgra) => {
668 Self::convert_rgb_to_rgba(src, dst)?;
669 Self::swizzle_rb_4chan(dst)
670 }
671 (Grey, Bgra) => {
672 Self::convert_grey_to_rgba(src, dst)?;
673 Self::swizzle_rb_4chan(dst)
674 }
675 (PlanarRgb, Bgra) => {
676 Self::convert_8bps_to_rgba(src, dst)?;
677 Self::swizzle_rb_4chan(dst)
678 }
679 (PlanarRgba, Bgra) => {
680 Self::convert_prgba_to_rgba(src, dst)?;
681 Self::swizzle_rb_4chan(dst)
682 }
683
684 (s, d) => Err(Error::NotSupported(format!("Conversion from {s} to {d}",))),
685 }
686 }
687
688 pub(crate) fn fill_image_outside_crop_u8(
690 dst: &mut Tensor<u8>,
691 rgba: [u8; 4],
692 crop: Rect,
693 ) -> Result<()> {
694 let dst_fmt = dst.format().unwrap();
695 let dst_w = dst.width().unwrap();
696 let dst_h = dst.height().unwrap();
697 let cm = crate::colorimetry::resolve_colorimetry(dst.colorimetry(), dst.height());
701 let cp = ColorParams {
702 matrix: crate::colorimetry::yuv_matrix(cm.encoding.unwrap()),
703 range: crate::colorimetry::yuv_range(cm.range.unwrap()),
704 encoding: cm.encoding.unwrap(),
705 range_kind: cm.range.unwrap(),
706 src_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
707 dst_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
708 };
709 let dst_stride = tensor_row_stride(dst);
710 let mut dst_map = dst.map_mut()?;
711 let dst_tup = (dst_map.as_mut_slice(), dst_w, dst_h, dst_stride);
712 Self::fill_outside_crop_dispatch(dst_tup, dst_fmt, rgba, crop, cp)
713 }
714
715 fn fill_outside_crop_dispatch(
720 dst: (&mut [u8], usize, usize, usize),
721 fmt: PixelFormat,
722 rgba: [u8; 4],
723 crop: Rect,
724 cp: ColorParams,
725 ) -> Result<()> {
726 use PixelFormat::*;
727 match fmt {
728 Rgba | Bgra => Self::fill_image_outside_crop_(dst, rgba, crop),
729 Rgb => Self::fill_image_outside_crop_(dst, Self::rgba_to_rgb(rgba), crop),
730 Grey => Self::fill_image_outside_crop_(dst, Self::rgba_to_grey(rgba), crop),
731 Yuyv => {
732 let (bytes, w, h, stride) = dst;
733 let yuyv = Self::rgba_to_yuyv(rgba, cp);
734 Self::fill_image_outside_crop_(
737 (&mut *bytes, w / 2, h, stride),
738 yuyv,
739 Rect::new(crop.left / 2, crop.top, crop.width.div_ceil(2), crop.height),
740 )?;
741 if w % 2 == 1 {
749 let last = w - 1;
750 let col_outside = last < crop.left || last >= crop.left + crop.width;
751 for y in 0..h {
752 let row_outside = y < crop.top || y >= crop.top + crop.height;
753 if row_outside || col_outside {
754 let off = y * stride + last * 2;
755 bytes[off] = yuyv[0];
756 bytes[off + 1] = yuyv[1];
757 }
758 }
759 }
760 Ok(())
761 }
762 PlanarRgb => Self::fill_image_outside_crop_planar(dst, Self::rgba_to_rgb(rgba), crop),
763 PlanarRgba => Self::fill_image_outside_crop_planar(dst, rgba, crop),
764 Nv16 => {
765 let yuyv = Self::rgba_to_yuyv(rgba, cp);
766 Self::fill_image_outside_crop_yuv_semiplanar(dst, yuyv[0], [yuyv[1], yuyv[3]], crop)
767 }
768 _ => Err(Error::Internal(format!(
769 "Found unexpected destination {fmt}",
770 ))),
771 }
772 }
773}
774
775impl ImageProcessorTrait for CPUProcessor {
776 fn convert(
777 &mut self,
778 src: &TensorDyn,
779 dst: &mut TensorDyn,
780 rotation: Rotation,
781 flip: Flip,
782 crop: Crop,
783 ) -> Result<()> {
784 let crop = crop.resolve(
785 src.width().unwrap_or(0),
786 src.height().unwrap_or(0),
787 dst.width().unwrap_or(0),
788 dst.height().unwrap_or(0),
789 )?;
790 self.convert_impl(src, dst, rotation, flip, crop)
791 }
792
793 fn draw_decoded_masks(
794 &mut self,
795 dst: &mut TensorDyn,
796 detect: &[DetectBox],
797 segmentation: &[Segmentation],
798 overlay: crate::MaskOverlay<'_>,
799 ) -> Result<()> {
800 prepare_dst_base_cpu(dst, overlay.background)?;
804 let dst = dst.as_typed_mut::<u8>().ok_or(Error::NotAnImage)?;
805 self.draw_decoded_masks_impl(
806 dst,
807 detect,
808 segmentation,
809 overlay.opacity,
810 overlay.color_mode,
811 )
812 }
813
814 fn draw_proto_masks(
815 &mut self,
816 dst: &mut TensorDyn,
817 detect: &[DetectBox],
818 proto_data: &ProtoData,
819 overlay: crate::MaskOverlay<'_>,
820 ) -> Result<()> {
821 prepare_dst_base_cpu(dst, overlay.background)?;
822 let dst = dst.as_typed_mut::<u8>().ok_or(Error::NotAnImage)?;
823 self.draw_proto_masks_impl(
824 dst,
825 detect,
826 proto_data,
827 overlay.opacity,
828 overlay.letterbox,
829 overlay.color_mode,
830 )
831 }
832
833 fn set_class_colors(&mut self, colors: &[[u8; 4]]) -> Result<()> {
834 for (c, new_c) in self.colors.iter_mut().zip(colors.iter()) {
835 *c = *new_c;
836 }
837 Ok(())
838 }
839}
840
841impl CPUProcessor {
843 pub(crate) fn convert_impl(
845 &mut self,
846 src: &TensorDyn,
847 dst: &mut TensorDyn,
848 rotation: Rotation,
849 flip: Flip,
850 crop: ResolvedCrop,
851 ) -> Result<()> {
852 let src_fmt = src.format().ok_or(Error::NotAnImage)?;
853 let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
854
855 let src_cm = crate::colorimetry::effective_colorimetry(src);
860 let dst_cm = crate::colorimetry::effective_colorimetry(dst);
861 let src_full = src_cm.range == Some(edgefirst_tensor::ColorRange::Full);
862 let dst_full = dst_cm.range == Some(edgefirst_tensor::ColorRange::Full);
863 let src_params = ColorParams {
864 matrix: crate::colorimetry::yuv_matrix(src_cm.encoding.unwrap()),
865 range: crate::colorimetry::yuv_range(src_cm.range.unwrap()),
866 encoding: src_cm.encoding.unwrap(),
867 range_kind: src_cm.range.unwrap(),
868 src_full_range: src_full,
869 dst_full_range: dst_full,
870 };
871 let dst_params = ColorParams {
872 matrix: crate::colorimetry::yuv_matrix(dst_cm.encoding.unwrap()),
873 range: crate::colorimetry::yuv_range(dst_cm.range.unwrap()),
874 encoding: dst_cm.encoding.unwrap(),
875 range_kind: dst_cm.range.unwrap(),
876 src_full_range: src_full,
877 dst_full_range: dst_full,
878 };
879 match (src.dtype(), dst.dtype()) {
880 (DType::U8, DType::U8) => {
881 let src = src.as_typed::<u8>().unwrap();
882 let dst = dst.as_typed_mut::<u8>().unwrap();
883 self.convert_u8(
884 src, dst, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
885 )
886 }
887 (DType::U8, DType::I8) => {
888 let src_u8 = src.as_typed::<u8>().unwrap();
891 let dst_i8 = dst.as_typed_mut::<i8>().unwrap();
892 let dst_u8 = unsafe { &mut *(dst_i8 as *mut Tensor<i8> as *mut Tensor<u8>) };
896 self.convert_u8(
897 src_u8, dst_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
898 )?;
899 apply_int8_xor_bias_rows(dst_u8, dst_fmt)
901 }
902 (DType::U8, d @ (DType::F32 | DType::F16)) => {
903 let src_u8 = src.as_typed::<u8>().unwrap();
904 let dw = dst.width().ok_or(Error::NotAnImage)?;
905 let dh = dst.height().ok_or(Error::NotAnImage)?;
906 let scratch_matches = self.widen_scratch.as_ref().is_some_and(|t| {
911 t.width() == Some(dw) && t.height() == Some(dh) && t.format() == Some(dst_fmt)
912 });
913 let mut tmp = if scratch_matches {
914 self.widen_scratch.take().unwrap()
915 } else {
916 TensorDyn::image(
917 dw,
918 dh,
919 dst_fmt,
920 DType::U8,
921 Some(TensorMemory::Mem),
922 edgefirst_tensor::CpuAccess::ReadWrite,
923 )?
924 };
925 {
926 let tmp_u8 = tmp.as_typed_mut::<u8>().unwrap();
927 self.convert_u8(
928 src_u8, tmp_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params,
929 dst_params,
930 )?;
931 }
932 {
942 let tmp_u8 = tmp.as_typed::<u8>().unwrap();
943 let (rows, row_len) = logical_surface(tmp_u8)?;
944 let src_stride = tensor_row_stride(tmp_u8);
945 let dst_stride_bytes = dst.effective_row_stride().ok_or(Error::NotAnImage)?;
946 let src_map = tmp_u8.map_read()?;
947 guard_plane(
948 src_map.as_slice().len(),
949 src_stride,
950 rows,
951 row_len,
952 "widen src",
953 )?;
954 let src_rows = src_map.as_slice().chunks(src_stride).take(rows);
955 let elem = d.size();
956 if !dst_stride_bytes.is_multiple_of(elem) {
957 return Err(Error::InvalidShape(format!(
958 "{d} destination row stride {dst_stride_bytes} is not a multiple of \
959 the element size {elem}"
960 )));
961 }
962 let dst_stride = dst_stride_bytes / elem;
963 match d {
964 DType::F32 => {
965 let dst_t = dst.as_typed_mut::<f32>().unwrap();
966 let mut dst_map = dst_t.map_mut()?;
967 guard_plane(
969 dst_map.as_slice().len(),
970 dst_stride,
971 rows,
972 row_len,
973 "widen f32 dst",
974 )?;
975 for (s, dr) in
979 src_rows.zip(dst_map.as_mut_slice().chunks_mut(dst_stride))
980 {
981 simd::widen_u8_to_f32_norm(s, &mut dr[..row_len]);
982 }
983 }
984 DType::F16 => {
985 let dst_t = dst.as_typed_mut::<half::f16>().unwrap();
986 let mut dst_map = dst_t.map_mut()?;
987 guard_plane(
988 dst_map.as_slice().len(),
989 dst_stride,
990 rows,
991 row_len,
992 "widen f16 dst",
993 )?;
994 for (s, dr) in
999 src_rows.zip(dst_map.as_mut_slice().chunks_mut(dst_stride))
1000 {
1001 simd::widen_u8_to_f16_norm(s, &mut dr[..row_len]);
1002 }
1003 }
1004 _ => unreachable!(),
1005 }
1006 }
1007 self.widen_scratch = Some(tmp);
1008 Ok(())
1009 }
1010 (s, d) => Err(Error::NotSupported(format!("dtype {s} -> {d}",))),
1011 }
1012 }
1013
1014 fn reuse_or_alloc_image(
1020 cached: Option<Tensor<u8>>,
1021 w: usize,
1022 h: usize,
1023 fmt: PixelFormat,
1024 ) -> Result<Tensor<u8>> {
1025 if let Some(t) = cached {
1026 if t.width() == Some(w) && t.height() == Some(h) && t.format() == Some(fmt) {
1027 return Ok(t);
1028 }
1029 }
1030 Ok(Tensor::<u8>::image(
1031 w,
1032 h,
1033 fmt,
1034 Some(TensorMemory::Mem),
1035 edgefirst_tensor::CpuAccess::ReadWrite,
1036 )?)
1037 }
1038
1039 fn pre_resize_region(
1060 &self,
1061 src_fmt: PixelFormat,
1062 (src_w, src_h): (usize, usize),
1063 (dst_w, dst_h): (usize, usize),
1064 rotation: Rotation,
1065 crop: ResolvedCrop,
1066 ) -> Option<Rect> {
1067 use PixelFormat::{Nv12, Nv16, Nv24};
1068
1069 if !matches!(src_fmt, Nv12 | Nv16 | Nv24) {
1070 return None;
1071 }
1072 let r = crop.src_rect?;
1073 let full_src = Rect {
1074 left: 0,
1075 top: 0,
1076 width: src_w,
1077 height: src_h,
1078 };
1079 if r == full_src {
1080 return None;
1081 }
1082
1083 let d = crop.dst_rect.unwrap_or(Rect {
1088 left: 0,
1089 top: 0,
1090 width: dst_w,
1091 height: dst_h,
1092 });
1093 let (dst_x, dst_y) = match rotation {
1094 Rotation::None | Rotation::Rotate180 => (d.width, d.height),
1095 Rotation::Clockwise90 | Rotation::CounterClockwise90 => (d.height, d.width),
1096 };
1097 let halo_x = self.filter_halo(r.width, dst_x)?;
1098 let halo_y = self.filter_halo(r.height, dst_y)?;
1099
1100 let mut left = r.left.saturating_sub(halo_x);
1101 let mut top = r.top.saturating_sub(halo_y);
1102 let mut right = (r.left + r.width + halo_x).min(src_w);
1103 let mut bottom = (r.top + r.height + halo_y).min(src_h);
1104
1105 let (align_x, align_y) = match src_fmt {
1111 Nv12 => (2, 2),
1112 Nv16 => (2, 1),
1113 _ => (1, 1),
1114 };
1115 left -= left % align_x;
1116 top -= top % align_y;
1117 right = right.next_multiple_of(align_x).min(src_w);
1118 bottom = bottom.next_multiple_of(align_y).min(src_h);
1119
1120 let grown = Rect {
1121 left,
1122 top,
1123 width: right - left,
1124 height: bottom - top,
1125 };
1126 (grown != full_src).then_some(grown)
1129 }
1130
1131 #[allow(clippy::too_many_arguments)]
1133 fn convert_u8(
1134 &mut self,
1135 src: &Tensor<u8>,
1136 dst: &mut Tensor<u8>,
1137 src_fmt: PixelFormat,
1138 dst_fmt: PixelFormat,
1139 rotation: Rotation,
1140 flip: Flip,
1141 crop: ResolvedCrop,
1142 src_params: ColorParams,
1143 dst_params: ColorParams,
1144 ) -> Result<()> {
1145 use PixelFormat::*;
1146
1147 #[cfg(test)]
1148 {
1149 self.last_tmp_dims = None;
1150 }
1151
1152 let src_w = src.width().unwrap();
1153 let src_h = src.height().unwrap();
1154 let dst_w = dst.width().unwrap();
1155 let dst_h = dst.height().unwrap();
1156
1157 crop.check_crop_dims(src_w, src_h, dst_w, dst_h)?;
1158
1159 let intermediate = match (src_fmt, dst_fmt) {
1161 (Nv12, Rgb) => Rgb,
1162 (Nv12, Rgba) => Rgba,
1163 (Nv12, Grey) => Grey,
1164 (Nv12, Yuyv) => Rgba,
1165 (Nv12, Nv16) => Rgba,
1166 (Nv12, PlanarRgb) => Rgb,
1167 (Nv12, PlanarRgba) => Rgba,
1168 (Nv16, PlanarRgb) => Rgb,
1169 (Nv16, PlanarRgba) => Rgba,
1170 (Nv24, PlanarRgb) => Rgb,
1171 (Nv24, PlanarRgba) => Rgba,
1172 (Yuyv, Rgb) => Rgb,
1173 (Yuyv, Rgba) => Rgba,
1174 (Yuyv, Grey) => Grey,
1175 (Yuyv, Yuyv) => Rgba,
1176 (Yuyv, PlanarRgb) => Rgb,
1177 (Yuyv, PlanarRgba) => Rgba,
1178 (Yuyv, Nv16) => Rgba,
1179 (Vyuy, Rgb) => Rgb,
1180 (Vyuy, Rgba) => Rgba,
1181 (Vyuy, Grey) => Grey,
1182 (Vyuy, Vyuy) => Rgba,
1183 (Vyuy, PlanarRgb) => Rgb,
1184 (Vyuy, PlanarRgba) => Rgba,
1185 (Vyuy, Nv16) => Rgba,
1186 (Rgba, Rgb) => Rgba,
1187 (Rgba, Rgba) => Rgba,
1188 (Rgba, Grey) => Grey,
1189 (Rgba, Yuyv) => Rgba,
1190 (Rgba, PlanarRgb) => Rgba,
1191 (Rgba, PlanarRgba) => Rgba,
1192 (Rgba, Nv16) => Rgba,
1193 (Rgb, Rgb) => Rgb,
1194 (Rgb, Rgba) => Rgb,
1195 (Rgb, Grey) => Grey,
1196 (Rgb, Yuyv) => Rgb,
1197 (Rgb, PlanarRgb) => Rgb,
1198 (Rgb, PlanarRgba) => Rgb,
1199 (Rgb, Nv16) => Rgb,
1200 (Grey, Rgb) => Rgb,
1201 (Grey, Rgba) => Rgba,
1202 (Grey, Grey) => Grey,
1203 (Grey, Yuyv) => Grey,
1204 (Grey, PlanarRgb) => Grey,
1205 (Grey, PlanarRgba) => Grey,
1206 (Grey, Nv16) => Grey,
1207 (Nv12, Bgra) => Rgba,
1208 (Yuyv, Bgra) => Rgba,
1209 (Vyuy, Bgra) => Rgba,
1210 (Rgba, Bgra) => Rgba,
1211 (Rgb, Bgra) => Rgb,
1212 (Grey, Bgra) => Grey,
1213 (Bgra, Bgra) => Bgra,
1214 (Nv16, Rgb) => Rgb,
1215 (Nv16, Rgba) => Rgba,
1216 (Nv16, Bgra) => Rgba,
1217 (Nv24, Rgb) => Rgb,
1218 (Nv24, Rgba) => Rgba,
1219 (Nv24, Grey) => Grey,
1220 (Nv24, Bgra) => Rgba,
1221 (PlanarRgb, Rgb) => Rgb,
1222 (PlanarRgb, Rgba) => Rgb,
1223 (PlanarRgb, Bgra) => Rgb,
1224 (PlanarRgba, Rgb) => Rgba,
1225 (PlanarRgba, Rgba) => Rgba,
1226 (PlanarRgba, Bgra) => Rgba,
1227 (s, d) => {
1228 return Err(Error::NotSupported(format!("Conversion from {s} to {d}",)));
1229 }
1230 };
1231
1232 let need_resize_flip_rotation = rotation != Rotation::None
1233 || flip != Flip::None
1234 || src_w != dst_w
1235 || src_h != dst_h
1236 || crop.src_rect.is_some_and(|c| {
1237 c != Rect {
1238 left: 0,
1239 top: 0,
1240 width: src_w,
1241 height: src_h,
1242 }
1243 })
1244 || crop.dst_rect.is_some_and(|c| {
1245 c != Rect {
1246 left: 0,
1247 top: 0,
1248 width: dst_w,
1249 height: dst_h,
1250 }
1251 });
1252
1253 let direct_is_yuv_src = matches!(src_fmt, Nv12 | Nv16 | Nv24 | Yuyv | Vyuy);
1256 let direct_params = if direct_is_yuv_src {
1257 src_params
1258 } else {
1259 dst_params
1260 };
1261
1262 let full_dst_rect = Rect {
1279 left: 0,
1280 top: 0,
1281 width: dst_w,
1282 height: dst_h,
1283 };
1284 let fused_region = if rotation != Rotation::None || flip != Flip::None {
1285 None
1286 } else {
1287 match crop.src_rect {
1288 None if src_w == dst_w && src_h == dst_h => Some(None),
1289 None => None,
1290 Some(r)
1291 if r.width == dst_w
1292 && r.height == dst_h
1293 && crop.dst_rect.is_none_or(|d| d == full_dst_rect)
1294 && chroma_alignment_ok(src_fmt, r) =>
1295 {
1296 Some(Some(r))
1297 }
1298 Some(_) => None,
1299 }
1300 };
1301 if let Some(region) = fused_region {
1302 if matches!(src_fmt, Nv12 | Nv16 | Nv24) && matches!(dst_fmt, PlanarRgb | PlanarRgba) {
1303 #[cfg(test)]
1304 {
1305 self.fused_hits += 1;
1306 }
1307 return self.convert_nv_to_planar_fused(
1308 src,
1309 dst,
1310 src_fmt,
1311 dst_fmt,
1312 direct_params,
1313 region,
1314 );
1315 }
1316 }
1317
1318 if !need_resize_flip_rotation && Self::support_conversion_pf(src_fmt, dst_fmt) {
1320 return Self::convert_format_pf(src, dst, src_fmt, dst_fmt, direct_params);
1321 }
1322
1323 if dst_fmt == Yuyv && !dst_w.is_multiple_of(2) {
1325 return Err(Error::NotSupported(format!(
1326 "{} destination must have width divisible by 2",
1327 dst_fmt,
1328 )));
1329 }
1330
1331 let mut cached_tmp = self.convert_tmp.take();
1339 let mut cached_tmp2 = self.convert_tmp2.take();
1340
1341 let pre_region = if intermediate != src_fmt {
1346 self.pre_resize_region(src_fmt, (src_w, src_h), (dst_w, dst_h), rotation, crop)
1347 } else {
1348 None
1349 };
1350
1351 let tmp_holder: Option<Tensor<u8>> = if intermediate != src_fmt {
1353 let _s = tracing::trace_span!(
1354 "image.convert.cpu.format_convert",
1355 from = ?src_fmt,
1356 to = ?intermediate,
1357 pass = "pre_resize",
1358 )
1359 .entered();
1360 let (tmp_w, tmp_h) = pre_region.map_or((src_w, src_h), |g| (g.width, g.height));
1361 let mut t = Self::reuse_or_alloc_image(cached_tmp.take(), tmp_w, tmp_h, intermediate)?;
1362 #[cfg(test)]
1363 {
1364 self.last_tmp_dims = Some((tmp_w, tmp_h));
1365 }
1366 match pre_region {
1367 Some(g) => {
1368 let mut sub = Self::reuse_or_alloc_image(
1369 self.convert_src_sub.take(),
1370 g.width,
1371 g.height,
1372 src_fmt,
1373 )?;
1374 {
1375 let _s = tracing::trace_span!(
1376 "image.convert.cpu.extract_region",
1377 region_w = g.width,
1378 region_h = g.height,
1379 )
1380 .entered();
1381 Self::extract_nv_region(src, &mut sub, src_fmt, g)?;
1382 }
1383 Self::convert_format_pf(&sub, &mut t, src_fmt, intermediate, src_params)?;
1384 self.convert_src_sub = Some(sub);
1385 }
1386 None => Self::convert_format_pf(src, &mut t, src_fmt, intermediate, src_params)?,
1387 }
1388 Some(t)
1389 } else {
1390 None
1391 };
1392
1393 let crop = match (pre_region, crop.src_rect) {
1398 (Some(g), Some(r)) => ResolvedCrop {
1399 src_rect: Some(Rect {
1400 left: r.left - g.left,
1401 top: r.top - g.top,
1402 ..r
1403 }),
1404 ..crop
1405 },
1406 _ => crop,
1407 };
1408 let (tmp, tmp_fmt): (&Tensor<u8>, PixelFormat) = match &tmp_holder {
1409 Some(t) => (t, intermediate),
1410 None => (src, src_fmt),
1411 };
1412
1413 debug_assert!(matches!(tmp_fmt, Rgb | Rgba | Grey));
1415 if tmp_fmt == dst_fmt {
1416 let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
1417 self.resize_flip_rotate_pf(tmp, dst, dst_fmt, rotation, flip, crop)?;
1418 } else if !need_resize_flip_rotation {
1419 let _s = tracing::trace_span!(
1420 "image.convert.cpu.format_convert",
1421 from = ?tmp_fmt,
1422 to = ?dst_fmt,
1423 pass = "direct",
1424 )
1425 .entered();
1426 Self::convert_format_pf(tmp, dst, tmp_fmt, dst_fmt, dst_params)?;
1427 } else {
1428 let mut tmp2 = Self::reuse_or_alloc_image(cached_tmp2.take(), dst_w, dst_h, tmp_fmt)?;
1429 if crop.dst_rect.is_some_and(|c| {
1430 c != Rect {
1431 left: 0,
1432 top: 0,
1433 width: dst_w,
1434 height: dst_h,
1435 }
1436 }) && crop.dst_color.is_none()
1437 {
1438 Self::convert_format_pf(dst, &mut tmp2, dst_fmt, tmp_fmt, dst_params)?;
1439 }
1440 {
1441 let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
1442 self.resize_flip_rotate_pf(tmp, &mut tmp2, tmp_fmt, rotation, flip, crop)?;
1443 }
1444 {
1445 let _s = tracing::trace_span!(
1446 "image.convert.cpu.format_convert",
1447 from = ?tmp_fmt,
1448 to = ?dst_fmt,
1449 pass = "post_resize",
1450 )
1451 .entered();
1452 Self::convert_format_pf(&tmp2, dst, tmp_fmt, dst_fmt, dst_params)?;
1453 }
1454 cached_tmp2 = Some(tmp2);
1455 }
1456 if let Some(t) = tmp_holder {
1459 cached_tmp = Some(t);
1460 }
1461 self.convert_tmp = cached_tmp;
1462 self.convert_tmp2 = cached_tmp2;
1463
1464 if let (Some(dst_rect), Some(dst_color)) = (crop.dst_rect, crop.dst_color) {
1465 let full_rect = Rect {
1466 left: 0,
1467 top: 0,
1468 width: dst_w,
1469 height: dst_h,
1470 };
1471 if dst_rect != full_rect {
1472 Self::fill_image_outside_crop_u8(dst, dst_color, dst_rect)?;
1473 }
1474 }
1475
1476 Ok(())
1477 }
1478
1479 fn draw_decoded_masks_impl(
1480 &mut self,
1481 dst: &mut Tensor<u8>,
1482 detect: &[DetectBox],
1483 segmentation: &[Segmentation],
1484 opacity: f32,
1485 color_mode: crate::ColorMode,
1486 ) -> Result<()> {
1487 let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
1488 if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
1489 return Err(crate::Error::NotSupported(
1490 "CPU image rendering only supports RGBA or RGB images".to_string(),
1491 ));
1492 }
1493
1494 let _timer = FunctionTimer::new("CPUProcessor::draw_decoded_masks");
1495
1496 let dst_w = dst.width().unwrap();
1497 let dst_h = dst.height().unwrap();
1498 let dst_rs = tensor_row_stride(dst);
1499 let dst_c = dst_fmt.channels();
1500
1501 let mut map = dst.map_mut()?;
1502 let dst_slice = map.as_mut_slice();
1503
1504 self.render_box(dst_w, dst_h, dst_rs, dst_c, dst_slice, detect, color_mode)?;
1505
1506 if segmentation.is_empty() {
1507 return Ok(());
1508 }
1509
1510 if let Some(bad) = segmentation
1517 .iter()
1518 .find(|s| s.segmentation.shape().len() != 3)
1519 {
1520 return Err(crate::Error::InvalidShape(format!(
1521 "segmentation must be [H, W, C], got {:?}",
1522 bad.segmentation.shape()
1523 )));
1524 }
1525
1526 let is_semantic = segmentation[0].segmentation.shape()[2] > 1;
1529
1530 if is_semantic {
1531 self.render_modelpack_segmentation(
1532 dst_w,
1533 dst_h,
1534 dst_rs,
1535 dst_c,
1536 dst_slice,
1537 &segmentation[0],
1538 opacity,
1539 )?;
1540 } else {
1541 for (idx, (seg, det)) in segmentation.iter().zip(detect).enumerate() {
1542 let color_index = color_mode.index(idx, det.label);
1543 self.render_yolo_segmentation(
1544 dst_w,
1545 dst_h,
1546 dst_rs,
1547 dst_c,
1548 dst_slice,
1549 seg,
1550 color_index,
1551 opacity,
1552 )?;
1553 }
1554 }
1555
1556 Ok(())
1557 }
1558
1559 fn draw_proto_masks_impl(
1560 &mut self,
1561 dst: &mut Tensor<u8>,
1562 detect: &[DetectBox],
1563 proto_data: &ProtoData,
1564 opacity: f32,
1565 letterbox: Option<[f32; 4]>,
1566 color_mode: crate::ColorMode,
1567 ) -> Result<()> {
1568 let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
1569 if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
1570 return Err(crate::Error::NotSupported(
1571 "CPU image rendering only supports RGBA or RGB images".to_string(),
1572 ));
1573 }
1574
1575 let _timer = FunctionTimer::new("CPUProcessor::draw_proto_masks");
1576
1577 let dst_w = dst.width().unwrap();
1578 let dst_h = dst.height().unwrap();
1579 let dst_rs = tensor_row_stride(dst);
1580 let channels = dst_fmt.channels();
1581
1582 let mut map = dst.map_mut()?;
1583 let dst_slice = map.as_mut_slice();
1584
1585 self.render_box(
1586 dst_w, dst_h, dst_rs, channels, dst_slice, detect, color_mode,
1587 )?;
1588
1589 if detect.is_empty() {
1590 return Ok(());
1591 }
1592 let proto_shape = proto_data.protos.shape();
1593 if proto_shape.len() != 3 {
1594 return Err(Error::InvalidShape(format!(
1595 "protos tensor must be rank-3, got {proto_shape:?}"
1596 )));
1597 }
1598 let proto_h = proto_shape[0];
1599 let proto_w = proto_shape[1];
1600 let num_protos = proto_shape[2];
1601 let coeff_shape = proto_data.mask_coefficients.shape();
1602 if coeff_shape.len() != 2 {
1603 return Err(Error::InvalidShape(format!(
1604 "mask_coefficients tensor must be rank-2, got {coeff_shape:?}"
1605 )));
1606 }
1607 if coeff_shape[0] == 0 {
1609 return Ok(());
1610 }
1611 if coeff_shape[1] != num_protos {
1612 return Err(Error::InvalidShape(format!(
1613 "mask_coefficients second dimension must match num_protos \
1614 ({num_protos}), got {coeff_shape:?}"
1615 )));
1616 }
1617
1618 let coeff_f32: Vec<f32> = match proto_data.mask_coefficients.dtype() {
1620 DType::F32 => {
1621 let t = proto_data.mask_coefficients.as_typed::<f32>().expect("F32");
1622 let m = t.map_read()?;
1623 m.as_slice().to_vec()
1624 }
1625 DType::F16 => {
1626 let t = proto_data
1627 .mask_coefficients
1628 .as_typed::<half::f16>()
1629 .expect("F16");
1630 let m = t.map_read()?;
1631 m.as_slice().iter().map(|v| v.to_f32()).collect()
1632 }
1633 DType::I8 => {
1634 let t = proto_data.mask_coefficients.as_typed::<i8>().expect("I8");
1635 let m = t.map_read()?;
1636 if let Some(q) = t.quantization() {
1637 use edgefirst_tensor::QuantMode;
1638 let (scale, zp) = match q.mode() {
1639 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1640 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1641 other => {
1642 return Err(Error::NotSupported(format!(
1643 "I8 mask_coefficients quantization mode {other:?} not supported"
1644 )));
1645 }
1646 };
1647 m.as_slice()
1648 .iter()
1649 .map(|&v| (v as f32 - zp) * scale)
1650 .collect()
1651 } else {
1652 m.as_slice().iter().map(|&v| v as f32).collect()
1653 }
1654 }
1655 DType::I16 => {
1656 let t = proto_data.mask_coefficients.as_typed::<i16>().expect("I16");
1657 let m = t.map_read()?;
1658 if let Some(q) = t.quantization() {
1659 use edgefirst_tensor::QuantMode;
1660 let (scale, zp) = match q.mode() {
1661 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1662 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1663 other => {
1664 return Err(Error::NotSupported(format!(
1665 "I16 mask_coefficients quantization mode {other:?} not supported"
1666 )));
1667 }
1668 };
1669 m.as_slice()
1670 .iter()
1671 .map(|&v| (v as f32 - zp) * scale)
1672 .collect()
1673 } else {
1674 m.as_slice().iter().map(|&v| v as f32).collect()
1675 }
1676 }
1677 other => {
1678 return Err(Error::InvalidShape(format!(
1679 "mask_coefficients dtype {other:?} not supported"
1680 )));
1681 }
1682 };
1683
1684 let (lx0, lx_range, ly0, ly_range) = match letterbox {
1686 Some([lx0, ly0, lx1, ly1]) => (lx0, lx1 - lx0, ly0, ly1 - ly0),
1687 None => (0.0_f32, 1.0_f32, 0.0_f32, 1.0_f32),
1688 };
1689
1690 match proto_data.protos.dtype() {
1693 DType::F32 => {
1694 let t = proto_data.protos.as_typed::<f32>().expect("F32");
1695 let m = t.map_read()?;
1696 self.draw_proto_masks_inner(
1697 dst_slice,
1698 dst_w,
1699 dst_h,
1700 dst_rs,
1701 channels,
1702 detect,
1703 m.as_slice(),
1704 &coeff_f32,
1705 proto_h,
1706 proto_w,
1707 num_protos,
1708 opacity,
1709 (lx0, lx_range, ly0, ly_range),
1710 color_mode,
1711 0.0_f32,
1712 |p: &f32, _| *p,
1713 );
1714 }
1715 DType::F16 => {
1716 let t = proto_data.protos.as_typed::<half::f16>().expect("F16");
1717 let m = t.map_read()?;
1718 self.draw_proto_masks_inner(
1719 dst_slice,
1720 dst_w,
1721 dst_h,
1722 dst_rs,
1723 channels,
1724 detect,
1725 m.as_slice(),
1726 &coeff_f32,
1727 proto_h,
1728 proto_w,
1729 num_protos,
1730 opacity,
1731 (lx0, lx_range, ly0, ly_range),
1732 color_mode,
1733 0.0_f32,
1734 |p: &half::f16, _| p.to_f32(),
1735 );
1736 }
1737 DType::I8 => {
1738 use edgefirst_tensor::QuantMode;
1739 let t = proto_data.protos.as_typed::<i8>().expect("I8");
1740 let m = t.map_read()?;
1741 let quant = t.quantization().ok_or_else(|| {
1742 Error::InvalidShape("I8 protos require quantization metadata".into())
1743 })?;
1744 let (scale, zp) = match quant.mode() {
1745 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1746 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1747 QuantMode::PerChannel { axis, .. }
1748 | QuantMode::PerChannelSymmetric { axis, .. } => {
1749 return Err(Error::NotSupported(format!(
1750 "per-channel quantization (axis={axis}) in draw_proto_masks \
1751 CPU path not yet supported"
1752 )));
1753 }
1754 };
1755 self.draw_proto_masks_inner(
1756 dst_slice,
1757 dst_w,
1758 dst_h,
1759 dst_rs,
1760 channels,
1761 detect,
1762 m.as_slice(),
1763 &coeff_f32,
1764 proto_h,
1765 proto_w,
1766 num_protos,
1767 opacity,
1768 (lx0, lx_range, ly0, ly_range),
1769 color_mode,
1770 scale,
1771 move |p: &i8, _| (*p as f32) - zp,
1772 );
1773 }
1774 other => {
1775 return Err(Error::InvalidShape(format!(
1776 "proto tensor dtype {other:?} not supported"
1777 )));
1778 }
1779 }
1780
1781 Ok(())
1782 }
1783
1784 #[allow(clippy::too_many_arguments)]
1785 fn draw_proto_masks_inner<P: Copy>(
1786 &self,
1787 dst_slice: &mut [u8],
1788 dst_w: usize,
1789 dst_h: usize,
1790 dst_rs: usize,
1791 channels: usize,
1792 detect: &[DetectBox],
1793 protos: &[P],
1794 coeff_all_f32: &[f32],
1795 proto_h: usize,
1796 proto_w: usize,
1797 num_protos: usize,
1798 opacity: f32,
1799 letterbox_xy: (f32, f32, f32, f32),
1800 color_mode: crate::ColorMode,
1801 acc_scale: f32,
1802 load_f32: impl Fn(&P, f32) -> f32 + Copy,
1803 ) {
1804 let (lx0, lx_range, ly0, ly_range) = letterbox_xy;
1805 let stride_y = proto_w * num_protos;
1806 for (idx, det) in detect.iter().enumerate() {
1807 let coeff = &coeff_all_f32[idx * num_protos..(idx + 1) * num_protos];
1808 let color_index = color_mode.index(idx, det.label);
1809 let color = self.colors[color_index % self.colors.len()];
1810 let alpha = if opacity == 1.0 {
1811 color[3] as u16
1812 } else {
1813 (color[3] as f32 * opacity).round() as u16
1814 };
1815
1816 let start_x = (dst_w as f32 * det.bbox.xmin).round() as usize;
1817 let start_y = (dst_h as f32 * det.bbox.ymin).round() as usize;
1818 let end_x = ((dst_w as f32 * det.bbox.xmax).round() as usize).min(dst_w);
1819 let end_y = ((dst_h as f32 * det.bbox.ymax).round() as usize).min(dst_h);
1820
1821 for y in start_y..end_y {
1822 for x in start_x..end_x {
1823 let px = (lx0 + (x as f32 / dst_w as f32) * lx_range) * proto_w as f32 - 0.5;
1824 let py = (ly0 + (y as f32 / dst_h as f32) * ly_range) * proto_h as f32 - 0.5;
1825
1826 let x0 = (px.floor() as isize).clamp(0, proto_w as isize - 1) as usize;
1830 let y0 = (py.floor() as isize).clamp(0, proto_h as isize - 1) as usize;
1831 let x1 = (x0 + 1).min(proto_w - 1);
1832 let y1 = (y0 + 1).min(proto_h - 1);
1833 let fx = px - px.floor();
1834 let fy = py - py.floor();
1835 let w00 = (1.0 - fx) * (1.0 - fy);
1836 let w10 = fx * (1.0 - fy);
1837 let w01 = (1.0 - fx) * fy;
1838 let w11 = fx * fy;
1839 let b00 = y0 * stride_y + x0 * num_protos;
1840 let b10 = y0 * stride_y + x1 * num_protos;
1841 let b01 = y1 * stride_y + x0 * num_protos;
1842 let b11 = y1 * stride_y + x1 * num_protos;
1843 let mut acc = 0.0_f32;
1844 for p in 0..num_protos {
1845 let v00 = load_f32(&protos[b00 + p], 0.0);
1846 let v10 = load_f32(&protos[b10 + p], 0.0);
1847 let v01 = load_f32(&protos[b01 + p], 0.0);
1848 let v11 = load_f32(&protos[b11 + p], 0.0);
1849 let val = w00 * v00 + w10 * v10 + w01 * v01 + w11 * v11;
1850 acc += coeff[p] * val;
1851 }
1852 let final_acc = if acc_scale == 0.0 {
1853 acc
1854 } else {
1855 acc_scale * acc
1856 };
1857 let mask = 1.0 / (1.0 + (-final_acc).exp());
1861 if mask < 0.5 {
1862 continue;
1863 }
1864 let dst_index = y * dst_rs + x * channels;
1865 for c in 0..3 {
1866 dst_slice[dst_index + c] = ((color[c] as u16 * alpha
1867 + dst_slice[dst_index + c] as u16 * (255 - alpha))
1868 / 255) as u8;
1869 }
1870 }
1871 }
1872 }
1873 }
1874}