1use crate::{
5 Crop, Error, Flip, FunctionTimer, ImageProcessorTrait, Rect, ResolvedCrop, Result, Rotation,
6};
7use edgefirst_decoder::{DetectBox, ProtoData, Segmentation};
8use edgefirst_tensor::{
9 DType, PixelFormat, Tensor, TensorDyn, TensorMapTrait, TensorMemory, TensorTrait,
10};
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_u8().ok_or(Error::NotAnImage)?;
197 let dst_u8 = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
198 let bg_map = bg_u8.map_read()?;
199 let mut dst_map = dst_u8.map_mut()?;
200 let bg_slice = bg_map.as_slice();
201 let dst_slice = dst_map.as_mut_slice();
202 if bg_slice.len() != dst_slice.len() {
203 return Err(Error::InvalidShape(
204 "background buffer size does not match dst".into(),
205 ));
206 }
207 dst_slice.copy_from_slice(bg_slice);
208 }
209 None => {
210 let dst_u8 = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
211 let mut dst_map = dst_u8.map_mut()?;
212 dst_map.as_mut_slice().fill(0);
213 }
214 }
215 Ok(())
216}
217
218fn chroma_alignment_ok(fmt: PixelFormat, r: Rect) -> bool {
226 match fmt {
227 PixelFormat::Nv12 => r.left.is_multiple_of(2) && r.top.is_multiple_of(2),
228 PixelFormat::Nv16 => r.left.is_multiple_of(2),
229 PixelFormat::Nv24 => true,
230 _ => false,
231 }
232}
233
234fn row_stride_for(width: usize, fmt: PixelFormat) -> usize {
236 use edgefirst_tensor::PixelLayout;
237 match fmt.layout() {
238 PixelLayout::Packed => width * fmt.channels(),
239 PixelLayout::Planar | PixelLayout::SemiPlanar => width,
240 _ => width, }
242}
243
244fn tensor_row_stride(tensor: &Tensor<u8>) -> usize {
249 tensor.effective_row_stride().unwrap_or_else(|| {
250 let w = tensor.width().unwrap_or(0);
251 let fmt = tensor.format().unwrap_or(PixelFormat::Rgb);
252 row_stride_for(w, fmt)
253 })
254}
255
256fn logical_surface(tensor: &Tensor<u8>) -> Result<(usize, usize)> {
266 use edgefirst_tensor::PixelLayout;
267 let fmt = tensor.format().ok_or(Error::NotAnImage)?;
268 let w = tensor.width().ok_or(Error::NotAnImage)?;
269 let h = tensor.height().ok_or(Error::NotAnImage)?;
270 Ok(match fmt.layout() {
271 PixelLayout::Packed => (h, w * fmt.channels()),
272 PixelLayout::Planar => (fmt.channels() * h, w),
273 PixelLayout::SemiPlanar => (fmt.combined_plane_height(h).unwrap_or(h), w),
274 _ => (h, row_stride_for(w, fmt)),
277 })
278}
279
280fn packed_row_pairs<'s, 'd>(
286 src: &'s [u8],
287 src_stride: usize,
288 src_row_bytes: usize,
289 dst: &'d mut [u8],
290 dst_stride: usize,
291 dst_row_bytes: usize,
292 rows: usize,
293) -> impl Iterator<Item = (&'s [u8], &'d mut [u8])> {
294 src.chunks(src_stride)
295 .zip(dst.chunks_mut(dst_stride))
296 .take(rows)
297 .map(move |(s, d)| (&s[..src_row_bytes], &mut d[..dst_row_bytes]))
298}
299
300fn split_semi_planar(
306 bytes: &[u8],
307 stride: usize,
308 src_h: usize,
309 fmt: PixelFormat,
310) -> Result<(&[u8], &[u8])> {
311 let total_h = fmt.combined_plane_height(src_h).unwrap_or(src_h);
312 let need = stride.checked_mul(total_h).ok_or_else(|| {
313 Error::InvalidShape(format!(
314 "{fmt:?} plane size overflow (stride={stride}, h={src_h})"
315 ))
316 })?;
317 if bytes.len() < need {
318 return Err(Error::InvalidShape(format!(
319 "{fmt:?} source has {} bytes but needs {need} (stride={stride}, h={src_h})",
320 bytes.len()
321 )));
322 }
323 Ok(bytes.split_at(stride * src_h))
324}
325
326fn split_semi_planar_mut(
333 bytes: &mut [u8],
334 stride: usize,
335 dst_h: usize,
336 fmt: PixelFormat,
337) -> Result<(&mut [u8], &mut [u8])> {
338 let total_h = fmt.combined_plane_height(dst_h).unwrap_or(dst_h);
339 let need = stride.checked_mul(total_h).ok_or_else(|| {
340 Error::InvalidShape(format!(
341 "{fmt:?} plane size overflow (stride={stride}, combined_h={total_h})"
342 ))
343 })?;
344 if bytes.len() < need {
345 return Err(Error::InvalidShape(format!(
346 "{fmt:?} destination has {} bytes but needs {need} (stride={stride}, combined_h={total_h})",
347 bytes.len()
348 )));
349 }
350 Ok(bytes.split_at_mut(stride * dst_h))
351}
352
353fn guard_plane(
359 buf_len: usize,
360 stride: usize,
361 rows: usize,
362 row_bytes: usize,
363 what: &str,
364) -> Result<()> {
365 let need = stride.checked_mul(rows).ok_or_else(|| {
366 Error::InvalidShape(format!(
367 "{what} plane size overflow (stride={stride}, rows={rows})"
368 ))
369 })?;
370 if row_bytes > stride || buf_len < need {
371 return Err(Error::InvalidShape(format!(
372 "{what} buffer too small: {buf_len} bytes, need {need} (stride={stride}, rows={rows}, row_bytes={row_bytes})"
373 )));
374 }
375 Ok(())
376}
377
378pub(crate) fn apply_int8_xor_bias(data: &mut [u8], fmt: PixelFormat) {
383 use edgefirst_tensor::PixelLayout;
384 if !fmt.has_alpha() {
385 for b in data.iter_mut() {
386 *b ^= 0x80;
387 }
388 } else if fmt.layout() == PixelLayout::Planar {
389 let channels = fmt.channels();
391 let plane_size = data.len() / channels;
392 for b in data[..plane_size * (channels - 1)].iter_mut() {
393 *b ^= 0x80;
394 }
395 } else {
396 let channels = fmt.channels();
398 for pixel in data.chunks_exact_mut(channels) {
399 for b in &mut pixel[..channels - 1] {
400 *b ^= 0x80;
401 }
402 }
403 }
404}
405
406fn apply_int8_xor_bias_rows(tensor: &mut Tensor<u8>, fmt: PixelFormat) -> Result<()> {
411 use edgefirst_tensor::PixelLayout;
412 let (rows, row_bytes) = logical_surface(tensor)?;
413 let stride = tensor_row_stride(tensor);
414 let color_rows = if fmt.has_alpha() && fmt.layout() == PixelLayout::Planar {
417 rows / fmt.channels() * (fmt.channels() - 1)
418 } else {
419 rows
420 };
421 let row_fmt = if fmt.layout() == PixelLayout::Planar {
424 PixelFormat::Grey
425 } else {
426 fmt
427 };
428 let mut map = tensor.map_mut()?;
429 let buf = map.as_mut_slice();
430 guard_plane(buf.len(), stride, rows, row_bytes, "int8 bias dst")?;
431 for row in buf.chunks_mut(stride).take(color_rows) {
432 apply_int8_xor_bias(&mut row[..row_bytes], row_fmt);
433 }
434 Ok(())
435}
436
437impl CPUProcessor {
438 pub fn new() -> Self {
440 Self::new_bilinear()
441 }
442
443 fn new_bilinear() -> Self {
445 let resizer = fast_image_resize::Resizer::new();
446 let options = fast_image_resize::ResizeOptions::new()
447 .resize_alg(fast_image_resize::ResizeAlg::Convolution(
448 fast_image_resize::FilterType::Bilinear,
449 ))
450 .use_alpha(false);
451
452 log::debug!("CPUConverter created");
453 Self {
454 resizer,
455 options,
456 colors: crate::DEFAULT_COLORS_U8,
457 widen_scratch: None,
458 resize_destride_scratch: Vec::new(),
459 resize_dst_destride_scratch: Vec::new(),
460 nv_strip_scratch: Vec::new(),
461 nv_strip_y_pack: Vec::new(),
462 nv_strip_uv_pack: Vec::new(),
463 #[cfg(test)]
464 fused_hits: 0,
465 #[cfg(test)]
466 last_tmp_dims: None,
467 convert_tmp: None,
468 convert_tmp2: None,
469 convert_src_sub: None,
470 }
471 }
472
473 pub fn new_nearest() -> Self {
475 let resizer = fast_image_resize::Resizer::new();
476 let options = fast_image_resize::ResizeOptions::new()
477 .resize_alg(fast_image_resize::ResizeAlg::Nearest)
478 .use_alpha(false);
479 log::debug!("CPUConverter created");
480 Self {
481 resizer,
482 options,
483 colors: crate::DEFAULT_COLORS_U8,
484 widen_scratch: None,
485 resize_destride_scratch: Vec::new(),
486 resize_dst_destride_scratch: Vec::new(),
487 nv_strip_scratch: Vec::new(),
488 nv_strip_y_pack: Vec::new(),
489 nv_strip_uv_pack: Vec::new(),
490 #[cfg(test)]
491 fused_hits: 0,
492 #[cfg(test)]
493 last_tmp_dims: None,
494 convert_tmp: None,
495 convert_tmp2: None,
496 convert_src_sub: None,
497 }
498 }
499
500 #[cfg(test)]
502 pub(super) fn fused_hits(&self) -> u64 {
503 self.fused_hits
504 }
505
506 #[cfg(test)]
509 pub(super) fn last_tmp_dims(&self) -> Option<(usize, usize)> {
510 self.last_tmp_dims
511 }
512
513 pub(crate) fn support_conversion_pf(src: PixelFormat, dst: PixelFormat) -> bool {
514 use PixelFormat::*;
515 matches!(
516 (src, dst),
517 (Nv12, Rgb)
518 | (Nv12, Rgba)
519 | (Nv12, Grey)
520 | (Nv16, Rgb)
521 | (Nv16, Rgba)
522 | (Nv16, Bgra)
523 | (Nv24, Rgb)
524 | (Nv24, Rgba)
525 | (Nv24, Grey)
526 | (Nv24, Bgra)
527 | (Yuyv, Rgb)
528 | (Yuyv, Rgba)
529 | (Yuyv, Grey)
530 | (Yuyv, Yuyv)
531 | (Yuyv, PlanarRgb)
532 | (Yuyv, PlanarRgba)
533 | (Yuyv, Nv16)
534 | (Vyuy, Rgb)
535 | (Vyuy, Rgba)
536 | (Vyuy, Grey)
537 | (Vyuy, Vyuy)
538 | (Vyuy, PlanarRgb)
539 | (Vyuy, PlanarRgba)
540 | (Vyuy, Nv16)
541 | (Rgba, Rgb)
542 | (Rgba, Rgba)
543 | (Rgba, Grey)
544 | (Rgba, Yuyv)
545 | (Rgba, PlanarRgb)
546 | (Rgba, PlanarRgba)
547 | (Rgba, Nv16)
548 | (Rgb, Rgb)
549 | (Rgb, Rgba)
550 | (Rgb, Grey)
551 | (Rgb, Yuyv)
552 | (Rgb, PlanarRgb)
553 | (Rgb, PlanarRgba)
554 | (Rgb, Nv16)
555 | (Grey, Rgb)
556 | (Grey, Rgba)
557 | (Grey, Grey)
558 | (Grey, Yuyv)
559 | (Grey, PlanarRgb)
560 | (Grey, PlanarRgba)
561 | (Grey, Nv16)
562 | (Nv12, Bgra)
563 | (Yuyv, Bgra)
564 | (Vyuy, Bgra)
565 | (Rgba, Bgra)
566 | (Rgb, Bgra)
567 | (Grey, Bgra)
568 | (Bgra, Bgra)
569 | (PlanarRgb, Rgb)
570 | (PlanarRgb, Rgba)
571 | (PlanarRgba, Rgb)
572 | (PlanarRgba, Rgba)
573 | (PlanarRgb, Bgra)
574 | (PlanarRgba, Bgra)
575 )
576 }
577
578 pub(crate) fn convert_format_pf(
580 src: &Tensor<u8>,
581 dst: &mut Tensor<u8>,
582 src_fmt: PixelFormat,
583 dst_fmt: PixelFormat,
584 cp: ColorParams,
585 ) -> Result<()> {
586 let _timer = FunctionTimer::new(format!(
587 "ImageProcessor::convert_format {} to {}",
588 src_fmt, dst_fmt,
589 ));
590
591 use PixelFormat::*;
592 match (src_fmt, dst_fmt) {
593 (Nv12, Rgb) => Self::convert_nv12_to_rgb(src, dst, cp),
594 (Nv12, Rgba) => Self::convert_nv12_to_rgba(src, dst, cp),
595 (Nv12, Grey) => Self::convert_nv12_to_grey(src, dst, cp),
596 (Yuyv, Rgb) => Self::convert_yuyv_to_rgb(src, dst, cp),
597 (Yuyv, Rgba) => Self::convert_yuyv_to_rgba(src, dst, cp),
598 (Yuyv, Grey) => Self::convert_yuyv_to_grey(src, dst, cp),
599 (Yuyv, Yuyv) => Self::copy_image(src, dst),
600 (Yuyv, PlanarRgb) => Self::convert_yuyv_to_8bps(src, dst, cp),
601 (Yuyv, PlanarRgba) => Self::convert_yuyv_to_prgba(src, dst, cp),
602 (Yuyv, Nv16) => Self::convert_yuyv_to_nv16(src, dst),
603 (Vyuy, Rgb) => Self::convert_vyuy_to_rgb(src, dst, cp),
604 (Vyuy, Rgba) => Self::convert_vyuy_to_rgba(src, dst, cp),
605 (Vyuy, Grey) => Self::convert_vyuy_to_grey(src, dst, cp),
606 (Vyuy, Vyuy) => Self::copy_image(src, dst),
607 (Vyuy, PlanarRgb) => Self::convert_vyuy_to_8bps(src, dst, cp),
608 (Vyuy, PlanarRgba) => Self::convert_vyuy_to_prgba(src, dst, cp),
609 (Vyuy, Nv16) => Self::convert_vyuy_to_nv16(src, dst),
610 (Rgba, Rgb) => Self::convert_rgba_to_rgb(src, dst),
611 (Rgba, Rgba) => Self::copy_image(src, dst),
612 (Rgba, Grey) => Self::convert_rgba_to_grey(src, dst),
613 (Rgba, Yuyv) => Self::convert_rgba_to_yuyv(src, dst, cp),
614 (Rgba, PlanarRgb) => Self::convert_rgba_to_8bps(src, dst),
615 (Rgba, PlanarRgba) => Self::convert_rgba_to_prgba(src, dst),
616 (Rgba, Nv16) => Self::convert_rgba_to_nv16(src, dst, cp),
617 (Rgb, Rgb) => Self::copy_image(src, dst),
618 (Rgb, Rgba) => Self::convert_rgb_to_rgba(src, dst),
619 (Rgb, Grey) => Self::convert_rgb_to_grey(src, dst),
620 (Rgb, Yuyv) => Self::convert_rgb_to_yuyv(src, dst, cp),
621 (Rgb, PlanarRgb) => Self::convert_rgb_to_8bps(src, dst),
622 (Rgb, PlanarRgba) => Self::convert_rgb_to_prgba(src, dst),
623 (Rgb, Nv16) => Self::convert_rgb_to_nv16(src, dst, cp),
624 (Grey, Rgb) => Self::convert_grey_to_rgb(src, dst),
625 (Grey, Rgba) => Self::convert_grey_to_rgba(src, dst),
626 (Grey, Grey) => Self::copy_image(src, dst),
627 (Grey, Yuyv) => Self::convert_grey_to_yuyv(src, dst, cp),
628 (Grey, PlanarRgb) => Self::convert_grey_to_8bps(src, dst),
629 (Grey, PlanarRgba) => Self::convert_grey_to_prgba(src, dst),
630 (Grey, Nv16) => Self::convert_grey_to_nv16(src, dst, cp),
631
632 (Nv16, Rgb) => Self::convert_nv16_to_rgb(src, dst, cp),
634 (Nv16, Rgba) => Self::convert_nv16_to_rgba(src, dst, cp),
635 (Nv24, Rgb) => Self::convert_nv24_to_rgb(src, dst, cp),
636 (Nv24, Rgba) => Self::convert_nv24_to_rgba(src, dst, cp),
637 (Nv24, Grey) => Self::convert_nv24_to_grey(src, dst, cp),
638 (PlanarRgb, Rgb) => Self::convert_8bps_to_rgb(src, dst),
639 (PlanarRgb, Rgba) => Self::convert_8bps_to_rgba(src, dst),
640 (PlanarRgba, Rgb) => Self::convert_prgba_to_rgb(src, dst),
641 (PlanarRgba, Rgba) => Self::convert_prgba_to_rgba(src, dst),
642
643 (Bgra, Bgra) => Self::copy_image(src, dst),
645 (Nv12, Bgra) => {
646 Self::convert_nv12_to_rgba(src, dst, cp)?;
647 Self::swizzle_rb_4chan(dst)
648 }
649 (Nv16, Bgra) => {
650 Self::convert_nv16_to_rgba(src, dst, cp)?;
651 Self::swizzle_rb_4chan(dst)
652 }
653 (Nv24, Bgra) => {
654 Self::convert_nv24_to_rgba(src, dst, cp)?;
655 Self::swizzle_rb_4chan(dst)
656 }
657 (Yuyv, Bgra) => {
658 Self::convert_yuyv_to_rgba(src, dst, cp)?;
659 Self::swizzle_rb_4chan(dst)
660 }
661 (Vyuy, Bgra) => {
662 Self::convert_vyuy_to_rgba(src, dst, cp)?;
663 Self::swizzle_rb_4chan(dst)
664 }
665 (Rgba, Bgra) => {
666 Self::copy_image(src, dst)?;
667 Self::swizzle_rb_4chan(dst)
668 }
669 (Rgb, Bgra) => {
670 Self::convert_rgb_to_rgba(src, dst)?;
671 Self::swizzle_rb_4chan(dst)
672 }
673 (Grey, Bgra) => {
674 Self::convert_grey_to_rgba(src, dst)?;
675 Self::swizzle_rb_4chan(dst)
676 }
677 (PlanarRgb, Bgra) => {
678 Self::convert_8bps_to_rgba(src, dst)?;
679 Self::swizzle_rb_4chan(dst)
680 }
681 (PlanarRgba, Bgra) => {
682 Self::convert_prgba_to_rgba(src, dst)?;
683 Self::swizzle_rb_4chan(dst)
684 }
685
686 (s, d) => Err(Error::NotSupported(format!("Conversion from {s} to {d}",))),
687 }
688 }
689
690 pub(crate) fn fill_image_outside_crop_u8(
692 dst: &mut Tensor<u8>,
693 rgba: [u8; 4],
694 crop: Rect,
695 ) -> Result<()> {
696 let dst_fmt = dst.format().unwrap();
697 let dst_w = dst.width().unwrap();
698 let dst_h = dst.height().unwrap();
699 let cm = crate::colorimetry::resolve_colorimetry(dst.colorimetry(), dst.height());
703 let cp = ColorParams {
704 matrix: crate::colorimetry::yuv_matrix(cm.encoding.unwrap()),
705 range: crate::colorimetry::yuv_range(cm.range.unwrap()),
706 encoding: cm.encoding.unwrap(),
707 range_kind: cm.range.unwrap(),
708 src_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
709 dst_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
710 };
711 let dst_stride = tensor_row_stride(dst);
712 let mut dst_map = dst.map_mut()?;
713 let dst_tup = (dst_map.as_mut_slice(), dst_w, dst_h, dst_stride);
714 Self::fill_outside_crop_dispatch(dst_tup, dst_fmt, rgba, crop, cp)
715 }
716
717 fn fill_outside_crop_dispatch(
722 dst: (&mut [u8], usize, usize, usize),
723 fmt: PixelFormat,
724 rgba: [u8; 4],
725 crop: Rect,
726 cp: ColorParams,
727 ) -> Result<()> {
728 use PixelFormat::*;
729 match fmt {
730 Rgba | Bgra => Self::fill_image_outside_crop_(dst, rgba, crop),
731 Rgb => Self::fill_image_outside_crop_(dst, Self::rgba_to_rgb(rgba), crop),
732 Grey => Self::fill_image_outside_crop_(dst, Self::rgba_to_grey(rgba), crop),
733 Yuyv => {
734 let (bytes, w, h, stride) = dst;
735 let yuyv = Self::rgba_to_yuyv(rgba, cp);
736 Self::fill_image_outside_crop_(
739 (&mut *bytes, w / 2, h, stride),
740 yuyv,
741 Rect::new(crop.left / 2, crop.top, crop.width.div_ceil(2), crop.height),
742 )?;
743 if w % 2 == 1 {
751 let last = w - 1;
752 let col_outside = last < crop.left || last >= crop.left + crop.width;
753 for y in 0..h {
754 let row_outside = y < crop.top || y >= crop.top + crop.height;
755 if row_outside || col_outside {
756 let off = y * stride + last * 2;
757 bytes[off] = yuyv[0];
758 bytes[off + 1] = yuyv[1];
759 }
760 }
761 }
762 Ok(())
763 }
764 PlanarRgb => Self::fill_image_outside_crop_planar(dst, Self::rgba_to_rgb(rgba), crop),
765 PlanarRgba => Self::fill_image_outside_crop_planar(dst, rgba, crop),
766 Nv16 => {
767 let yuyv = Self::rgba_to_yuyv(rgba, cp);
768 Self::fill_image_outside_crop_yuv_semiplanar(dst, yuyv[0], [yuyv[1], yuyv[3]], crop)
769 }
770 _ => Err(Error::Internal(format!(
771 "Found unexpected destination {fmt}",
772 ))),
773 }
774 }
775}
776
777impl ImageProcessorTrait for CPUProcessor {
778 fn convert(
779 &mut self,
780 src: &TensorDyn,
781 dst: &mut TensorDyn,
782 rotation: Rotation,
783 flip: Flip,
784 crop: Crop,
785 ) -> Result<()> {
786 let crop = crop.resolve(
787 src.width().unwrap_or(0),
788 src.height().unwrap_or(0),
789 dst.width().unwrap_or(0),
790 dst.height().unwrap_or(0),
791 )?;
792 self.convert_impl(src, dst, rotation, flip, crop)
793 }
794
795 fn draw_decoded_masks(
796 &mut self,
797 dst: &mut TensorDyn,
798 detect: &[DetectBox],
799 segmentation: &[Segmentation],
800 overlay: crate::MaskOverlay<'_>,
801 ) -> Result<()> {
802 prepare_dst_base_cpu(dst, overlay.background)?;
806 let dst = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
807 self.draw_decoded_masks_impl(
808 dst,
809 detect,
810 segmentation,
811 overlay.opacity,
812 overlay.color_mode,
813 )
814 }
815
816 fn draw_proto_masks(
817 &mut self,
818 dst: &mut TensorDyn,
819 detect: &[DetectBox],
820 proto_data: &ProtoData,
821 overlay: crate::MaskOverlay<'_>,
822 ) -> Result<()> {
823 prepare_dst_base_cpu(dst, overlay.background)?;
824 let dst = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
825 self.draw_proto_masks_impl(
826 dst,
827 detect,
828 proto_data,
829 overlay.opacity,
830 overlay.letterbox,
831 overlay.color_mode,
832 )
833 }
834
835 fn set_class_colors(&mut self, colors: &[[u8; 4]]) -> Result<()> {
836 for (c, new_c) in self.colors.iter_mut().zip(colors.iter()) {
837 *c = *new_c;
838 }
839 Ok(())
840 }
841}
842
843impl CPUProcessor {
845 pub(crate) fn convert_impl(
847 &mut self,
848 src: &TensorDyn,
849 dst: &mut TensorDyn,
850 rotation: Rotation,
851 flip: Flip,
852 crop: ResolvedCrop,
853 ) -> Result<()> {
854 let src_fmt = src.format().ok_or(Error::NotAnImage)?;
855 let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
856
857 let src_cm = crate::colorimetry::effective_colorimetry(src);
862 let dst_cm = crate::colorimetry::effective_colorimetry(dst);
863 let src_full = src_cm.range == Some(edgefirst_tensor::ColorRange::Full);
864 let dst_full = dst_cm.range == Some(edgefirst_tensor::ColorRange::Full);
865 let src_params = ColorParams {
866 matrix: crate::colorimetry::yuv_matrix(src_cm.encoding.unwrap()),
867 range: crate::colorimetry::yuv_range(src_cm.range.unwrap()),
868 encoding: src_cm.encoding.unwrap(),
869 range_kind: src_cm.range.unwrap(),
870 src_full_range: src_full,
871 dst_full_range: dst_full,
872 };
873 let dst_params = ColorParams {
874 matrix: crate::colorimetry::yuv_matrix(dst_cm.encoding.unwrap()),
875 range: crate::colorimetry::yuv_range(dst_cm.range.unwrap()),
876 encoding: dst_cm.encoding.unwrap(),
877 range_kind: dst_cm.range.unwrap(),
878 src_full_range: src_full,
879 dst_full_range: dst_full,
880 };
881 match (src.dtype(), dst.dtype()) {
882 (DType::U8, DType::U8) => {
883 let src = src.as_u8().unwrap();
884 let dst = dst.as_u8_mut().unwrap();
885 self.convert_u8(
886 src, dst, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
887 )
888 }
889 (DType::U8, DType::I8) => {
890 let src_u8 = src.as_u8().unwrap();
893 let dst_i8 = dst.as_i8_mut().unwrap();
894 let dst_u8 = unsafe { &mut *(dst_i8 as *mut Tensor<i8> as *mut Tensor<u8>) };
898 self.convert_u8(
899 src_u8, dst_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
900 )?;
901 apply_int8_xor_bias_rows(dst_u8, dst_fmt)
903 }
904 (DType::U8, d @ (DType::F32 | DType::F16)) => {
905 let src_u8 = src.as_u8().unwrap();
906 let dw = dst.width().ok_or(Error::NotAnImage)?;
907 let dh = dst.height().ok_or(Error::NotAnImage)?;
908 let scratch_matches = self.widen_scratch.as_ref().is_some_and(|t| {
913 t.width() == Some(dw) && t.height() == Some(dh) && t.format() == Some(dst_fmt)
914 });
915 let mut tmp = if scratch_matches {
916 self.widen_scratch.take().unwrap()
917 } else {
918 TensorDyn::image(
919 dw,
920 dh,
921 dst_fmt,
922 DType::U8,
923 Some(TensorMemory::Mem),
924 edgefirst_tensor::CpuAccess::ReadWrite,
925 )?
926 };
927 {
928 let tmp_u8 = tmp.as_u8_mut().unwrap();
929 self.convert_u8(
930 src_u8, tmp_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params,
931 dst_params,
932 )?;
933 }
934 {
944 let tmp_u8 = tmp.as_u8().unwrap();
945 let (rows, row_len) = logical_surface(tmp_u8)?;
946 let src_stride = tensor_row_stride(tmp_u8);
947 let dst_stride_bytes = dst.effective_row_stride().ok_or(Error::NotAnImage)?;
948 let src_map = tmp_u8.map_read()?;
949 guard_plane(
950 src_map.as_slice().len(),
951 src_stride,
952 rows,
953 row_len,
954 "widen src",
955 )?;
956 let src_rows = src_map.as_slice().chunks(src_stride).take(rows);
957 let elem = d.size();
958 if !dst_stride_bytes.is_multiple_of(elem) {
959 return Err(Error::InvalidShape(format!(
960 "{d} destination row stride {dst_stride_bytes} is not a multiple of \
961 the element size {elem}"
962 )));
963 }
964 let dst_stride = dst_stride_bytes / elem;
965 match d {
966 DType::F32 => {
967 let dst_t = dst.as_f32_mut().unwrap();
968 let mut dst_map = dst_t.map_mut()?;
969 guard_plane(
971 dst_map.as_slice().len(),
972 dst_stride,
973 rows,
974 row_len,
975 "widen f32 dst",
976 )?;
977 for (s, dr) in
981 src_rows.zip(dst_map.as_mut_slice().chunks_mut(dst_stride))
982 {
983 simd::widen_u8_to_f32_norm(s, &mut dr[..row_len]);
984 }
985 }
986 DType::F16 => {
987 let dst_t = dst.as_f16_mut().unwrap();
988 let mut dst_map = dst_t.map_mut()?;
989 guard_plane(
990 dst_map.as_slice().len(),
991 dst_stride,
992 rows,
993 row_len,
994 "widen f16 dst",
995 )?;
996 for (s, dr) in
1001 src_rows.zip(dst_map.as_mut_slice().chunks_mut(dst_stride))
1002 {
1003 simd::widen_u8_to_f16_norm(s, &mut dr[..row_len]);
1004 }
1005 }
1006 _ => unreachable!(),
1007 }
1008 }
1009 self.widen_scratch = Some(tmp);
1010 Ok(())
1011 }
1012 (s, d) => Err(Error::NotSupported(format!("dtype {s} -> {d}",))),
1013 }
1014 }
1015
1016 fn reuse_or_alloc_image(
1022 cached: Option<Tensor<u8>>,
1023 w: usize,
1024 h: usize,
1025 fmt: PixelFormat,
1026 ) -> Result<Tensor<u8>> {
1027 if let Some(t) = cached {
1028 if t.width() == Some(w) && t.height() == Some(h) && t.format() == Some(fmt) {
1029 return Ok(t);
1030 }
1031 }
1032 Ok(Tensor::<u8>::image(
1033 w,
1034 h,
1035 fmt,
1036 Some(TensorMemory::Mem),
1037 edgefirst_tensor::CpuAccess::ReadWrite,
1038 )?)
1039 }
1040
1041 fn pre_resize_region(
1062 &self,
1063 src_fmt: PixelFormat,
1064 (src_w, src_h): (usize, usize),
1065 (dst_w, dst_h): (usize, usize),
1066 rotation: Rotation,
1067 crop: ResolvedCrop,
1068 ) -> Option<Rect> {
1069 use PixelFormat::{Nv12, Nv16, Nv24};
1070
1071 if !matches!(src_fmt, Nv12 | Nv16 | Nv24) {
1072 return None;
1073 }
1074 let r = crop.src_rect?;
1075 let full_src = Rect {
1076 left: 0,
1077 top: 0,
1078 width: src_w,
1079 height: src_h,
1080 };
1081 if r == full_src {
1082 return None;
1083 }
1084
1085 let d = crop.dst_rect.unwrap_or(Rect {
1090 left: 0,
1091 top: 0,
1092 width: dst_w,
1093 height: dst_h,
1094 });
1095 let (dst_x, dst_y) = match rotation {
1096 Rotation::None | Rotation::Rotate180 => (d.width, d.height),
1097 Rotation::Clockwise90 | Rotation::CounterClockwise90 => (d.height, d.width),
1098 };
1099 let halo_x = self.filter_halo(r.width, dst_x)?;
1100 let halo_y = self.filter_halo(r.height, dst_y)?;
1101
1102 let mut left = r.left.saturating_sub(halo_x);
1103 let mut top = r.top.saturating_sub(halo_y);
1104 let mut right = (r.left + r.width + halo_x).min(src_w);
1105 let mut bottom = (r.top + r.height + halo_y).min(src_h);
1106
1107 let (align_x, align_y) = match src_fmt {
1113 Nv12 => (2, 2),
1114 Nv16 => (2, 1),
1115 _ => (1, 1),
1116 };
1117 left -= left % align_x;
1118 top -= top % align_y;
1119 right = right.next_multiple_of(align_x).min(src_w);
1120 bottom = bottom.next_multiple_of(align_y).min(src_h);
1121
1122 let grown = Rect {
1123 left,
1124 top,
1125 width: right - left,
1126 height: bottom - top,
1127 };
1128 (grown != full_src).then_some(grown)
1131 }
1132
1133 #[allow(clippy::too_many_arguments)]
1135 fn convert_u8(
1136 &mut self,
1137 src: &Tensor<u8>,
1138 dst: &mut Tensor<u8>,
1139 src_fmt: PixelFormat,
1140 dst_fmt: PixelFormat,
1141 rotation: Rotation,
1142 flip: Flip,
1143 crop: ResolvedCrop,
1144 src_params: ColorParams,
1145 dst_params: ColorParams,
1146 ) -> Result<()> {
1147 use PixelFormat::*;
1148
1149 #[cfg(test)]
1150 {
1151 self.last_tmp_dims = None;
1152 }
1153
1154 let src_w = src.width().unwrap();
1155 let src_h = src.height().unwrap();
1156 let dst_w = dst.width().unwrap();
1157 let dst_h = dst.height().unwrap();
1158
1159 crop.check_crop_dims(src_w, src_h, dst_w, dst_h)?;
1160
1161 let intermediate = match (src_fmt, dst_fmt) {
1163 (Nv12, Rgb) => Rgb,
1164 (Nv12, Rgba) => Rgba,
1165 (Nv12, Grey) => Grey,
1166 (Nv12, Yuyv) => Rgba,
1167 (Nv12, Nv16) => Rgba,
1168 (Nv12, PlanarRgb) => Rgb,
1169 (Nv12, PlanarRgba) => Rgba,
1170 (Nv16, PlanarRgb) => Rgb,
1171 (Nv16, PlanarRgba) => Rgba,
1172 (Nv24, PlanarRgb) => Rgb,
1173 (Nv24, PlanarRgba) => Rgba,
1174 (Yuyv, Rgb) => Rgb,
1175 (Yuyv, Rgba) => Rgba,
1176 (Yuyv, Grey) => Grey,
1177 (Yuyv, Yuyv) => Rgba,
1178 (Yuyv, PlanarRgb) => Rgb,
1179 (Yuyv, PlanarRgba) => Rgba,
1180 (Yuyv, Nv16) => Rgba,
1181 (Vyuy, Rgb) => Rgb,
1182 (Vyuy, Rgba) => Rgba,
1183 (Vyuy, Grey) => Grey,
1184 (Vyuy, Vyuy) => Rgba,
1185 (Vyuy, PlanarRgb) => Rgb,
1186 (Vyuy, PlanarRgba) => Rgba,
1187 (Vyuy, Nv16) => Rgba,
1188 (Rgba, Rgb) => Rgba,
1189 (Rgba, Rgba) => Rgba,
1190 (Rgba, Grey) => Grey,
1191 (Rgba, Yuyv) => Rgba,
1192 (Rgba, PlanarRgb) => Rgba,
1193 (Rgba, PlanarRgba) => Rgba,
1194 (Rgba, Nv16) => Rgba,
1195 (Rgb, Rgb) => Rgb,
1196 (Rgb, Rgba) => Rgb,
1197 (Rgb, Grey) => Grey,
1198 (Rgb, Yuyv) => Rgb,
1199 (Rgb, PlanarRgb) => Rgb,
1200 (Rgb, PlanarRgba) => Rgb,
1201 (Rgb, Nv16) => Rgb,
1202 (Grey, Rgb) => Rgb,
1203 (Grey, Rgba) => Rgba,
1204 (Grey, Grey) => Grey,
1205 (Grey, Yuyv) => Grey,
1206 (Grey, PlanarRgb) => Grey,
1207 (Grey, PlanarRgba) => Grey,
1208 (Grey, Nv16) => Grey,
1209 (Nv12, Bgra) => Rgba,
1210 (Yuyv, Bgra) => Rgba,
1211 (Vyuy, Bgra) => Rgba,
1212 (Rgba, Bgra) => Rgba,
1213 (Rgb, Bgra) => Rgb,
1214 (Grey, Bgra) => Grey,
1215 (Bgra, Bgra) => Bgra,
1216 (Nv16, Rgb) => Rgb,
1217 (Nv16, Rgba) => Rgba,
1218 (Nv16, Bgra) => Rgba,
1219 (Nv24, Rgb) => Rgb,
1220 (Nv24, Rgba) => Rgba,
1221 (Nv24, Grey) => Grey,
1222 (Nv24, Bgra) => Rgba,
1223 (PlanarRgb, Rgb) => Rgb,
1224 (PlanarRgb, Rgba) => Rgb,
1225 (PlanarRgb, Bgra) => Rgb,
1226 (PlanarRgba, Rgb) => Rgba,
1227 (PlanarRgba, Rgba) => Rgba,
1228 (PlanarRgba, Bgra) => Rgba,
1229 (s, d) => {
1230 return Err(Error::NotSupported(format!("Conversion from {s} to {d}",)));
1231 }
1232 };
1233
1234 let need_resize_flip_rotation = rotation != Rotation::None
1235 || flip != Flip::None
1236 || src_w != dst_w
1237 || src_h != dst_h
1238 || crop.src_rect.is_some_and(|c| {
1239 c != Rect {
1240 left: 0,
1241 top: 0,
1242 width: src_w,
1243 height: src_h,
1244 }
1245 })
1246 || crop.dst_rect.is_some_and(|c| {
1247 c != Rect {
1248 left: 0,
1249 top: 0,
1250 width: dst_w,
1251 height: dst_h,
1252 }
1253 });
1254
1255 let direct_is_yuv_src = matches!(src_fmt, Nv12 | Nv16 | Nv24 | Yuyv | Vyuy);
1258 let direct_params = if direct_is_yuv_src {
1259 src_params
1260 } else {
1261 dst_params
1262 };
1263
1264 let full_dst_rect = Rect {
1281 left: 0,
1282 top: 0,
1283 width: dst_w,
1284 height: dst_h,
1285 };
1286 let fused_region = if rotation != Rotation::None || flip != Flip::None {
1287 None
1288 } else {
1289 match crop.src_rect {
1290 None if src_w == dst_w && src_h == dst_h => Some(None),
1291 None => None,
1292 Some(r)
1293 if r.width == dst_w
1294 && r.height == dst_h
1295 && crop.dst_rect.is_none_or(|d| d == full_dst_rect)
1296 && chroma_alignment_ok(src_fmt, r) =>
1297 {
1298 Some(Some(r))
1299 }
1300 Some(_) => None,
1301 }
1302 };
1303 if let Some(region) = fused_region {
1304 if matches!(src_fmt, Nv12 | Nv16 | Nv24) && matches!(dst_fmt, PlanarRgb | PlanarRgba) {
1305 #[cfg(test)]
1306 {
1307 self.fused_hits += 1;
1308 }
1309 return self.convert_nv_to_planar_fused(
1310 src,
1311 dst,
1312 src_fmt,
1313 dst_fmt,
1314 direct_params,
1315 region,
1316 );
1317 }
1318 }
1319
1320 if !need_resize_flip_rotation && Self::support_conversion_pf(src_fmt, dst_fmt) {
1322 return Self::convert_format_pf(src, dst, src_fmt, dst_fmt, direct_params);
1323 }
1324
1325 if dst_fmt == Yuyv && !dst_w.is_multiple_of(2) {
1327 return Err(Error::NotSupported(format!(
1328 "{} destination must have width divisible by 2",
1329 dst_fmt,
1330 )));
1331 }
1332
1333 let mut cached_tmp = self.convert_tmp.take();
1341 let mut cached_tmp2 = self.convert_tmp2.take();
1342
1343 let pre_region = if intermediate != src_fmt {
1348 self.pre_resize_region(src_fmt, (src_w, src_h), (dst_w, dst_h), rotation, crop)
1349 } else {
1350 None
1351 };
1352
1353 let tmp_holder: Option<Tensor<u8>> = if intermediate != src_fmt {
1355 let _s = tracing::trace_span!(
1356 "image.convert.cpu.format_convert",
1357 from = ?src_fmt,
1358 to = ?intermediate,
1359 pass = "pre_resize",
1360 )
1361 .entered();
1362 let (tmp_w, tmp_h) = pre_region.map_or((src_w, src_h), |g| (g.width, g.height));
1363 let mut t = Self::reuse_or_alloc_image(cached_tmp.take(), tmp_w, tmp_h, intermediate)?;
1364 #[cfg(test)]
1365 {
1366 self.last_tmp_dims = Some((tmp_w, tmp_h));
1367 }
1368 match pre_region {
1369 Some(g) => {
1370 let mut sub = Self::reuse_or_alloc_image(
1371 self.convert_src_sub.take(),
1372 g.width,
1373 g.height,
1374 src_fmt,
1375 )?;
1376 {
1377 let _s = tracing::trace_span!(
1378 "image.convert.cpu.extract_region",
1379 region_w = g.width,
1380 region_h = g.height,
1381 )
1382 .entered();
1383 Self::extract_nv_region(src, &mut sub, src_fmt, g)?;
1384 }
1385 Self::convert_format_pf(&sub, &mut t, src_fmt, intermediate, src_params)?;
1386 self.convert_src_sub = Some(sub);
1387 }
1388 None => Self::convert_format_pf(src, &mut t, src_fmt, intermediate, src_params)?,
1389 }
1390 Some(t)
1391 } else {
1392 None
1393 };
1394
1395 let crop = match (pre_region, crop.src_rect) {
1400 (Some(g), Some(r)) => ResolvedCrop {
1401 src_rect: Some(Rect {
1402 left: r.left - g.left,
1403 top: r.top - g.top,
1404 ..r
1405 }),
1406 ..crop
1407 },
1408 _ => crop,
1409 };
1410 let (tmp, tmp_fmt): (&Tensor<u8>, PixelFormat) = match &tmp_holder {
1411 Some(t) => (t, intermediate),
1412 None => (src, src_fmt),
1413 };
1414
1415 debug_assert!(matches!(tmp_fmt, Rgb | Rgba | Grey));
1417 if tmp_fmt == dst_fmt {
1418 let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
1419 self.resize_flip_rotate_pf(tmp, dst, dst_fmt, rotation, flip, crop)?;
1420 } else if !need_resize_flip_rotation {
1421 let _s = tracing::trace_span!(
1422 "image.convert.cpu.format_convert",
1423 from = ?tmp_fmt,
1424 to = ?dst_fmt,
1425 pass = "direct",
1426 )
1427 .entered();
1428 Self::convert_format_pf(tmp, dst, tmp_fmt, dst_fmt, dst_params)?;
1429 } else {
1430 let mut tmp2 = Self::reuse_or_alloc_image(cached_tmp2.take(), dst_w, dst_h, tmp_fmt)?;
1431 if crop.dst_rect.is_some_and(|c| {
1432 c != Rect {
1433 left: 0,
1434 top: 0,
1435 width: dst_w,
1436 height: dst_h,
1437 }
1438 }) && crop.dst_color.is_none()
1439 {
1440 Self::convert_format_pf(dst, &mut tmp2, dst_fmt, tmp_fmt, dst_params)?;
1441 }
1442 {
1443 let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
1444 self.resize_flip_rotate_pf(tmp, &mut tmp2, tmp_fmt, rotation, flip, crop)?;
1445 }
1446 {
1447 let _s = tracing::trace_span!(
1448 "image.convert.cpu.format_convert",
1449 from = ?tmp_fmt,
1450 to = ?dst_fmt,
1451 pass = "post_resize",
1452 )
1453 .entered();
1454 Self::convert_format_pf(&tmp2, dst, tmp_fmt, dst_fmt, dst_params)?;
1455 }
1456 cached_tmp2 = Some(tmp2);
1457 }
1458 if let Some(t) = tmp_holder {
1461 cached_tmp = Some(t);
1462 }
1463 self.convert_tmp = cached_tmp;
1464 self.convert_tmp2 = cached_tmp2;
1465
1466 if let (Some(dst_rect), Some(dst_color)) = (crop.dst_rect, crop.dst_color) {
1467 let full_rect = Rect {
1468 left: 0,
1469 top: 0,
1470 width: dst_w,
1471 height: dst_h,
1472 };
1473 if dst_rect != full_rect {
1474 Self::fill_image_outside_crop_u8(dst, dst_color, dst_rect)?;
1475 }
1476 }
1477
1478 Ok(())
1479 }
1480
1481 fn draw_decoded_masks_impl(
1482 &mut self,
1483 dst: &mut Tensor<u8>,
1484 detect: &[DetectBox],
1485 segmentation: &[Segmentation],
1486 opacity: f32,
1487 color_mode: crate::ColorMode,
1488 ) -> Result<()> {
1489 let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
1490 if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
1491 return Err(crate::Error::NotSupported(
1492 "CPU image rendering only supports RGBA or RGB images".to_string(),
1493 ));
1494 }
1495
1496 let _timer = FunctionTimer::new("CPUProcessor::draw_decoded_masks");
1497
1498 let dst_w = dst.width().unwrap();
1499 let dst_h = dst.height().unwrap();
1500 let dst_rs = tensor_row_stride(dst);
1501 let dst_c = dst_fmt.channels();
1502
1503 let mut map = dst.map_mut()?;
1504 let dst_slice = map.as_mut_slice();
1505
1506 self.render_box(dst_w, dst_h, dst_rs, dst_c, dst_slice, detect, color_mode)?;
1507
1508 if segmentation.is_empty() {
1509 return Ok(());
1510 }
1511
1512 let is_semantic = segmentation[0].segmentation.shape()[2] > 1;
1515
1516 if is_semantic {
1517 self.render_modelpack_segmentation(
1518 dst_w,
1519 dst_h,
1520 dst_rs,
1521 dst_c,
1522 dst_slice,
1523 &segmentation[0],
1524 opacity,
1525 )?;
1526 } else {
1527 for (idx, (seg, det)) in segmentation.iter().zip(detect).enumerate() {
1528 let color_index = color_mode.index(idx, det.label);
1529 self.render_yolo_segmentation(
1530 dst_w,
1531 dst_h,
1532 dst_rs,
1533 dst_c,
1534 dst_slice,
1535 seg,
1536 color_index,
1537 opacity,
1538 )?;
1539 }
1540 }
1541
1542 Ok(())
1543 }
1544
1545 fn draw_proto_masks_impl(
1546 &mut self,
1547 dst: &mut Tensor<u8>,
1548 detect: &[DetectBox],
1549 proto_data: &ProtoData,
1550 opacity: f32,
1551 letterbox: Option<[f32; 4]>,
1552 color_mode: crate::ColorMode,
1553 ) -> Result<()> {
1554 let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
1555 if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
1556 return Err(crate::Error::NotSupported(
1557 "CPU image rendering only supports RGBA or RGB images".to_string(),
1558 ));
1559 }
1560
1561 let _timer = FunctionTimer::new("CPUProcessor::draw_proto_masks");
1562
1563 let dst_w = dst.width().unwrap();
1564 let dst_h = dst.height().unwrap();
1565 let dst_rs = tensor_row_stride(dst);
1566 let channels = dst_fmt.channels();
1567
1568 let mut map = dst.map_mut()?;
1569 let dst_slice = map.as_mut_slice();
1570
1571 self.render_box(
1572 dst_w, dst_h, dst_rs, channels, dst_slice, detect, color_mode,
1573 )?;
1574
1575 if detect.is_empty() {
1576 return Ok(());
1577 }
1578 let proto_shape = proto_data.protos.shape();
1579 if proto_shape.len() != 3 {
1580 return Err(Error::InvalidShape(format!(
1581 "protos tensor must be rank-3, got {proto_shape:?}"
1582 )));
1583 }
1584 let proto_h = proto_shape[0];
1585 let proto_w = proto_shape[1];
1586 let num_protos = proto_shape[2];
1587 let coeff_shape = proto_data.mask_coefficients.shape();
1588 if coeff_shape.len() != 2 {
1589 return Err(Error::InvalidShape(format!(
1590 "mask_coefficients tensor must be rank-2, got {coeff_shape:?}"
1591 )));
1592 }
1593 if coeff_shape[0] == 0 {
1595 return Ok(());
1596 }
1597 if coeff_shape[1] != num_protos {
1598 return Err(Error::InvalidShape(format!(
1599 "mask_coefficients second dimension must match num_protos \
1600 ({num_protos}), got {coeff_shape:?}"
1601 )));
1602 }
1603
1604 let coeff_f32: Vec<f32> = match proto_data.mask_coefficients.dtype() {
1606 DType::F32 => {
1607 let t = proto_data.mask_coefficients.as_f32().expect("F32");
1608 let m = t.map_read()?;
1609 m.as_slice().to_vec()
1610 }
1611 DType::F16 => {
1612 let t = proto_data.mask_coefficients.as_f16().expect("F16");
1613 let m = t.map_read()?;
1614 m.as_slice().iter().map(|v| v.to_f32()).collect()
1615 }
1616 DType::I8 => {
1617 let t = proto_data.mask_coefficients.as_i8().expect("I8");
1618 let m = t.map_read()?;
1619 if let Some(q) = t.quantization() {
1620 use edgefirst_tensor::QuantMode;
1621 let (scale, zp) = match q.mode() {
1622 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1623 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1624 other => {
1625 return Err(Error::NotSupported(format!(
1626 "I8 mask_coefficients quantization mode {other:?} not supported"
1627 )));
1628 }
1629 };
1630 m.as_slice()
1631 .iter()
1632 .map(|&v| (v as f32 - zp) * scale)
1633 .collect()
1634 } else {
1635 m.as_slice().iter().map(|&v| v as f32).collect()
1636 }
1637 }
1638 DType::I16 => {
1639 let t = proto_data.mask_coefficients.as_i16().expect("I16");
1640 let m = t.map_read()?;
1641 if let Some(q) = t.quantization() {
1642 use edgefirst_tensor::QuantMode;
1643 let (scale, zp) = match q.mode() {
1644 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1645 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1646 other => {
1647 return Err(Error::NotSupported(format!(
1648 "I16 mask_coefficients quantization mode {other:?} not supported"
1649 )));
1650 }
1651 };
1652 m.as_slice()
1653 .iter()
1654 .map(|&v| (v as f32 - zp) * scale)
1655 .collect()
1656 } else {
1657 m.as_slice().iter().map(|&v| v as f32).collect()
1658 }
1659 }
1660 other => {
1661 return Err(Error::InvalidShape(format!(
1662 "mask_coefficients dtype {other:?} not supported"
1663 )));
1664 }
1665 };
1666
1667 let (lx0, lx_range, ly0, ly_range) = match letterbox {
1669 Some([lx0, ly0, lx1, ly1]) => (lx0, lx1 - lx0, ly0, ly1 - ly0),
1670 None => (0.0_f32, 1.0_f32, 0.0_f32, 1.0_f32),
1671 };
1672
1673 match proto_data.protos.dtype() {
1676 DType::F32 => {
1677 let t = proto_data.protos.as_f32().expect("F32");
1678 let m = t.map_read()?;
1679 self.draw_proto_masks_inner(
1680 dst_slice,
1681 dst_w,
1682 dst_h,
1683 dst_rs,
1684 channels,
1685 detect,
1686 m.as_slice(),
1687 &coeff_f32,
1688 proto_h,
1689 proto_w,
1690 num_protos,
1691 opacity,
1692 (lx0, lx_range, ly0, ly_range),
1693 color_mode,
1694 0.0_f32,
1695 |p: &f32, _| *p,
1696 );
1697 }
1698 DType::F16 => {
1699 let t = proto_data.protos.as_f16().expect("F16");
1700 let m = t.map_read()?;
1701 self.draw_proto_masks_inner(
1702 dst_slice,
1703 dst_w,
1704 dst_h,
1705 dst_rs,
1706 channels,
1707 detect,
1708 m.as_slice(),
1709 &coeff_f32,
1710 proto_h,
1711 proto_w,
1712 num_protos,
1713 opacity,
1714 (lx0, lx_range, ly0, ly_range),
1715 color_mode,
1716 0.0_f32,
1717 |p: &half::f16, _| p.to_f32(),
1718 );
1719 }
1720 DType::I8 => {
1721 use edgefirst_tensor::QuantMode;
1722 let t = proto_data.protos.as_i8().expect("I8");
1723 let m = t.map_read()?;
1724 let quant = t.quantization().ok_or_else(|| {
1725 Error::InvalidShape("I8 protos require quantization metadata".into())
1726 })?;
1727 let (scale, zp) = match quant.mode() {
1728 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1729 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1730 QuantMode::PerChannel { axis, .. }
1731 | QuantMode::PerChannelSymmetric { axis, .. } => {
1732 return Err(Error::NotSupported(format!(
1733 "per-channel quantization (axis={axis}) in draw_proto_masks \
1734 CPU path not yet supported"
1735 )));
1736 }
1737 };
1738 self.draw_proto_masks_inner(
1739 dst_slice,
1740 dst_w,
1741 dst_h,
1742 dst_rs,
1743 channels,
1744 detect,
1745 m.as_slice(),
1746 &coeff_f32,
1747 proto_h,
1748 proto_w,
1749 num_protos,
1750 opacity,
1751 (lx0, lx_range, ly0, ly_range),
1752 color_mode,
1753 scale,
1754 move |p: &i8, _| (*p as f32) - zp,
1755 );
1756 }
1757 other => {
1758 return Err(Error::InvalidShape(format!(
1759 "proto tensor dtype {other:?} not supported"
1760 )));
1761 }
1762 }
1763
1764 Ok(())
1765 }
1766
1767 #[allow(clippy::too_many_arguments)]
1768 fn draw_proto_masks_inner<P: Copy>(
1769 &self,
1770 dst_slice: &mut [u8],
1771 dst_w: usize,
1772 dst_h: usize,
1773 dst_rs: usize,
1774 channels: usize,
1775 detect: &[DetectBox],
1776 protos: &[P],
1777 coeff_all_f32: &[f32],
1778 proto_h: usize,
1779 proto_w: usize,
1780 num_protos: usize,
1781 opacity: f32,
1782 letterbox_xy: (f32, f32, f32, f32),
1783 color_mode: crate::ColorMode,
1784 acc_scale: f32,
1785 load_f32: impl Fn(&P, f32) -> f32 + Copy,
1786 ) {
1787 let (lx0, lx_range, ly0, ly_range) = letterbox_xy;
1788 let stride_y = proto_w * num_protos;
1789 for (idx, det) in detect.iter().enumerate() {
1790 let coeff = &coeff_all_f32[idx * num_protos..(idx + 1) * num_protos];
1791 let color_index = color_mode.index(idx, det.label);
1792 let color = self.colors[color_index % self.colors.len()];
1793 let alpha = if opacity == 1.0 {
1794 color[3] as u16
1795 } else {
1796 (color[3] as f32 * opacity).round() as u16
1797 };
1798
1799 let start_x = (dst_w as f32 * det.bbox.xmin).round() as usize;
1800 let start_y = (dst_h as f32 * det.bbox.ymin).round() as usize;
1801 let end_x = ((dst_w as f32 * det.bbox.xmax).round() as usize).min(dst_w);
1802 let end_y = ((dst_h as f32 * det.bbox.ymax).round() as usize).min(dst_h);
1803
1804 for y in start_y..end_y {
1805 for x in start_x..end_x {
1806 let px = (lx0 + (x as f32 / dst_w as f32) * lx_range) * proto_w as f32 - 0.5;
1807 let py = (ly0 + (y as f32 / dst_h as f32) * ly_range) * proto_h as f32 - 0.5;
1808
1809 let x0 = (px.floor() as isize).clamp(0, proto_w as isize - 1) as usize;
1813 let y0 = (py.floor() as isize).clamp(0, proto_h as isize - 1) as usize;
1814 let x1 = (x0 + 1).min(proto_w - 1);
1815 let y1 = (y0 + 1).min(proto_h - 1);
1816 let fx = px - px.floor();
1817 let fy = py - py.floor();
1818 let w00 = (1.0 - fx) * (1.0 - fy);
1819 let w10 = fx * (1.0 - fy);
1820 let w01 = (1.0 - fx) * fy;
1821 let w11 = fx * fy;
1822 let b00 = y0 * stride_y + x0 * num_protos;
1823 let b10 = y0 * stride_y + x1 * num_protos;
1824 let b01 = y1 * stride_y + x0 * num_protos;
1825 let b11 = y1 * stride_y + x1 * num_protos;
1826 let mut acc = 0.0_f32;
1827 for p in 0..num_protos {
1828 let v00 = load_f32(&protos[b00 + p], 0.0);
1829 let v10 = load_f32(&protos[b10 + p], 0.0);
1830 let v01 = load_f32(&protos[b01 + p], 0.0);
1831 let v11 = load_f32(&protos[b11 + p], 0.0);
1832 let val = w00 * v00 + w10 * v10 + w01 * v01 + w11 * v11;
1833 acc += coeff[p] * val;
1834 }
1835 let final_acc = if acc_scale == 0.0 {
1836 acc
1837 } else {
1838 acc_scale * acc
1839 };
1840 let mask = 1.0 / (1.0 + (-final_acc).exp());
1844 if mask < 0.5 {
1845 continue;
1846 }
1847 let dst_index = y * dst_rs + x * channels;
1848 for c in 0..3 {
1849 dst_slice[dst_index + c] = ((color[c] as u16 * alpha
1850 + dst_slice[dst_index + c] as u16 * (255 - alpha))
1851 / 255) as u8;
1852 }
1853 }
1854 }
1855 }
1856 }
1857}