1use crate::{
5 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 normalize;
15mod resize;
16mod simd;
17mod tests;
18
19#[derive(Debug, Clone, Copy)]
31pub(crate) struct ColorParams {
32 pub matrix: yuv::YuvStandardMatrix,
34 pub range: yuv::YuvRange,
36 pub encoding: edgefirst_tensor::ColorEncoding,
41 pub range_kind: edgefirst_tensor::ColorRange,
42 pub src_full_range: bool,
45 pub dst_full_range: bool,
48}
49
50#[derive(Debug)]
53pub struct CPUProcessor {
54 resizer: fast_image_resize::Resizer,
55 options: fast_image_resize::ResizeOptions,
56 colors: [[u8; 4]; 20],
57 box_line_width: usize,
59 widen_scratch: Option<TensorDyn>,
65 resize_destride_scratch: Vec<u8>,
73 resize_dst_destride_scratch: Vec<u8>,
84 nv_strip_scratch: Vec<u8>,
90 nv_strip_y_pack: Vec<u8>,
106 nv_strip_uv_pack: Vec<u8>,
108 #[cfg(test)]
115 fused_hits: u64,
116 #[cfg(test)]
123 last_tmp_dims: Option<(usize, usize)>,
124 convert_tmp: Option<Tensor<u8>>,
131 convert_tmp2: Option<Tensor<u8>>,
132 convert_src_sub: Option<Tensor<u8>>,
137}
138
139impl Clone for CPUProcessor {
145 fn clone(&self) -> Self {
146 Self {
147 resizer: self.resizer.clone(),
148 options: self.options,
149 colors: self.colors,
150 box_line_width: self.box_line_width,
151 widen_scratch: None,
152 resize_destride_scratch: Vec::new(),
153 resize_dst_destride_scratch: Vec::new(),
154 nv_strip_scratch: Vec::new(),
155 nv_strip_y_pack: Vec::new(),
156 nv_strip_uv_pack: Vec::new(),
157 #[cfg(test)]
158 fused_hits: 0,
159 #[cfg(test)]
160 last_tmp_dims: None,
161 convert_tmp: None,
162 convert_tmp2: None,
163 convert_src_sub: None,
164 }
165 }
166}
167
168unsafe impl Send for CPUProcessor {}
169unsafe impl Sync for CPUProcessor {}
170
171impl Default for CPUProcessor {
172 fn default() -> Self {
173 Self::new_bilinear()
174 }
175}
176
177fn prepare_dst_base_cpu(dst: &mut TensorDyn, background: Option<&TensorDyn>) -> Result<()> {
188 match background {
189 Some(bg) => {
190 if bg.shape() != dst.shape() {
191 return Err(Error::InvalidShape(
192 "background shape does not match dst".into(),
193 ));
194 }
195 if bg.format() != dst.format() {
196 return Err(Error::InvalidShape(
197 "background pixel format does not match dst".into(),
198 ));
199 }
200 let bg_u8 = bg.as_typed::<u8>().ok_or(Error::NotAnImage)?;
201 let dst_u8 = dst.as_typed_mut::<u8>().ok_or(Error::NotAnImage)?;
202 CPUProcessor::copy_image(bg_u8, dst_u8)?;
209 }
210 None => {
211 let dst_u8 = dst.as_typed_mut::<u8>().ok_or(Error::NotAnImage)?;
212 let mut dst_map = dst_u8.map_mut()?;
213 dst_map.as_mut_slice().fill(0);
214 }
215 }
216 Ok(())
217}
218
219fn chroma_alignment_ok(fmt: PixelFormat, r: Rect) -> bool {
227 match fmt {
228 PixelFormat::Nv12 => r.left.is_multiple_of(2) && r.top.is_multiple_of(2),
229 PixelFormat::Nv16 => r.left.is_multiple_of(2),
230 PixelFormat::Nv24 => true,
231 _ => false,
232 }
233}
234
235fn row_stride_for(width: usize, fmt: PixelFormat) -> usize {
237 use edgefirst_tensor::PixelLayout;
238 match fmt.layout() {
239 PixelLayout::Packed => width * fmt.channels(),
240 PixelLayout::Planar | PixelLayout::SemiPlanar => width,
241 _ => width, }
243}
244
245fn tensor_row_stride(tensor: &Tensor<u8>) -> usize {
250 tensor.effective_row_stride().unwrap_or_else(|| {
251 let w = tensor.width().unwrap_or(0);
252 let fmt = tensor.format().unwrap_or(PixelFormat::Rgb);
253 row_stride_for(w, fmt)
254 })
255}
256
257fn logical_surface(tensor: &Tensor<u8>) -> Result<(usize, usize)> {
267 use edgefirst_tensor::PixelLayout;
268 let fmt = tensor.format().ok_or(Error::NotAnImage)?;
269 let w = tensor.width().ok_or(Error::NotAnImage)?;
270 let h = tensor.height().ok_or(Error::NotAnImage)?;
271 Ok(match fmt.layout() {
272 PixelLayout::Packed => (h, w * fmt.channels()),
273 PixelLayout::Planar => (fmt.channels() * h, w),
274 PixelLayout::SemiPlanar => (fmt.combined_plane_height(h).unwrap_or(h), w),
275 _ => (h, row_stride_for(w, fmt)),
278 })
279}
280
281fn packed_row_pairs<'s, 'd>(
287 src: &'s [u8],
288 src_stride: usize,
289 src_row_bytes: usize,
290 dst: &'d mut [u8],
291 dst_stride: usize,
292 dst_row_bytes: usize,
293 rows: usize,
294) -> impl Iterator<Item = (&'s [u8], &'d mut [u8])> {
295 src.chunks(src_stride)
296 .zip(dst.chunks_mut(dst_stride))
297 .take(rows)
298 .map(move |(s, d)| (&s[..src_row_bytes], &mut d[..dst_row_bytes]))
299}
300
301fn split_semi_planar(
307 bytes: &[u8],
308 stride: usize,
309 src_h: usize,
310 fmt: PixelFormat,
311) -> Result<(&[u8], &[u8])> {
312 let total_h = fmt.combined_plane_height(src_h).unwrap_or(src_h);
313 let need = stride.checked_mul(total_h).ok_or_else(|| {
314 Error::InvalidShape(format!(
315 "{fmt:?} plane size overflow (stride={stride}, h={src_h})"
316 ))
317 })?;
318 if bytes.len() < need {
319 return Err(Error::InvalidShape(format!(
320 "{fmt:?} source has {} bytes but needs {need} (stride={stride}, h={src_h})",
321 bytes.len()
322 )));
323 }
324 Ok(bytes.split_at(stride * src_h))
325}
326
327fn split_semi_planar_mut(
334 bytes: &mut [u8],
335 stride: usize,
336 dst_h: usize,
337 fmt: PixelFormat,
338) -> Result<(&mut [u8], &mut [u8])> {
339 let total_h = fmt.combined_plane_height(dst_h).unwrap_or(dst_h);
340 let need = stride.checked_mul(total_h).ok_or_else(|| {
341 Error::InvalidShape(format!(
342 "{fmt:?} plane size overflow (stride={stride}, combined_h={total_h})"
343 ))
344 })?;
345 if bytes.len() < need {
346 return Err(Error::InvalidShape(format!(
347 "{fmt:?} destination has {} bytes but needs {need} (stride={stride}, combined_h={total_h})",
348 bytes.len()
349 )));
350 }
351 Ok(bytes.split_at_mut(stride * dst_h))
352}
353
354fn guard_plane(
360 buf_len: usize,
361 stride: usize,
362 rows: usize,
363 row_bytes: usize,
364 what: &str,
365) -> Result<()> {
366 let need = stride.checked_mul(rows).ok_or_else(|| {
367 Error::InvalidShape(format!(
368 "{what} plane size overflow (stride={stride}, rows={rows})"
369 ))
370 })?;
371 if row_bytes > stride || buf_len < need {
372 return Err(Error::InvalidShape(format!(
373 "{what} buffer too small: {buf_len} bytes, need {need} (stride={stride}, rows={rows}, row_bytes={row_bytes})"
374 )));
375 }
376 Ok(())
377}
378
379pub(crate) fn apply_int8_xor_bias(data: &mut [u8], fmt: PixelFormat) {
384 use edgefirst_tensor::PixelLayout;
385 if !fmt.has_alpha() {
386 for b in data.iter_mut() {
387 *b ^= 0x80;
388 }
389 } else if fmt.layout() == PixelLayout::Planar {
390 let channels = fmt.channels();
392 let plane_size = data.len() / channels;
393 for b in data[..plane_size * (channels - 1)].iter_mut() {
394 *b ^= 0x80;
395 }
396 } else {
397 let channels = fmt.channels();
399 for pixel in data.chunks_exact_mut(channels) {
400 for b in &mut pixel[..channels - 1] {
401 *b ^= 0x80;
402 }
403 }
404 }
405}
406
407pub(crate) fn apply_int8_xor_bias_rows(tensor: &mut Tensor<u8>, fmt: PixelFormat) -> Result<()> {
413 use edgefirst_tensor::PixelLayout;
414 let (rows, row_bytes) = logical_surface(tensor)?;
415 let stride = tensor_row_stride(tensor);
416 let color_rows = if fmt.has_alpha() && fmt.layout() == PixelLayout::Planar {
419 rows / fmt.channels() * (fmt.channels() - 1)
420 } else {
421 rows
422 };
423 let row_fmt = if fmt.layout() == PixelLayout::Planar {
426 PixelFormat::Grey
427 } else {
428 fmt
429 };
430 let mut map = tensor.map_mut()?;
432 let buf = map.as_mut_slice();
433 guard_plane(buf.len(), stride, rows, row_bytes, "int8 bias dst")?;
434 for row in buf.chunks_mut(stride).take(color_rows) {
435 apply_int8_xor_bias(&mut row[..row_bytes], row_fmt);
436 }
437 Ok(())
438}
439
440impl CPUProcessor {
441 pub fn new() -> Self {
443 Self::new_bilinear()
444 }
445
446 fn resize_alg(mode: crate::Interpolation) -> fast_image_resize::ResizeAlg {
449 match mode.supported() {
450 crate::Interpolation::Nearest => fast_image_resize::ResizeAlg::Nearest,
451 _ => {
452 fast_image_resize::ResizeAlg::Interpolation(fast_image_resize::FilterType::Bilinear)
453 }
454 }
455 }
456
457 fn new_bilinear() -> Self {
459 Self::with_interpolation(crate::Interpolation::Bilinear)
460 }
461
462 fn with_interpolation(mode: crate::Interpolation) -> Self {
464 let resizer = fast_image_resize::Resizer::new();
465 let options = fast_image_resize::ResizeOptions::new()
466 .resize_alg(Self::resize_alg(mode))
467 .use_alpha(false);
468
469 log::debug!("CPUConverter created");
470 Self {
471 resizer,
472 options,
473 colors: crate::DEFAULT_COLORS_U8,
474 box_line_width: crate::DEFAULT_BOX_LINE_WIDTH as usize,
475 widen_scratch: None,
476 resize_destride_scratch: Vec::new(),
477 resize_dst_destride_scratch: Vec::new(),
478 nv_strip_scratch: Vec::new(),
479 nv_strip_y_pack: Vec::new(),
480 nv_strip_uv_pack: Vec::new(),
481 #[cfg(test)]
482 fused_hits: 0,
483 #[cfg(test)]
484 last_tmp_dims: None,
485 convert_tmp: None,
486 convert_tmp2: None,
487 convert_src_sub: None,
488 }
489 }
490
491 pub fn set_interpolation(&mut self, mode: crate::Interpolation) {
494 self.options = self.options.clone().resize_alg(Self::resize_alg(mode));
495 }
496
497 pub fn new_nearest() -> Self {
499 Self::with_interpolation(crate::Interpolation::Nearest)
500 }
501
502 #[cfg(test)]
504 pub(super) fn fused_hits(&self) -> u64 {
505 self.fused_hits
506 }
507
508 #[cfg(test)]
511 pub(super) fn last_tmp_dims(&self) -> Option<(usize, usize)> {
512 self.last_tmp_dims
513 }
514
515 pub(crate) fn support_conversion_pf(src: PixelFormat, dst: PixelFormat) -> bool {
516 use PixelFormat::*;
517 matches!(
518 (src, dst),
519 (Nv12, Rgb)
520 | (Nv12, Rgba)
521 | (Nv12, Grey)
522 | (Nv16, Rgb)
523 | (Nv16, Rgba)
524 | (Nv16, Bgra)
525 | (Nv24, Rgb)
526 | (Nv24, Rgba)
527 | (Nv24, Grey)
528 | (Nv24, Bgra)
529 | (Yuyv, Rgb)
530 | (Yuyv, Rgba)
531 | (Yuyv, Grey)
532 | (Yuyv, Yuyv)
533 | (Yuyv, PlanarRgb)
534 | (Yuyv, PlanarRgba)
535 | (Yuyv, Nv16)
536 | (Vyuy, Rgb)
537 | (Vyuy, Rgba)
538 | (Vyuy, Grey)
539 | (Vyuy, Vyuy)
540 | (Vyuy, PlanarRgb)
541 | (Vyuy, PlanarRgba)
542 | (Vyuy, Nv16)
543 | (Rgba, Rgb)
544 | (Rgba, Rgba)
545 | (Rgba, Grey)
546 | (Rgba, Yuyv)
547 | (Rgba, PlanarRgb)
548 | (Rgba, PlanarRgba)
549 | (Rgba, Nv16)
550 | (Rgb, Rgb)
551 | (Rgb, Rgba)
552 | (Rgb, Grey)
553 | (Rgb, Yuyv)
554 | (Rgb, PlanarRgb)
555 | (Rgb, PlanarRgba)
556 | (Rgb, Nv16)
557 | (Grey, Rgb)
558 | (Grey, Rgba)
559 | (Grey, Grey)
560 | (Grey, Yuyv)
561 | (Grey, PlanarRgb)
562 | (Grey, PlanarRgba)
563 | (Grey, Nv16)
564 | (Nv12, Bgra)
565 | (Yuyv, Bgra)
566 | (Vyuy, Bgra)
567 | (Rgba, Bgra)
568 | (Rgb, Bgra)
569 | (Grey, Bgra)
570 | (Bgra, Bgra)
571 | (PlanarRgb, Rgb)
572 | (PlanarRgb, Rgba)
573 | (PlanarRgba, Rgb)
574 | (PlanarRgba, Rgba)
575 | (PlanarRgb, Bgra)
576 | (PlanarRgba, Bgra)
577 )
578 }
579
580 pub(crate) fn convert_format_pf(
582 src: &Tensor<u8>,
583 dst: &mut Tensor<u8>,
584 src_fmt: PixelFormat,
585 dst_fmt: PixelFormat,
586 cp: ColorParams,
587 ) -> Result<()> {
588 let _timer = FunctionTimer::new(format!(
589 "ImageProcessor::convert_format {} to {}",
590 src_fmt, dst_fmt,
591 ));
592
593 use PixelFormat::*;
594 match (src_fmt, dst_fmt) {
595 (Nv12, Rgb) => Self::convert_nv12_to_rgb(src, dst, cp),
596 (Nv12, Rgba) => Self::convert_nv12_to_rgba(src, dst, cp),
597 (Nv12, Grey) => Self::convert_nv12_to_grey(src, dst, cp),
598 (Yuyv, Rgb) => Self::convert_yuyv_to_rgb(src, dst, cp),
599 (Yuyv, Rgba) => Self::convert_yuyv_to_rgba(src, dst, cp),
600 (Yuyv, Grey) => Self::convert_yuyv_to_grey(src, dst, cp),
601 (Yuyv, Yuyv) => Self::copy_image(src, dst),
602 (Yuyv, PlanarRgb) => Self::convert_yuyv_to_8bps(src, dst, cp),
603 (Yuyv, PlanarRgba) => Self::convert_yuyv_to_prgba(src, dst, cp),
604 (Yuyv, Nv16) => Self::convert_yuyv_to_nv16(src, dst),
605 (Vyuy, Rgb) => Self::convert_vyuy_to_rgb(src, dst, cp),
606 (Vyuy, Rgba) => Self::convert_vyuy_to_rgba(src, dst, cp),
607 (Vyuy, Grey) => Self::convert_vyuy_to_grey(src, dst, cp),
608 (Vyuy, Vyuy) => Self::copy_image(src, dst),
609 (Vyuy, PlanarRgb) => Self::convert_vyuy_to_8bps(src, dst, cp),
610 (Vyuy, PlanarRgba) => Self::convert_vyuy_to_prgba(src, dst, cp),
611 (Vyuy, Nv16) => Self::convert_vyuy_to_nv16(src, dst),
612 (Rgba, Rgb) => Self::convert_rgba_to_rgb(src, dst),
613 (Rgba, Rgba) => Self::copy_image(src, dst),
614 (Rgba, Grey) => Self::convert_rgba_to_grey(src, dst),
615 (Rgba, Yuyv) => Self::convert_rgba_to_yuyv(src, dst, cp),
616 (Rgba, PlanarRgb) => Self::convert_rgba_to_8bps(src, dst),
617 (Rgba, PlanarRgba) => Self::convert_rgba_to_prgba(src, dst),
618 (Rgba, Nv16) => Self::convert_rgba_to_nv16(src, dst, cp),
619 (Rgb, Rgb) => Self::copy_image(src, dst),
620 (Rgb, Rgba) => Self::convert_rgb_to_rgba(src, dst),
621 (Rgb, Grey) => Self::convert_rgb_to_grey(src, dst),
622 (Rgb, Yuyv) => Self::convert_rgb_to_yuyv(src, dst, cp),
623 (Rgb, PlanarRgb) => Self::convert_rgb_to_8bps(src, dst),
624 (Rgb, PlanarRgba) => Self::convert_rgb_to_prgba(src, dst),
625 (Rgb, Nv16) => Self::convert_rgb_to_nv16(src, dst, cp),
626 (Grey, Rgb) => Self::convert_grey_to_rgb(src, dst),
627 (Grey, Rgba) => Self::convert_grey_to_rgba(src, dst),
628 (Grey, Grey) => Self::copy_image(src, dst),
629 (Grey, Yuyv) => Self::convert_grey_to_yuyv(src, dst, cp),
630 (Grey, PlanarRgb) => Self::convert_grey_to_8bps(src, dst),
631 (Grey, PlanarRgba) => Self::convert_grey_to_prgba(src, dst),
632 (Grey, Nv16) => Self::convert_grey_to_nv16(src, dst, cp),
633
634 (Nv16, Rgb) => Self::convert_nv16_to_rgb(src, dst, cp),
636 (Nv16, Rgba) => Self::convert_nv16_to_rgba(src, dst, cp),
637 (Nv24, Rgb) => Self::convert_nv24_to_rgb(src, dst, cp),
638 (Nv24, Rgba) => Self::convert_nv24_to_rgba(src, dst, cp),
639 (Nv24, Grey) => Self::convert_nv24_to_grey(src, dst, cp),
640 (PlanarRgb, Rgb) => Self::convert_8bps_to_rgb(src, dst),
641 (PlanarRgb, Rgba) => Self::convert_8bps_to_rgba(src, dst),
642 (PlanarRgba, Rgb) => Self::convert_prgba_to_rgb(src, dst),
643 (PlanarRgba, Rgba) => Self::convert_prgba_to_rgba(src, dst),
644
645 (Bgra, Bgra) => Self::copy_image(src, dst),
647 (Nv12, Bgra) => {
648 Self::convert_nv12_to_rgba(src, dst, cp)?;
649 Self::swizzle_rb_4chan(dst)
650 }
651 (Nv16, Bgra) => {
652 Self::convert_nv16_to_rgba(src, dst, cp)?;
653 Self::swizzle_rb_4chan(dst)
654 }
655 (Nv24, Bgra) => {
656 Self::convert_nv24_to_rgba(src, dst, cp)?;
657 Self::swizzle_rb_4chan(dst)
658 }
659 (Yuyv, Bgra) => {
660 Self::convert_yuyv_to_rgba(src, dst, cp)?;
661 Self::swizzle_rb_4chan(dst)
662 }
663 (Vyuy, Bgra) => {
664 Self::convert_vyuy_to_rgba(src, dst, cp)?;
665 Self::swizzle_rb_4chan(dst)
666 }
667 (Rgba, Bgra) => {
668 Self::copy_image(src, dst)?;
669 Self::swizzle_rb_4chan(dst)
670 }
671 (Rgb, Bgra) => {
672 Self::convert_rgb_to_rgba(src, dst)?;
673 Self::swizzle_rb_4chan(dst)
674 }
675 (Grey, Bgra) => {
676 Self::convert_grey_to_rgba(src, dst)?;
677 Self::swizzle_rb_4chan(dst)
678 }
679 (PlanarRgb, Bgra) => {
680 Self::convert_8bps_to_rgba(src, dst)?;
681 Self::swizzle_rb_4chan(dst)
682 }
683 (PlanarRgba, Bgra) => {
684 Self::convert_prgba_to_rgba(src, dst)?;
685 Self::swizzle_rb_4chan(dst)
686 }
687
688 (s, d) => Err(Error::NotSupported(format!("Conversion from {s} to {d}",))),
689 }
690 }
691
692 pub(crate) fn fill_image_outside_crop_u8(
695 dst: &mut Tensor<u8>,
696 rgba: [u8; 4],
697 crop: Rect,
698 device_interior: bool,
699 ) -> Result<()> {
700 let dst_fmt = dst.format().unwrap();
701 let dst_w = dst.width().unwrap();
702 let dst_h = dst.height().unwrap();
703 let cm = crate::colorimetry::resolve_colorimetry(dst.colorimetry(), dst.height());
707 let cp = ColorParams {
708 matrix: crate::colorimetry::yuv_matrix(cm.encoding.unwrap()),
709 range: crate::colorimetry::yuv_range(cm.range.unwrap()),
710 encoding: cm.encoding.unwrap(),
711 range_kind: cm.range.unwrap(),
712 src_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
713 dst_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
714 };
715 let dst_stride = tensor_row_stride(dst);
716 let mut dst_map = if device_interior {
717 dst.map_for_partial_write()?
718 } else {
719 dst.map_for_write()?
720 };
721 let dst_tup = (dst_map.as_mut_slice(), dst_w, dst_h, dst_stride);
722 Self::fill_outside_crop_dispatch(dst_tup, dst_fmt, rgba, crop, cp)
723 }
724
725 fn fill_outside_crop_dispatch(
730 dst: (&mut [u8], usize, usize, usize),
731 fmt: PixelFormat,
732 rgba: [u8; 4],
733 crop: Rect,
734 cp: ColorParams,
735 ) -> Result<()> {
736 use PixelFormat::*;
737 match fmt {
738 Rgba | Bgra => Self::fill_image_outside_crop_(dst, rgba, crop),
739 Rgb => Self::fill_image_outside_crop_(dst, Self::rgba_to_rgb(rgba), crop),
740 Grey => Self::fill_image_outside_crop_(dst, Self::rgba_to_grey(rgba), crop),
741 Yuyv => {
742 let (bytes, w, h, stride) = dst;
743 let yuyv = Self::rgba_to_yuyv(rgba, cp);
744 Self::fill_image_outside_crop_(
747 (&mut *bytes, w / 2, h, stride),
748 yuyv,
749 Rect::new(crop.left / 2, crop.top, crop.width.div_ceil(2), crop.height),
750 )?;
751 if w % 2 == 1 {
759 let last = w - 1;
760 let col_outside = last < crop.left || last >= crop.left + crop.width;
761 for y in 0..h {
762 let row_outside = y < crop.top || y >= crop.top + crop.height;
763 if row_outside || col_outside {
764 let off = y * stride + last * 2;
765 bytes[off] = yuyv[0];
766 bytes[off + 1] = yuyv[1];
767 }
768 }
769 }
770 Ok(())
771 }
772 PlanarRgb => Self::fill_image_outside_crop_planar(dst, Self::rgba_to_rgb(rgba), crop),
773 PlanarRgba => Self::fill_image_outside_crop_planar(dst, rgba, crop),
774 Nv16 => {
775 let yuyv = Self::rgba_to_yuyv(rgba, cp);
776 Self::fill_image_outside_crop_yuv_semiplanar(dst, yuyv[0], [yuyv[1], yuyv[3]], crop)
777 }
778 _ => Err(Error::Internal(format!(
779 "Found unexpected destination {fmt}",
780 ))),
781 }
782 }
783}
784
785fn color_params(src: &TensorDyn, dst: &TensorDyn) -> (ColorParams, ColorParams) {
788 let src_cm = crate::colorimetry::effective_colorimetry(src);
793 let dst_cm = crate::colorimetry::effective_colorimetry(dst);
794 let src_full = src_cm.range == Some(edgefirst_tensor::ColorRange::Full);
795 let dst_full = dst_cm.range == Some(edgefirst_tensor::ColorRange::Full);
796 let src_params = ColorParams {
797 matrix: crate::colorimetry::yuv_matrix(src_cm.encoding.unwrap()),
798 range: crate::colorimetry::yuv_range(src_cm.range.unwrap()),
799 encoding: src_cm.encoding.unwrap(),
800 range_kind: src_cm.range.unwrap(),
801 src_full_range: src_full,
802 dst_full_range: dst_full,
803 };
804 let dst_params = ColorParams {
805 matrix: crate::colorimetry::yuv_matrix(dst_cm.encoding.unwrap()),
806 range: crate::colorimetry::yuv_range(dst_cm.range.unwrap()),
807 encoding: dst_cm.encoding.unwrap(),
808 range_kind: dst_cm.range.unwrap(),
809 src_full_range: src_full,
810 dst_full_range: dst_full,
811 };
812 (src_params, dst_params)
813}
814
815fn partial_dst_access(dst: &TensorDyn, crop: &ResolvedCrop) -> edgefirst_tensor::CpuAccess {
819 let full = |r: &Rect| {
820 r.left == 0 && r.top == 0 && Some(r.width) == dst.width() && Some(r.height) == dst.height()
821 };
822 let keeps_outside = crop.dst_color.is_none() && crop.dst_rect.is_some_and(|r| !full(&r));
823 if dst.is_view() || keeps_outside {
824 edgefirst_tensor::CpuAccess::ReadWrite
825 } else {
826 edgefirst_tensor::CpuAccess::None
827 }
828}
829
830pub(crate) fn convert_dst_access(
835 src: Option<PixelFormat>,
836 dst_fmt: PixelFormat,
837 dst_dtype: DType,
838) -> edgefirst_tensor::CpuAccess {
839 use edgefirst_tensor::CpuAccess;
840 let swizzles = dst_fmt == PixelFormat::Bgra && src != Some(PixelFormat::Bgra);
841 match dst_dtype {
842 DType::I8 => CpuAccess::ReadWrite,
843 DType::U8 if swizzles => CpuAccess::ReadWrite,
844 _ => CpuAccess::Write,
845 }
846}
847
848fn require_draw_access(dst: &TensorDyn, background: Option<&TensorDyn>) -> Result<()> {
851 dst.cpu_access()
852 .require(edgefirst_tensor::CpuAccess::ReadWrite)?;
853 if let Some(bg) = background {
854 bg.cpu_access().require(edgefirst_tensor::CpuAccess::Read)?;
855 }
856 Ok(())
857}
858
859impl ImageProcessorTrait for CPUProcessor {
860 fn convert(
861 &mut self,
862 src: &TensorDyn,
863 dst: &mut TensorDyn,
864 options: Option<crate::ConvertOptions>,
865 ) -> Result<()> {
866 src.cpu_access()
867 .require(edgefirst_tensor::CpuAccess::Read)?;
868 dst.cpu_access()
869 .require(edgefirst_tensor::CpuAccess::Write)?;
870 let options = options.unwrap_or_default();
871 let (rotation, flip, crop) = (
872 options.get_rotation(),
873 options.get_flip(),
874 options.get_crop(),
875 );
876 let crop = crop.resolve(
877 src.width().unwrap_or(0),
878 src.height().unwrap_or(0),
879 dst.width().unwrap_or(0),
880 dst.height().unwrap_or(0),
881 options.get_placement(),
882 )?;
883 dst.cpu_access().require(partial_dst_access(dst, &crop))?;
884 let output_map = crate::options::output_map_for(&options, dst, crop.dst_rect.is_some())?;
885 if output_map.is_none() {
886 if let Some(dst_fmt) = dst.format() {
887 dst.cpu_access()
888 .require(convert_dst_access(src.format(), dst_fmt, dst.dtype()))?;
889 }
890 }
891 let processor_options = self.options;
892 if let Some(mode) = options.get_interpolation() {
893 self.options = self.options.resize_alg(Self::resize_alg(mode));
894 }
895 let result = match output_map {
896 None => self.convert_impl(src, dst, rotation, flip, crop),
897 Some((map, pad_output)) => {
898 self.convert_mapped(src, dst, rotation, flip, crop, &map, pad_output)
899 }
900 };
901 self.options = processor_options;
902 result
903 }
904
905 fn draw_decoded_masks(
906 &mut self,
907 dst: &mut TensorDyn,
908 detect: &[DetectBox],
909 segmentation: &[Segmentation],
910 overlay: crate::MaskOverlay<'_>,
911 ) -> Result<()> {
912 require_draw_access(dst, overlay.background)?;
916 prepare_dst_base_cpu(dst, overlay.background)?;
917 let dst = dst.as_typed_mut::<u8>().ok_or(Error::NotAnImage)?;
918 self.draw_decoded_masks_impl(
919 dst,
920 detect,
921 segmentation,
922 overlay.opacity,
923 overlay.color_mode,
924 )
925 }
926
927 fn draw_proto_masks(
928 &mut self,
929 dst: &mut TensorDyn,
930 detect: &[DetectBox],
931 proto_data: &ProtoData,
932 overlay: crate::MaskOverlay<'_>,
933 ) -> Result<()> {
934 require_draw_access(dst, overlay.background)?;
935 prepare_dst_base_cpu(dst, overlay.background)?;
936 let dst = dst.as_typed_mut::<u8>().ok_or(Error::NotAnImage)?;
937 self.draw_proto_masks_impl(
938 dst,
939 detect,
940 proto_data,
941 overlay.opacity,
942 overlay.content_region,
943 overlay.color_mode,
944 )
945 }
946
947 fn set_class_colors(&mut self, colors: &[[u8; 4]]) -> Result<()> {
948 for (c, new_c) in self.colors.iter_mut().zip(colors.iter()) {
949 *c = *new_c;
950 }
951 Ok(())
952 }
953
954 fn set_box_line_width(&mut self, width: u32) -> Result<()> {
955 crate::check_box_line_width(width)?;
956 self.box_line_width = usize::try_from(width).unwrap_or(usize::MAX);
957 Ok(())
958 }
959}
960
961impl CPUProcessor {
963 #[allow(clippy::too_many_arguments)]
966 fn convert_mapped(
967 &mut self,
968 src: &TensorDyn,
969 dst: &mut TensorDyn,
970 rotation: Rotation,
971 flip: Flip,
972 crop: ResolvedCrop,
973 map: &crate::options::OutputMap,
974 pad_output: Option<[f32; 4]>,
975 ) -> Result<()> {
976 let src_u8 = src
977 .as_typed::<u8>()
978 .ok_or_else(|| Error::NotSupported(format!("dtype {} source", src.dtype())))?;
979 let src_fmt = src.format().ok_or(Error::NotAnImage)?;
980 let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
981 let dw = dst.width().ok_or(Error::NotAnImage)?;
982 let dh = dst.height().ok_or(Error::NotAnImage)?;
983 let mut tmp = self.take_scratch(dw, dh, dst_fmt)?;
984 let (src_params, dst_params) = color_params(src, dst);
985 let result = (|| {
986 let tmp_u8 = tmp.as_typed_mut::<u8>().unwrap();
987 self.convert_u8(
988 src_u8, tmp_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
989 )?;
990 normalize::write_mapped(tmp.as_typed::<u8>().unwrap(), dst, map)?;
991 if let (Some(pad), Some(rect)) = (pad_output, crop.dst_rect) {
992 let values = map.pad_values([0; 4], Some(pad));
993 normalize::fill_outside(dst, rect, values)?;
994 }
995 Ok(())
996 })();
997 self.widen_scratch = Some(tmp);
998 result
999 }
1000
1001 fn take_scratch(&mut self, w: usize, h: usize, fmt: PixelFormat) -> Result<TensorDyn> {
1003 let matches = self.widen_scratch.as_ref().is_some_and(|t| {
1004 t.width() == Some(w) && t.height() == Some(h) && t.format() == Some(fmt)
1005 });
1006 if matches {
1007 return Ok(self.widen_scratch.take().unwrap());
1008 }
1009 Ok(TensorDyn::image(
1010 w,
1011 h,
1012 fmt,
1013 DType::U8,
1014 Some(TensorMemory::Mem),
1015 edgefirst_tensor::CpuAccess::ReadWrite,
1016 )?)
1017 }
1018
1019 pub(crate) fn convert_impl(
1021 &mut self,
1022 src: &TensorDyn,
1023 dst: &mut TensorDyn,
1024 rotation: Rotation,
1025 flip: Flip,
1026 crop: ResolvedCrop,
1027 ) -> Result<()> {
1028 let src_fmt = src.format().ok_or(Error::NotAnImage)?;
1029 let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
1030
1031 let (src_params, dst_params) = color_params(src, dst);
1032 match (src.dtype(), dst.dtype()) {
1033 (DType::U8, DType::U8) => {
1034 let src = src.as_typed::<u8>().unwrap();
1035 let dst = dst.as_typed_mut::<u8>().unwrap();
1036 self.convert_u8(
1037 src, dst, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
1038 )
1039 }
1040 (DType::U8, DType::I8) => {
1041 let src_u8 = src.as_typed::<u8>().unwrap();
1044 let dst_i8 = dst.as_typed_mut::<i8>().unwrap();
1045 let dst_u8 = unsafe { &mut *(dst_i8 as *mut Tensor<i8> as *mut Tensor<u8>) };
1049 self.convert_u8(
1050 src_u8, dst_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
1051 )?;
1052 apply_int8_xor_bias_rows(dst_u8, dst_fmt)
1054 }
1055 (DType::U8, d @ (DType::F32 | DType::F16)) => {
1056 let src_u8 = src.as_typed::<u8>().unwrap();
1057 let dw = dst.width().ok_or(Error::NotAnImage)?;
1058 let dh = dst.height().ok_or(Error::NotAnImage)?;
1059 let mut tmp = self.take_scratch(dw, dh, dst_fmt)?;
1064 {
1065 let tmp_u8 = tmp.as_typed_mut::<u8>().unwrap();
1066 self.convert_u8(
1067 src_u8, tmp_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params,
1068 dst_params,
1069 )?;
1070 }
1071 {
1081 let tmp_u8 = tmp.as_typed::<u8>().unwrap();
1082 let (rows, row_len) = logical_surface(tmp_u8)?;
1083 let src_stride = tensor_row_stride(tmp_u8);
1084 let dst_stride_bytes = dst.effective_row_stride().ok_or(Error::NotAnImage)?;
1085 let src_map = tmp_u8.map_read()?;
1086 guard_plane(
1087 src_map.as_slice().len(),
1088 src_stride,
1089 rows,
1090 row_len,
1091 "widen src",
1092 )?;
1093 let src_rows = src_map.as_slice().chunks(src_stride).take(rows);
1094 let elem = d.size();
1095 if !dst_stride_bytes.is_multiple_of(elem) {
1096 return Err(Error::InvalidShape(format!(
1097 "{d} destination row stride {dst_stride_bytes} is not a multiple of \
1098 the element size {elem}"
1099 )));
1100 }
1101 let dst_stride = dst_stride_bytes / elem;
1102 match d {
1103 DType::F32 => {
1104 let dst_t = dst.as_typed_mut::<f32>().unwrap();
1105 let mut dst_map = dst_t.map_for_write()?;
1106 guard_plane(
1108 dst_map.as_slice().len(),
1109 dst_stride,
1110 rows,
1111 row_len,
1112 "widen f32 dst",
1113 )?;
1114 for (s, dr) in
1118 src_rows.zip(dst_map.as_mut_slice().chunks_mut(dst_stride))
1119 {
1120 simd::widen_u8_to_f32_norm(s, &mut dr[..row_len]);
1121 }
1122 }
1123 DType::F16 => {
1124 let dst_t = dst.as_typed_mut::<half::f16>().unwrap();
1125 let mut dst_map = dst_t.map_for_write()?;
1126 guard_plane(
1127 dst_map.as_slice().len(),
1128 dst_stride,
1129 rows,
1130 row_len,
1131 "widen f16 dst",
1132 )?;
1133 for (s, dr) in
1138 src_rows.zip(dst_map.as_mut_slice().chunks_mut(dst_stride))
1139 {
1140 simd::widen_u8_to_f16_norm(s, &mut dr[..row_len]);
1141 }
1142 }
1143 _ => unreachable!(),
1144 }
1145 }
1146 self.widen_scratch = Some(tmp);
1147 Ok(())
1148 }
1149 (s, d) => Err(Error::NotSupported(format!("dtype {s} -> {d}",))),
1150 }
1151 }
1152
1153 fn reuse_or_alloc_image(
1159 cached: Option<Tensor<u8>>,
1160 w: usize,
1161 h: usize,
1162 fmt: PixelFormat,
1163 ) -> Result<Tensor<u8>> {
1164 if let Some(t) = cached {
1165 if t.width() == Some(w) && t.height() == Some(h) && t.format() == Some(fmt) {
1166 return Ok(t);
1167 }
1168 }
1169 Ok(Tensor::<u8>::image(
1170 w,
1171 h,
1172 fmt,
1173 Some(TensorMemory::Mem),
1174 edgefirst_tensor::CpuAccess::ReadWrite,
1175 )?)
1176 }
1177
1178 fn pre_resize_region(
1199 &self,
1200 src_fmt: PixelFormat,
1201 (src_w, src_h): (usize, usize),
1202 (dst_w, dst_h): (usize, usize),
1203 rotation: Rotation,
1204 crop: ResolvedCrop,
1205 ) -> Option<Rect> {
1206 use PixelFormat::{Nv12, Nv16, Nv24};
1207
1208 if !matches!(src_fmt, Nv12 | Nv16 | Nv24) {
1209 return None;
1210 }
1211 let r = crop.src_rect?;
1212 let full_src = Rect {
1213 left: 0,
1214 top: 0,
1215 width: src_w,
1216 height: src_h,
1217 };
1218 if r == full_src {
1219 return None;
1220 }
1221
1222 let d = crop.dst_rect.unwrap_or(Rect {
1227 left: 0,
1228 top: 0,
1229 width: dst_w,
1230 height: dst_h,
1231 });
1232 let (dst_x, dst_y) = match rotation {
1233 Rotation::None | Rotation::Rotate180 => (d.width, d.height),
1234 Rotation::Clockwise90 | Rotation::CounterClockwise90 => (d.height, d.width),
1235 };
1236 let halo_x = self.filter_halo(r.width, dst_x)?;
1237 let halo_y = self.filter_halo(r.height, dst_y)?;
1238
1239 let mut left = r.left.saturating_sub(halo_x);
1240 let mut top = r.top.saturating_sub(halo_y);
1241 let mut right = (r.left + r.width + halo_x).min(src_w);
1242 let mut bottom = (r.top + r.height + halo_y).min(src_h);
1243
1244 let (align_x, align_y) = match src_fmt {
1250 Nv12 => (2, 2),
1251 Nv16 => (2, 1),
1252 _ => (1, 1),
1253 };
1254 left -= left % align_x;
1255 top -= top % align_y;
1256 right = right.next_multiple_of(align_x).min(src_w);
1257 bottom = bottom.next_multiple_of(align_y).min(src_h);
1258
1259 let grown = Rect {
1260 left,
1261 top,
1262 width: right - left,
1263 height: bottom - top,
1264 };
1265 (grown != full_src).then_some(grown)
1268 }
1269
1270 #[allow(clippy::too_many_arguments)]
1272 fn convert_u8(
1273 &mut self,
1274 src: &Tensor<u8>,
1275 dst: &mut Tensor<u8>,
1276 src_fmt: PixelFormat,
1277 dst_fmt: PixelFormat,
1278 rotation: Rotation,
1279 flip: Flip,
1280 crop: ResolvedCrop,
1281 src_params: ColorParams,
1282 dst_params: ColorParams,
1283 ) -> Result<()> {
1284 use PixelFormat::*;
1285
1286 #[cfg(test)]
1287 {
1288 self.last_tmp_dims = None;
1289 }
1290
1291 let src_w = src.width().unwrap();
1292 let src_h = src.height().unwrap();
1293 let dst_w = dst.width().unwrap();
1294 let dst_h = dst.height().unwrap();
1295
1296 crop.check_crop_dims(src_w, src_h, dst_w, dst_h)?;
1297
1298 let intermediate = match (src_fmt, dst_fmt) {
1300 (Nv12, Rgb) => Rgb,
1301 (Nv12, Rgba) => Rgba,
1302 (Nv12, Grey) => Grey,
1303 (Nv12, Yuyv) => Rgba,
1304 (Nv12, Nv16) => Rgba,
1305 (Nv12, PlanarRgb) => Rgb,
1306 (Nv12, PlanarRgba) => Rgba,
1307 (Nv16, PlanarRgb) => Rgb,
1308 (Nv16, PlanarRgba) => Rgba,
1309 (Nv24, PlanarRgb) => Rgb,
1310 (Nv24, PlanarRgba) => Rgba,
1311 (Yuyv, Rgb) => Rgb,
1312 (Yuyv, Rgba) => Rgba,
1313 (Yuyv, Grey) => Grey,
1314 (Yuyv, Yuyv) => Rgba,
1315 (Yuyv, PlanarRgb) => Rgb,
1316 (Yuyv, PlanarRgba) => Rgba,
1317 (Yuyv, Nv16) => Rgba,
1318 (Vyuy, Rgb) => Rgb,
1319 (Vyuy, Rgba) => Rgba,
1320 (Vyuy, Grey) => Grey,
1321 (Vyuy, Vyuy) => Rgba,
1322 (Vyuy, PlanarRgb) => Rgb,
1323 (Vyuy, PlanarRgba) => Rgba,
1324 (Vyuy, Nv16) => Rgba,
1325 (Rgba, Rgb) => Rgba,
1326 (Rgba, Rgba) => Rgba,
1327 (Rgba, Grey) => Grey,
1328 (Rgba, Yuyv) => Rgba,
1329 (Rgba, PlanarRgb) => Rgba,
1330 (Rgba, PlanarRgba) => Rgba,
1331 (Rgba, Nv16) => Rgba,
1332 (Rgb, Rgb) => Rgb,
1333 (Rgb, Rgba) => Rgb,
1334 (Rgb, Grey) => Grey,
1335 (Rgb, Yuyv) => Rgb,
1336 (Rgb, PlanarRgb) => Rgb,
1337 (Rgb, PlanarRgba) => Rgb,
1338 (Rgb, Nv16) => Rgb,
1339 (Grey, Rgb) => Rgb,
1340 (Grey, Rgba) => Rgba,
1341 (Grey, Grey) => Grey,
1342 (Grey, Yuyv) => Grey,
1343 (Grey, PlanarRgb) => Grey,
1344 (Grey, PlanarRgba) => Grey,
1345 (Grey, Nv16) => Grey,
1346 (Nv12, Bgra) => Rgba,
1347 (Yuyv, Bgra) => Rgba,
1348 (Vyuy, Bgra) => Rgba,
1349 (Rgba, Bgra) => Rgba,
1350 (Rgb, Bgra) => Rgb,
1351 (Grey, Bgra) => Grey,
1352 (Bgra, Bgra) => Bgra,
1353 (Nv16, Rgb) => Rgb,
1354 (Nv16, Rgba) => Rgba,
1355 (Nv16, Bgra) => Rgba,
1356 (Nv24, Rgb) => Rgb,
1357 (Nv24, Rgba) => Rgba,
1358 (Nv24, Grey) => Grey,
1359 (Nv24, Bgra) => Rgba,
1360 (PlanarRgb, Rgb) => Rgb,
1361 (PlanarRgb, Rgba) => Rgb,
1362 (PlanarRgb, Bgra) => Rgb,
1363 (PlanarRgba, Rgb) => Rgba,
1364 (PlanarRgba, Rgba) => Rgba,
1365 (PlanarRgba, Bgra) => Rgba,
1366 (s, d) => {
1367 return Err(Error::NotSupported(format!("Conversion from {s} to {d}",)));
1368 }
1369 };
1370
1371 let need_resize_flip_rotation = rotation != Rotation::None
1372 || flip != Flip::None
1373 || src_w != dst_w
1374 || src_h != dst_h
1375 || crop.src_rect.is_some_and(|c| {
1376 c != Rect {
1377 left: 0,
1378 top: 0,
1379 width: src_w,
1380 height: src_h,
1381 }
1382 })
1383 || crop.dst_rect.is_some_and(|c| {
1384 c != Rect {
1385 left: 0,
1386 top: 0,
1387 width: dst_w,
1388 height: dst_h,
1389 }
1390 });
1391
1392 let direct_is_yuv_src = matches!(src_fmt, Nv12 | Nv16 | Nv24 | Yuyv | Vyuy);
1395 let direct_params = if direct_is_yuv_src {
1396 src_params
1397 } else {
1398 dst_params
1399 };
1400
1401 let full_dst_rect = Rect {
1418 left: 0,
1419 top: 0,
1420 width: dst_w,
1421 height: dst_h,
1422 };
1423 let fused_region = if rotation != Rotation::None || flip != Flip::None {
1424 None
1425 } else {
1426 match crop.src_rect {
1427 None if src_w == dst_w && src_h == dst_h => Some(None),
1428 None => None,
1429 Some(r)
1430 if r.width == dst_w
1431 && r.height == dst_h
1432 && crop.dst_rect.is_none_or(|d| d == full_dst_rect)
1433 && chroma_alignment_ok(src_fmt, r) =>
1434 {
1435 Some(Some(r))
1436 }
1437 Some(_) => None,
1438 }
1439 };
1440 if let Some(region) = fused_region {
1441 if matches!(src_fmt, Nv12 | Nv16 | Nv24) && matches!(dst_fmt, PlanarRgb | PlanarRgba) {
1442 #[cfg(test)]
1443 {
1444 self.fused_hits += 1;
1445 }
1446 return self.convert_nv_to_planar_fused(
1447 src,
1448 dst,
1449 src_fmt,
1450 dst_fmt,
1451 direct_params,
1452 region,
1453 );
1454 }
1455 }
1456
1457 if !need_resize_flip_rotation && Self::support_conversion_pf(src_fmt, dst_fmt) {
1459 return Self::convert_format_pf(src, dst, src_fmt, dst_fmt, direct_params);
1460 }
1461
1462 if dst_fmt == Yuyv && !dst_w.is_multiple_of(2) {
1464 return Err(Error::NotSupported(format!(
1465 "{} destination must have width divisible by 2",
1466 dst_fmt,
1467 )));
1468 }
1469
1470 let mut cached_tmp = self.convert_tmp.take();
1478 let mut cached_tmp2 = self.convert_tmp2.take();
1479
1480 let pre_region = if intermediate != src_fmt {
1485 self.pre_resize_region(src_fmt, (src_w, src_h), (dst_w, dst_h), rotation, crop)
1486 } else {
1487 None
1488 };
1489
1490 let tmp_holder: Option<Tensor<u8>> = if intermediate != src_fmt {
1492 let _s = tracing::trace_span!(
1493 "image.convert.cpu.format_convert",
1494 from = ?src_fmt,
1495 to = ?intermediate,
1496 pass = "pre_resize",
1497 )
1498 .entered();
1499 let (tmp_w, tmp_h) = pre_region.map_or((src_w, src_h), |g| (g.width, g.height));
1500 let mut t = Self::reuse_or_alloc_image(cached_tmp.take(), tmp_w, tmp_h, intermediate)?;
1501 #[cfg(test)]
1502 {
1503 self.last_tmp_dims = Some((tmp_w, tmp_h));
1504 }
1505 match pre_region {
1506 Some(g) => {
1507 let mut sub = Self::reuse_or_alloc_image(
1508 self.convert_src_sub.take(),
1509 g.width,
1510 g.height,
1511 src_fmt,
1512 )?;
1513 {
1514 let _s = tracing::trace_span!(
1515 "image.convert.cpu.extract_region",
1516 region_w = g.width,
1517 region_h = g.height,
1518 )
1519 .entered();
1520 Self::extract_nv_region(src, &mut sub, src_fmt, g)?;
1521 }
1522 Self::convert_format_pf(&sub, &mut t, src_fmt, intermediate, src_params)?;
1523 self.convert_src_sub = Some(sub);
1524 }
1525 None => Self::convert_format_pf(src, &mut t, src_fmt, intermediate, src_params)?,
1526 }
1527 Some(t)
1528 } else {
1529 None
1530 };
1531
1532 let crop = match (pre_region, crop.src_rect) {
1537 (Some(g), Some(r)) => ResolvedCrop {
1538 src_rect: Some(Rect {
1539 left: r.left - g.left,
1540 top: r.top - g.top,
1541 ..r
1542 }),
1543 ..crop
1544 },
1545 _ => crop,
1546 };
1547 let (tmp, tmp_fmt): (&Tensor<u8>, PixelFormat) = match &tmp_holder {
1548 Some(t) => (t, intermediate),
1549 None => (src, src_fmt),
1550 };
1551
1552 debug_assert!(matches!(tmp_fmt, Rgb | Rgba | Grey));
1554 if tmp_fmt == dst_fmt {
1555 let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
1556 self.resize_flip_rotate_pf(tmp, dst, dst_fmt, rotation, flip, crop)?;
1557 } else if !need_resize_flip_rotation {
1558 let _s = tracing::trace_span!(
1559 "image.convert.cpu.format_convert",
1560 from = ?tmp_fmt,
1561 to = ?dst_fmt,
1562 pass = "direct",
1563 )
1564 .entered();
1565 Self::convert_format_pf(tmp, dst, tmp_fmt, dst_fmt, dst_params)?;
1566 } else {
1567 let mut tmp2 = Self::reuse_or_alloc_image(cached_tmp2.take(), dst_w, dst_h, tmp_fmt)?;
1568 if crop.dst_rect.is_some_and(|c| {
1569 c != Rect {
1570 left: 0,
1571 top: 0,
1572 width: dst_w,
1573 height: dst_h,
1574 }
1575 }) && crop.dst_color.is_none()
1576 {
1577 Self::convert_format_pf(dst, &mut tmp2, dst_fmt, tmp_fmt, dst_params)?;
1578 }
1579 {
1580 let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
1581 self.resize_flip_rotate_pf(tmp, &mut tmp2, tmp_fmt, rotation, flip, crop)?;
1582 }
1583 {
1584 let _s = tracing::trace_span!(
1585 "image.convert.cpu.format_convert",
1586 from = ?tmp_fmt,
1587 to = ?dst_fmt,
1588 pass = "post_resize",
1589 )
1590 .entered();
1591 Self::convert_format_pf(&tmp2, dst, tmp_fmt, dst_fmt, dst_params)?;
1592 }
1593 cached_tmp2 = Some(tmp2);
1594 }
1595 if let Some(t) = tmp_holder {
1598 cached_tmp = Some(t);
1599 }
1600 self.convert_tmp = cached_tmp;
1601 self.convert_tmp2 = cached_tmp2;
1602
1603 if let (Some(dst_rect), Some(dst_color)) = (crop.dst_rect, crop.dst_color) {
1604 let full_rect = Rect {
1605 left: 0,
1606 top: 0,
1607 width: dst_w,
1608 height: dst_h,
1609 };
1610 if dst_rect != full_rect {
1611 Self::fill_image_outside_crop_u8(dst, dst_color, dst_rect, false)?;
1612 }
1613 }
1614
1615 Ok(())
1616 }
1617
1618 fn draw_decoded_masks_impl(
1619 &mut self,
1620 dst: &mut Tensor<u8>,
1621 detect: &[DetectBox],
1622 segmentation: &[Segmentation],
1623 opacity: f32,
1624 color_mode: crate::ColorMode,
1625 ) -> Result<()> {
1626 let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
1627 if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
1628 return Err(crate::Error::NotSupported(
1629 "CPU image rendering only supports RGBA or RGB images".to_string(),
1630 ));
1631 }
1632
1633 let _timer = FunctionTimer::new("CPUProcessor::draw_decoded_masks");
1634
1635 let dst_w = dst.width().unwrap();
1636 let dst_h = dst.height().unwrap();
1637 let dst_rs = tensor_row_stride(dst);
1638 let dst_c = dst_fmt.channels();
1639
1640 let mut map = dst.map_mut()?;
1641 let dst_slice = map.as_mut_slice();
1642
1643 self.render_box(
1644 dst_w, dst_h, dst_rs, dst_c, dst_slice, detect, opacity, color_mode,
1645 )?;
1646
1647 if segmentation.is_empty() {
1648 return Ok(());
1649 }
1650
1651 if let Some(bad) = segmentation
1658 .iter()
1659 .find(|s| s.segmentation.shape().len() != 3)
1660 {
1661 return Err(crate::Error::InvalidShape(format!(
1662 "segmentation must be [H, W, C], got {:?}",
1663 bad.segmentation.shape()
1664 )));
1665 }
1666
1667 let is_semantic = segmentation[0].segmentation.shape()[2] > 1;
1670
1671 if is_semantic {
1672 self.render_modelpack_segmentation(
1673 dst_w,
1674 dst_h,
1675 dst_rs,
1676 dst_c,
1677 dst_slice,
1678 &segmentation[0],
1679 opacity,
1680 )?;
1681 } else {
1682 for (idx, (seg, det)) in segmentation.iter().zip(detect).enumerate() {
1683 let color_index = color_mode.index(idx, det.label);
1684 self.render_yolo_segmentation(
1685 dst_w,
1686 dst_h,
1687 dst_rs,
1688 dst_c,
1689 dst_slice,
1690 seg,
1691 color_index,
1692 opacity,
1693 )?;
1694 }
1695 }
1696
1697 Ok(())
1698 }
1699
1700 fn draw_proto_masks_impl(
1701 &mut self,
1702 dst: &mut Tensor<u8>,
1703 detect: &[DetectBox],
1704 proto_data: &ProtoData,
1705 opacity: f32,
1706 content_region: Option<[f32; 4]>,
1707 color_mode: crate::ColorMode,
1708 ) -> Result<()> {
1709 let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
1710 if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
1711 return Err(crate::Error::NotSupported(
1712 "CPU image rendering only supports RGBA or RGB images".to_string(),
1713 ));
1714 }
1715
1716 let _timer = FunctionTimer::new("CPUProcessor::draw_proto_masks");
1717
1718 let dst_w = dst.width().unwrap();
1719 let dst_h = dst.height().unwrap();
1720 let dst_rs = tensor_row_stride(dst);
1721 let channels = dst_fmt.channels();
1722
1723 let mut map = dst.map_mut()?;
1724 let dst_slice = map.as_mut_slice();
1725
1726 self.render_box(
1727 dst_w, dst_h, dst_rs, channels, dst_slice, detect, opacity, color_mode,
1728 )?;
1729
1730 if detect.is_empty() {
1731 return Ok(());
1732 }
1733 let proto_shape = proto_data.protos.shape();
1734 if proto_shape.len() != 3 {
1735 return Err(Error::InvalidShape(format!(
1736 "protos tensor must be rank-3, got {proto_shape:?}"
1737 )));
1738 }
1739 let proto_h = proto_shape[0];
1740 let proto_w = proto_shape[1];
1741 let num_protos = proto_shape[2];
1742 let coeff_shape = proto_data.mask_coefficients.shape();
1743 if coeff_shape.len() != 2 {
1744 return Err(Error::InvalidShape(format!(
1745 "mask_coefficients tensor must be rank-2, got {coeff_shape:?}"
1746 )));
1747 }
1748 if coeff_shape[0] == 0 {
1750 return Ok(());
1751 }
1752 if coeff_shape[1] != num_protos {
1753 return Err(Error::InvalidShape(format!(
1754 "mask_coefficients second dimension must match num_protos \
1755 ({num_protos}), got {coeff_shape:?}"
1756 )));
1757 }
1758
1759 let coeff_f32: Vec<f32> = match proto_data.mask_coefficients.dtype() {
1761 DType::F32 => {
1762 let t = proto_data.mask_coefficients.as_typed::<f32>().expect("F32");
1763 let m = t.map_read()?;
1764 m.as_slice().to_vec()
1765 }
1766 DType::F16 => {
1767 let t = proto_data
1768 .mask_coefficients
1769 .as_typed::<half::f16>()
1770 .expect("F16");
1771 let m = t.map_read()?;
1772 m.as_slice().iter().map(|v| v.to_f32()).collect()
1773 }
1774 DType::I8 => {
1775 let t = proto_data.mask_coefficients.as_typed::<i8>().expect("I8");
1776 let m = t.map_read()?;
1777 if let Some(q) = t.quantization() {
1778 use edgefirst_tensor::QuantMode;
1779 let (scale, zp) = match q.mode() {
1780 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1781 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1782 other => {
1783 return Err(Error::NotSupported(format!(
1784 "I8 mask_coefficients quantization mode {other:?} not supported"
1785 )));
1786 }
1787 };
1788 m.as_slice()
1789 .iter()
1790 .map(|&v| (v as f32 - zp) * scale)
1791 .collect()
1792 } else {
1793 m.as_slice().iter().map(|&v| v as f32).collect()
1794 }
1795 }
1796 DType::I16 => {
1797 let t = proto_data.mask_coefficients.as_typed::<i16>().expect("I16");
1798 let m = t.map_read()?;
1799 if let Some(q) = t.quantization() {
1800 use edgefirst_tensor::QuantMode;
1801 let (scale, zp) = match q.mode() {
1802 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1803 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1804 other => {
1805 return Err(Error::NotSupported(format!(
1806 "I16 mask_coefficients quantization mode {other:?} not supported"
1807 )));
1808 }
1809 };
1810 m.as_slice()
1811 .iter()
1812 .map(|&v| (v as f32 - zp) * scale)
1813 .collect()
1814 } else {
1815 m.as_slice().iter().map(|&v| v as f32).collect()
1816 }
1817 }
1818 other => {
1819 return Err(Error::InvalidShape(format!(
1820 "mask_coefficients dtype {other:?} not supported"
1821 )));
1822 }
1823 };
1824
1825 let (lx0, lx_range, ly0, ly_range) = match content_region {
1827 Some([lx0, ly0, lx1, ly1]) => (lx0, lx1 - lx0, ly0, ly1 - ly0),
1828 None => (0.0_f32, 1.0_f32, 0.0_f32, 1.0_f32),
1829 };
1830
1831 match proto_data.protos.dtype() {
1834 DType::F32 => {
1835 let t = proto_data.protos.as_typed::<f32>().expect("F32");
1836 let m = t.map_read()?;
1837 self.draw_proto_masks_inner(
1838 dst_slice,
1839 dst_w,
1840 dst_h,
1841 dst_rs,
1842 channels,
1843 detect,
1844 m.as_slice(),
1845 &coeff_f32,
1846 proto_h,
1847 proto_w,
1848 num_protos,
1849 opacity,
1850 (lx0, lx_range, ly0, ly_range),
1851 color_mode,
1852 0.0_f32,
1853 |p: &f32, _| *p,
1854 );
1855 }
1856 DType::F16 => {
1857 let t = proto_data.protos.as_typed::<half::f16>().expect("F16");
1858 let m = t.map_read()?;
1859 self.draw_proto_masks_inner(
1860 dst_slice,
1861 dst_w,
1862 dst_h,
1863 dst_rs,
1864 channels,
1865 detect,
1866 m.as_slice(),
1867 &coeff_f32,
1868 proto_h,
1869 proto_w,
1870 num_protos,
1871 opacity,
1872 (lx0, lx_range, ly0, ly_range),
1873 color_mode,
1874 0.0_f32,
1875 |p: &half::f16, _| p.to_f32(),
1876 );
1877 }
1878 DType::I8 => {
1879 use edgefirst_tensor::QuantMode;
1880 let t = proto_data.protos.as_typed::<i8>().expect("I8");
1881 let m = t.map_read()?;
1882 let quant = t.quantization().ok_or_else(|| {
1883 Error::InvalidShape("I8 protos require quantization metadata".into())
1884 })?;
1885 let (scale, zp) = match quant.mode() {
1886 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1887 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1888 QuantMode::PerChannel { axis, .. }
1889 | QuantMode::PerChannelSymmetric { axis, .. } => {
1890 return Err(Error::NotSupported(format!(
1891 "per-channel quantization (axis={axis}) in draw_proto_masks \
1892 CPU path not yet supported"
1893 )));
1894 }
1895 };
1896 self.draw_proto_masks_inner(
1897 dst_slice,
1898 dst_w,
1899 dst_h,
1900 dst_rs,
1901 channels,
1902 detect,
1903 m.as_slice(),
1904 &coeff_f32,
1905 proto_h,
1906 proto_w,
1907 num_protos,
1908 opacity,
1909 (lx0, lx_range, ly0, ly_range),
1910 color_mode,
1911 scale,
1912 move |p: &i8, _| (*p as f32) - zp,
1913 );
1914 }
1915 other => {
1916 return Err(Error::InvalidShape(format!(
1917 "proto tensor dtype {other:?} not supported"
1918 )));
1919 }
1920 }
1921
1922 Ok(())
1923 }
1924
1925 #[allow(clippy::too_many_arguments)]
1926 fn draw_proto_masks_inner<P: Copy>(
1927 &self,
1928 dst_slice: &mut [u8],
1929 dst_w: usize,
1930 dst_h: usize,
1931 dst_rs: usize,
1932 channels: usize,
1933 detect: &[DetectBox],
1934 protos: &[P],
1935 coeff_all_f32: &[f32],
1936 proto_h: usize,
1937 proto_w: usize,
1938 num_protos: usize,
1939 opacity: f32,
1940 letterbox_xy: (f32, f32, f32, f32),
1941 color_mode: crate::ColorMode,
1942 acc_scale: f32,
1943 load_f32: impl Fn(&P, f32) -> f32 + Copy,
1944 ) {
1945 let (lx0, lx_range, ly0, ly_range) = letterbox_xy;
1946 let stride_y = proto_w * num_protos;
1947 for (idx, det) in detect.iter().enumerate() {
1948 let coeff = &coeff_all_f32[idx * num_protos..(idx + 1) * num_protos];
1949 let color_index = color_mode.index(idx, det.label);
1950 let color = self.colors[color_index % self.colors.len()];
1951 let alpha = if opacity == 1.0 {
1952 color[3] as u16
1953 } else {
1954 (color[3] as f32 * opacity).round() as u16
1955 };
1956
1957 let start_x = (dst_w as f32 * det.bbox.xmin).round() as usize;
1958 let start_y = (dst_h as f32 * det.bbox.ymin).round() as usize;
1959 let end_x = ((dst_w as f32 * det.bbox.xmax).round() as usize).min(dst_w);
1960 let end_y = ((dst_h as f32 * det.bbox.ymax).round() as usize).min(dst_h);
1961
1962 for y in start_y..end_y {
1963 for x in start_x..end_x {
1964 let px = (lx0 + (x as f32 / dst_w as f32) * lx_range) * proto_w as f32 - 0.5;
1965 let py = (ly0 + (y as f32 / dst_h as f32) * ly_range) * proto_h as f32 - 0.5;
1966
1967 let x0 = (px.floor() as isize).clamp(0, proto_w as isize - 1) as usize;
1971 let y0 = (py.floor() as isize).clamp(0, proto_h as isize - 1) as usize;
1972 let x1 = (x0 + 1).min(proto_w - 1);
1973 let y1 = (y0 + 1).min(proto_h - 1);
1974 let fx = px - px.floor();
1975 let fy = py - py.floor();
1976 let w00 = (1.0 - fx) * (1.0 - fy);
1977 let w10 = fx * (1.0 - fy);
1978 let w01 = (1.0 - fx) * fy;
1979 let w11 = fx * fy;
1980 let b00 = y0 * stride_y + x0 * num_protos;
1981 let b10 = y0 * stride_y + x1 * num_protos;
1982 let b01 = y1 * stride_y + x0 * num_protos;
1983 let b11 = y1 * stride_y + x1 * num_protos;
1984 let mut acc = 0.0_f32;
1985 for p in 0..num_protos {
1986 let v00 = load_f32(&protos[b00 + p], 0.0);
1987 let v10 = load_f32(&protos[b10 + p], 0.0);
1988 let v01 = load_f32(&protos[b01 + p], 0.0);
1989 let v11 = load_f32(&protos[b11 + p], 0.0);
1990 let val = w00 * v00 + w10 * v10 + w01 * v01 + w11 * v11;
1991 acc += coeff[p] * val;
1992 }
1993 let final_acc = if acc_scale == 0.0 {
1994 acc
1995 } else {
1996 acc_scale * acc
1997 };
1998 let mask = 1.0 / (1.0 + (-final_acc).exp());
2002 if mask < 0.5 {
2003 continue;
2004 }
2005 let dst_index = y * dst_rs + x * channels;
2006 for c in 0..3 {
2007 dst_slice[dst_index + c] = ((color[c] as u16 * alpha
2008 + dst_slice[dst_index + c] as u16 * (255 - alpha))
2009 / 255) as u8;
2010 }
2011 }
2012 }
2013 }
2014 }
2015}