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 split_semi_planar(
262 bytes: &[u8],
263 stride: usize,
264 src_h: usize,
265 fmt: PixelFormat,
266) -> Result<(&[u8], &[u8])> {
267 let total_h = fmt.combined_plane_height(src_h).unwrap_or(src_h);
268 let need = stride.checked_mul(total_h).ok_or_else(|| {
269 Error::InvalidShape(format!(
270 "{fmt:?} plane size overflow (stride={stride}, h={src_h})"
271 ))
272 })?;
273 if bytes.len() < need {
274 return Err(Error::InvalidShape(format!(
275 "{fmt:?} source has {} bytes but needs {need} (stride={stride}, h={src_h})",
276 bytes.len()
277 )));
278 }
279 Ok(bytes.split_at(stride * src_h))
280}
281
282fn split_semi_planar_mut(
289 bytes: &mut [u8],
290 stride: usize,
291 dst_h: usize,
292 fmt: PixelFormat,
293) -> Result<(&mut [u8], &mut [u8])> {
294 let total_h = fmt.combined_plane_height(dst_h).unwrap_or(dst_h);
295 let need = stride.checked_mul(total_h).ok_or_else(|| {
296 Error::InvalidShape(format!(
297 "{fmt:?} plane size overflow (stride={stride}, combined_h={total_h})"
298 ))
299 })?;
300 if bytes.len() < need {
301 return Err(Error::InvalidShape(format!(
302 "{fmt:?} destination has {} bytes but needs {need} (stride={stride}, combined_h={total_h})",
303 bytes.len()
304 )));
305 }
306 Ok(bytes.split_at_mut(stride * dst_h))
307}
308
309fn guard_plane(
315 buf_len: usize,
316 stride: usize,
317 rows: usize,
318 row_bytes: usize,
319 what: &str,
320) -> Result<()> {
321 let need = stride.checked_mul(rows).ok_or_else(|| {
322 Error::InvalidShape(format!(
323 "{what} plane size overflow (stride={stride}, rows={rows})"
324 ))
325 })?;
326 if row_bytes > stride || buf_len < need {
327 return Err(Error::InvalidShape(format!(
328 "{what} buffer too small: {buf_len} bytes, need {need} (stride={stride}, rows={rows}, row_bytes={row_bytes})"
329 )));
330 }
331 Ok(())
332}
333
334pub(crate) fn apply_int8_xor_bias(data: &mut [u8], fmt: PixelFormat) {
339 use edgefirst_tensor::PixelLayout;
340 if !fmt.has_alpha() {
341 for b in data.iter_mut() {
342 *b ^= 0x80;
343 }
344 } else if fmt.layout() == PixelLayout::Planar {
345 let channels = fmt.channels();
347 let plane_size = data.len() / channels;
348 for b in data[..plane_size * (channels - 1)].iter_mut() {
349 *b ^= 0x80;
350 }
351 } else {
352 let channels = fmt.channels();
354 for pixel in data.chunks_exact_mut(channels) {
355 for b in &mut pixel[..channels - 1] {
356 *b ^= 0x80;
357 }
358 }
359 }
360}
361
362impl CPUProcessor {
363 pub fn new() -> Self {
365 Self::new_bilinear()
366 }
367
368 fn new_bilinear() -> Self {
370 let resizer = fast_image_resize::Resizer::new();
371 let options = fast_image_resize::ResizeOptions::new()
372 .resize_alg(fast_image_resize::ResizeAlg::Convolution(
373 fast_image_resize::FilterType::Bilinear,
374 ))
375 .use_alpha(false);
376
377 log::debug!("CPUConverter created");
378 Self {
379 resizer,
380 options,
381 colors: crate::DEFAULT_COLORS_U8,
382 widen_scratch: None,
383 resize_destride_scratch: Vec::new(),
384 resize_dst_destride_scratch: Vec::new(),
385 nv_strip_scratch: Vec::new(),
386 nv_strip_y_pack: Vec::new(),
387 nv_strip_uv_pack: Vec::new(),
388 #[cfg(test)]
389 fused_hits: 0,
390 #[cfg(test)]
391 last_tmp_dims: None,
392 convert_tmp: None,
393 convert_tmp2: None,
394 convert_src_sub: None,
395 }
396 }
397
398 pub fn new_nearest() -> Self {
400 let resizer = fast_image_resize::Resizer::new();
401 let options = fast_image_resize::ResizeOptions::new()
402 .resize_alg(fast_image_resize::ResizeAlg::Nearest)
403 .use_alpha(false);
404 log::debug!("CPUConverter created");
405 Self {
406 resizer,
407 options,
408 colors: crate::DEFAULT_COLORS_U8,
409 widen_scratch: None,
410 resize_destride_scratch: Vec::new(),
411 resize_dst_destride_scratch: Vec::new(),
412 nv_strip_scratch: Vec::new(),
413 nv_strip_y_pack: Vec::new(),
414 nv_strip_uv_pack: Vec::new(),
415 #[cfg(test)]
416 fused_hits: 0,
417 #[cfg(test)]
418 last_tmp_dims: None,
419 convert_tmp: None,
420 convert_tmp2: None,
421 convert_src_sub: None,
422 }
423 }
424
425 #[cfg(test)]
427 pub(super) fn fused_hits(&self) -> u64 {
428 self.fused_hits
429 }
430
431 #[cfg(test)]
434 pub(super) fn last_tmp_dims(&self) -> Option<(usize, usize)> {
435 self.last_tmp_dims
436 }
437
438 pub(crate) fn support_conversion_pf(src: PixelFormat, dst: PixelFormat) -> bool {
439 use PixelFormat::*;
440 matches!(
441 (src, dst),
442 (Nv12, Rgb)
443 | (Nv12, Rgba)
444 | (Nv12, Grey)
445 | (Nv16, Rgb)
446 | (Nv16, Rgba)
447 | (Nv16, Bgra)
448 | (Nv24, Rgb)
449 | (Nv24, Rgba)
450 | (Nv24, Grey)
451 | (Nv24, Bgra)
452 | (Yuyv, Rgb)
453 | (Yuyv, Rgba)
454 | (Yuyv, Grey)
455 | (Yuyv, Yuyv)
456 | (Yuyv, PlanarRgb)
457 | (Yuyv, PlanarRgba)
458 | (Yuyv, Nv16)
459 | (Vyuy, Rgb)
460 | (Vyuy, Rgba)
461 | (Vyuy, Grey)
462 | (Vyuy, Vyuy)
463 | (Vyuy, PlanarRgb)
464 | (Vyuy, PlanarRgba)
465 | (Vyuy, Nv16)
466 | (Rgba, Rgb)
467 | (Rgba, Rgba)
468 | (Rgba, Grey)
469 | (Rgba, Yuyv)
470 | (Rgba, PlanarRgb)
471 | (Rgba, PlanarRgba)
472 | (Rgba, Nv16)
473 | (Rgb, Rgb)
474 | (Rgb, Rgba)
475 | (Rgb, Grey)
476 | (Rgb, Yuyv)
477 | (Rgb, PlanarRgb)
478 | (Rgb, PlanarRgba)
479 | (Rgb, Nv16)
480 | (Grey, Rgb)
481 | (Grey, Rgba)
482 | (Grey, Grey)
483 | (Grey, Yuyv)
484 | (Grey, PlanarRgb)
485 | (Grey, PlanarRgba)
486 | (Grey, Nv16)
487 | (Nv12, Bgra)
488 | (Yuyv, Bgra)
489 | (Vyuy, Bgra)
490 | (Rgba, Bgra)
491 | (Rgb, Bgra)
492 | (Grey, Bgra)
493 | (Bgra, Bgra)
494 | (PlanarRgb, Rgb)
495 | (PlanarRgb, Rgba)
496 | (PlanarRgba, Rgb)
497 | (PlanarRgba, Rgba)
498 | (PlanarRgb, Bgra)
499 | (PlanarRgba, Bgra)
500 )
501 }
502
503 pub(crate) fn convert_format_pf(
505 src: &Tensor<u8>,
506 dst: &mut Tensor<u8>,
507 src_fmt: PixelFormat,
508 dst_fmt: PixelFormat,
509 cp: ColorParams,
510 ) -> Result<()> {
511 let _timer = FunctionTimer::new(format!(
512 "ImageProcessor::convert_format {} to {}",
513 src_fmt, dst_fmt,
514 ));
515
516 use PixelFormat::*;
517 match (src_fmt, dst_fmt) {
518 (Nv12, Rgb) => Self::convert_nv12_to_rgb(src, dst, cp),
519 (Nv12, Rgba) => Self::convert_nv12_to_rgba(src, dst, cp),
520 (Nv12, Grey) => Self::convert_nv12_to_grey(src, dst, cp),
521 (Yuyv, Rgb) => Self::convert_yuyv_to_rgb(src, dst, cp),
522 (Yuyv, Rgba) => Self::convert_yuyv_to_rgba(src, dst, cp),
523 (Yuyv, Grey) => Self::convert_yuyv_to_grey(src, dst, cp),
524 (Yuyv, Yuyv) => Self::copy_image(src, dst),
525 (Yuyv, PlanarRgb) => Self::convert_yuyv_to_8bps(src, dst, cp),
526 (Yuyv, PlanarRgba) => Self::convert_yuyv_to_prgba(src, dst, cp),
527 (Yuyv, Nv16) => Self::convert_yuyv_to_nv16(src, dst),
528 (Vyuy, Rgb) => Self::convert_vyuy_to_rgb(src, dst, cp),
529 (Vyuy, Rgba) => Self::convert_vyuy_to_rgba(src, dst, cp),
530 (Vyuy, Grey) => Self::convert_vyuy_to_grey(src, dst, cp),
531 (Vyuy, Vyuy) => Self::copy_image(src, dst),
532 (Vyuy, PlanarRgb) => Self::convert_vyuy_to_8bps(src, dst, cp),
533 (Vyuy, PlanarRgba) => Self::convert_vyuy_to_prgba(src, dst, cp),
534 (Vyuy, Nv16) => Self::convert_vyuy_to_nv16(src, dst),
535 (Rgba, Rgb) => Self::convert_rgba_to_rgb(src, dst),
536 (Rgba, Rgba) => Self::copy_image(src, dst),
537 (Rgba, Grey) => Self::convert_rgba_to_grey(src, dst),
538 (Rgba, Yuyv) => Self::convert_rgba_to_yuyv(src, dst, cp),
539 (Rgba, PlanarRgb) => Self::convert_rgba_to_8bps(src, dst),
540 (Rgba, PlanarRgba) => Self::convert_rgba_to_prgba(src, dst),
541 (Rgba, Nv16) => Self::convert_rgba_to_nv16(src, dst, cp),
542 (Rgb, Rgb) => Self::copy_image(src, dst),
543 (Rgb, Rgba) => Self::convert_rgb_to_rgba(src, dst),
544 (Rgb, Grey) => Self::convert_rgb_to_grey(src, dst),
545 (Rgb, Yuyv) => Self::convert_rgb_to_yuyv(src, dst, cp),
546 (Rgb, PlanarRgb) => Self::convert_rgb_to_8bps(src, dst),
547 (Rgb, PlanarRgba) => Self::convert_rgb_to_prgba(src, dst),
548 (Rgb, Nv16) => Self::convert_rgb_to_nv16(src, dst, cp),
549 (Grey, Rgb) => Self::convert_grey_to_rgb(src, dst),
550 (Grey, Rgba) => Self::convert_grey_to_rgba(src, dst),
551 (Grey, Grey) => Self::copy_image(src, dst),
552 (Grey, Yuyv) => Self::convert_grey_to_yuyv(src, dst, cp),
553 (Grey, PlanarRgb) => Self::convert_grey_to_8bps(src, dst),
554 (Grey, PlanarRgba) => Self::convert_grey_to_prgba(src, dst),
555 (Grey, Nv16) => Self::convert_grey_to_nv16(src, dst, cp),
556
557 (Nv16, Rgb) => Self::convert_nv16_to_rgb(src, dst, cp),
559 (Nv16, Rgba) => Self::convert_nv16_to_rgba(src, dst, cp),
560 (Nv24, Rgb) => Self::convert_nv24_to_rgb(src, dst, cp),
561 (Nv24, Rgba) => Self::convert_nv24_to_rgba(src, dst, cp),
562 (Nv24, Grey) => Self::convert_nv24_to_grey(src, dst, cp),
563 (PlanarRgb, Rgb) => Self::convert_8bps_to_rgb(src, dst),
564 (PlanarRgb, Rgba) => Self::convert_8bps_to_rgba(src, dst),
565 (PlanarRgba, Rgb) => Self::convert_prgba_to_rgb(src, dst),
566 (PlanarRgba, Rgba) => Self::convert_prgba_to_rgba(src, dst),
567
568 (Bgra, Bgra) => Self::copy_image(src, dst),
570 (Nv12, Bgra) => {
571 Self::convert_nv12_to_rgba(src, dst, cp)?;
572 Self::swizzle_rb_4chan(dst)
573 }
574 (Nv16, Bgra) => {
575 Self::convert_nv16_to_rgba(src, dst, cp)?;
576 Self::swizzle_rb_4chan(dst)
577 }
578 (Nv24, Bgra) => {
579 Self::convert_nv24_to_rgba(src, dst, cp)?;
580 Self::swizzle_rb_4chan(dst)
581 }
582 (Yuyv, Bgra) => {
583 Self::convert_yuyv_to_rgba(src, dst, cp)?;
584 Self::swizzle_rb_4chan(dst)
585 }
586 (Vyuy, Bgra) => {
587 Self::convert_vyuy_to_rgba(src, dst, cp)?;
588 Self::swizzle_rb_4chan(dst)
589 }
590 (Rgba, Bgra) => {
591 dst.map_mut()?.copy_from_slice(&src.map_read()?);
592 Self::swizzle_rb_4chan(dst)
593 }
594 (Rgb, Bgra) => {
595 Self::convert_rgb_to_rgba(src, dst)?;
596 Self::swizzle_rb_4chan(dst)
597 }
598 (Grey, Bgra) => {
599 Self::convert_grey_to_rgba(src, dst)?;
600 Self::swizzle_rb_4chan(dst)
601 }
602 (PlanarRgb, Bgra) => {
603 Self::convert_8bps_to_rgba(src, dst)?;
604 Self::swizzle_rb_4chan(dst)
605 }
606 (PlanarRgba, Bgra) => {
607 Self::convert_prgba_to_rgba(src, dst)?;
608 Self::swizzle_rb_4chan(dst)
609 }
610
611 (s, d) => Err(Error::NotSupported(format!("Conversion from {s} to {d}",))),
612 }
613 }
614
615 pub(crate) fn fill_image_outside_crop_u8(
617 dst: &mut Tensor<u8>,
618 rgba: [u8; 4],
619 crop: Rect,
620 ) -> Result<()> {
621 let dst_fmt = dst.format().unwrap();
622 let dst_w = dst.width().unwrap();
623 let dst_h = dst.height().unwrap();
624 let cm = crate::colorimetry::resolve_colorimetry(dst.colorimetry(), dst.height());
628 let cp = ColorParams {
629 matrix: crate::colorimetry::yuv_matrix(cm.encoding.unwrap()),
630 range: crate::colorimetry::yuv_range(cm.range.unwrap()),
631 encoding: cm.encoding.unwrap(),
632 range_kind: cm.range.unwrap(),
633 src_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
634 dst_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
635 };
636 let mut dst_map = dst.map_mut()?;
637 let dst_tup = (dst_map.as_mut_slice(), dst_w, dst_h);
638 Self::fill_outside_crop_dispatch(dst_tup, dst_fmt, rgba, crop, cp)
639 }
640
641 fn fill_outside_crop_dispatch(
643 dst: (&mut [u8], usize, usize),
644 fmt: PixelFormat,
645 rgba: [u8; 4],
646 crop: Rect,
647 cp: ColorParams,
648 ) -> Result<()> {
649 use PixelFormat::*;
650 match fmt {
651 Rgba | Bgra => Self::fill_image_outside_crop_(dst, rgba, crop),
652 Rgb => Self::fill_image_outside_crop_(dst, Self::rgba_to_rgb(rgba), crop),
653 Grey => Self::fill_image_outside_crop_(dst, Self::rgba_to_grey(rgba), crop),
654 Yuyv => Self::fill_image_outside_crop_(
655 (dst.0, dst.1 / 2, dst.2),
656 Self::rgba_to_yuyv(rgba, cp),
657 Rect::new(crop.left / 2, crop.top, crop.width.div_ceil(2), crop.height),
658 ),
659 PlanarRgb => Self::fill_image_outside_crop_planar(dst, Self::rgba_to_rgb(rgba), crop),
660 PlanarRgba => Self::fill_image_outside_crop_planar(dst, rgba, crop),
661 Nv16 => {
662 let yuyv = Self::rgba_to_yuyv(rgba, cp);
663 Self::fill_image_outside_crop_yuv_semiplanar(dst, yuyv[0], [yuyv[1], yuyv[3]], crop)
664 }
665 _ => Err(Error::Internal(format!(
666 "Found unexpected destination {fmt}",
667 ))),
668 }
669 }
670}
671
672impl ImageProcessorTrait for CPUProcessor {
673 fn convert(
674 &mut self,
675 src: &TensorDyn,
676 dst: &mut TensorDyn,
677 rotation: Rotation,
678 flip: Flip,
679 crop: Crop,
680 ) -> Result<()> {
681 let crop = crop.resolve(
682 src.width().unwrap_or(0),
683 src.height().unwrap_or(0),
684 dst.width().unwrap_or(0),
685 dst.height().unwrap_or(0),
686 )?;
687 self.convert_impl(src, dst, rotation, flip, crop)
688 }
689
690 fn draw_decoded_masks(
691 &mut self,
692 dst: &mut TensorDyn,
693 detect: &[DetectBox],
694 segmentation: &[Segmentation],
695 overlay: crate::MaskOverlay<'_>,
696 ) -> Result<()> {
697 prepare_dst_base_cpu(dst, overlay.background)?;
701 let dst = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
702 self.draw_decoded_masks_impl(
703 dst,
704 detect,
705 segmentation,
706 overlay.opacity,
707 overlay.color_mode,
708 )
709 }
710
711 fn draw_proto_masks(
712 &mut self,
713 dst: &mut TensorDyn,
714 detect: &[DetectBox],
715 proto_data: &ProtoData,
716 overlay: crate::MaskOverlay<'_>,
717 ) -> Result<()> {
718 prepare_dst_base_cpu(dst, overlay.background)?;
719 let dst = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
720 self.draw_proto_masks_impl(
721 dst,
722 detect,
723 proto_data,
724 overlay.opacity,
725 overlay.letterbox,
726 overlay.color_mode,
727 )
728 }
729
730 fn set_class_colors(&mut self, colors: &[[u8; 4]]) -> Result<()> {
731 for (c, new_c) in self.colors.iter_mut().zip(colors.iter()) {
732 *c = *new_c;
733 }
734 Ok(())
735 }
736}
737
738impl CPUProcessor {
740 pub(crate) fn convert_impl(
742 &mut self,
743 src: &TensorDyn,
744 dst: &mut TensorDyn,
745 rotation: Rotation,
746 flip: Flip,
747 crop: ResolvedCrop,
748 ) -> Result<()> {
749 let src_fmt = src.format().ok_or(Error::NotAnImage)?;
750 let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
751
752 let src_cm = crate::colorimetry::effective_colorimetry(src);
757 let dst_cm = crate::colorimetry::effective_colorimetry(dst);
758 let src_full = src_cm.range == Some(edgefirst_tensor::ColorRange::Full);
759 let dst_full = dst_cm.range == Some(edgefirst_tensor::ColorRange::Full);
760 let src_params = ColorParams {
761 matrix: crate::colorimetry::yuv_matrix(src_cm.encoding.unwrap()),
762 range: crate::colorimetry::yuv_range(src_cm.range.unwrap()),
763 encoding: src_cm.encoding.unwrap(),
764 range_kind: src_cm.range.unwrap(),
765 src_full_range: src_full,
766 dst_full_range: dst_full,
767 };
768 let dst_params = ColorParams {
769 matrix: crate::colorimetry::yuv_matrix(dst_cm.encoding.unwrap()),
770 range: crate::colorimetry::yuv_range(dst_cm.range.unwrap()),
771 encoding: dst_cm.encoding.unwrap(),
772 range_kind: dst_cm.range.unwrap(),
773 src_full_range: src_full,
774 dst_full_range: dst_full,
775 };
776 match (src.dtype(), dst.dtype()) {
777 (DType::U8, DType::U8) => {
778 let src = src.as_u8().unwrap();
779 let dst = dst.as_u8_mut().unwrap();
780 self.convert_u8(
781 src, dst, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
782 )
783 }
784 (DType::U8, DType::I8) => {
785 let src_u8 = src.as_u8().unwrap();
788 let dst_i8 = dst.as_i8_mut().unwrap();
789 let dst_u8 = unsafe { &mut *(dst_i8 as *mut Tensor<i8> as *mut Tensor<u8>) };
793 self.convert_u8(
794 src_u8, dst_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
795 )?;
796 let mut map = dst_u8.map_mut()?;
798 apply_int8_xor_bias(map.as_mut_slice(), dst_fmt);
799 Ok(())
800 }
801 (DType::U8, d @ (DType::F32 | DType::F16)) => {
802 let src_u8 = src.as_u8().unwrap();
803 let dw = dst.width().ok_or(Error::NotAnImage)?;
804 let dh = dst.height().ok_or(Error::NotAnImage)?;
805 let scratch_matches = self.widen_scratch.as_ref().is_some_and(|t| {
810 t.width() == Some(dw) && t.height() == Some(dh) && t.format() == Some(dst_fmt)
811 });
812 let mut tmp = if scratch_matches {
813 self.widen_scratch.take().unwrap()
814 } else {
815 TensorDyn::image(
816 dw,
817 dh,
818 dst_fmt,
819 DType::U8,
820 Some(TensorMemory::Mem),
821 edgefirst_tensor::CpuAccess::ReadWrite,
822 )?
823 };
824 {
825 let tmp_u8 = tmp.as_u8_mut().unwrap();
826 self.convert_u8(
827 src_u8, tmp_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params,
828 dst_params,
829 )?;
830 }
831 {
834 let tmp_u8 = tmp.as_u8().unwrap();
835 let src_map = tmp_u8.map_read()?;
836 match d {
837 DType::F32 => {
838 let dst_t = dst.as_f32_mut().unwrap();
839 let mut dst_map = dst_t.map_mut()?;
840 debug_assert_eq!(src_map.as_slice().len(), dst_map.as_slice().len());
841 simd::widen_u8_to_f32_norm(src_map.as_slice(), dst_map.as_mut_slice());
845 }
846 DType::F16 => {
847 let dst_t = dst.as_f16_mut().unwrap();
848 let mut dst_map = dst_t.map_mut()?;
849 debug_assert_eq!(src_map.as_slice().len(), dst_map.as_slice().len());
850 simd::widen_u8_to_f16_norm(src_map.as_slice(), dst_map.as_mut_slice());
855 }
856 _ => unreachable!(),
857 }
858 }
859 self.widen_scratch = Some(tmp);
860 Ok(())
861 }
862 (s, d) => Err(Error::NotSupported(format!("dtype {s} -> {d}",))),
863 }
864 }
865
866 fn reuse_or_alloc_image(
872 cached: Option<Tensor<u8>>,
873 w: usize,
874 h: usize,
875 fmt: PixelFormat,
876 ) -> Result<Tensor<u8>> {
877 if let Some(t) = cached {
878 if t.width() == Some(w) && t.height() == Some(h) && t.format() == Some(fmt) {
879 return Ok(t);
880 }
881 }
882 Ok(Tensor::<u8>::image(
883 w,
884 h,
885 fmt,
886 Some(TensorMemory::Mem),
887 edgefirst_tensor::CpuAccess::ReadWrite,
888 )?)
889 }
890
891 fn pre_resize_region(
912 &self,
913 src_fmt: PixelFormat,
914 (src_w, src_h): (usize, usize),
915 (dst_w, dst_h): (usize, usize),
916 rotation: Rotation,
917 crop: ResolvedCrop,
918 ) -> Option<Rect> {
919 use PixelFormat::{Nv12, Nv16, Nv24};
920
921 if !matches!(src_fmt, Nv12 | Nv16 | Nv24) {
922 return None;
923 }
924 let r = crop.src_rect?;
925 let full_src = Rect {
926 left: 0,
927 top: 0,
928 width: src_w,
929 height: src_h,
930 };
931 if r == full_src {
932 return None;
933 }
934
935 let d = crop.dst_rect.unwrap_or(Rect {
940 left: 0,
941 top: 0,
942 width: dst_w,
943 height: dst_h,
944 });
945 let (dst_x, dst_y) = match rotation {
946 Rotation::None | Rotation::Rotate180 => (d.width, d.height),
947 Rotation::Clockwise90 | Rotation::CounterClockwise90 => (d.height, d.width),
948 };
949 let halo_x = self.filter_halo(r.width, dst_x)?;
950 let halo_y = self.filter_halo(r.height, dst_y)?;
951
952 let mut left = r.left.saturating_sub(halo_x);
953 let mut top = r.top.saturating_sub(halo_y);
954 let mut right = (r.left + r.width + halo_x).min(src_w);
955 let mut bottom = (r.top + r.height + halo_y).min(src_h);
956
957 let (align_x, align_y) = match src_fmt {
963 Nv12 => (2, 2),
964 Nv16 => (2, 1),
965 _ => (1, 1),
966 };
967 left -= left % align_x;
968 top -= top % align_y;
969 right = right.next_multiple_of(align_x).min(src_w);
970 bottom = bottom.next_multiple_of(align_y).min(src_h);
971
972 let grown = Rect {
973 left,
974 top,
975 width: right - left,
976 height: bottom - top,
977 };
978 (grown != full_src).then_some(grown)
981 }
982
983 #[allow(clippy::too_many_arguments)]
985 fn convert_u8(
986 &mut self,
987 src: &Tensor<u8>,
988 dst: &mut Tensor<u8>,
989 src_fmt: PixelFormat,
990 dst_fmt: PixelFormat,
991 rotation: Rotation,
992 flip: Flip,
993 crop: ResolvedCrop,
994 src_params: ColorParams,
995 dst_params: ColorParams,
996 ) -> Result<()> {
997 use PixelFormat::*;
998
999 #[cfg(test)]
1000 {
1001 self.last_tmp_dims = None;
1002 }
1003
1004 let src_w = src.width().unwrap();
1005 let src_h = src.height().unwrap();
1006 let dst_w = dst.width().unwrap();
1007 let dst_h = dst.height().unwrap();
1008
1009 crop.check_crop_dims(src_w, src_h, dst_w, dst_h)?;
1010
1011 let intermediate = match (src_fmt, dst_fmt) {
1013 (Nv12, Rgb) => Rgb,
1014 (Nv12, Rgba) => Rgba,
1015 (Nv12, Grey) => Grey,
1016 (Nv12, Yuyv) => Rgba,
1017 (Nv12, Nv16) => Rgba,
1018 (Nv12, PlanarRgb) => Rgb,
1019 (Nv12, PlanarRgba) => Rgba,
1020 (Nv16, PlanarRgb) => Rgb,
1021 (Nv16, PlanarRgba) => Rgba,
1022 (Nv24, PlanarRgb) => Rgb,
1023 (Nv24, PlanarRgba) => Rgba,
1024 (Yuyv, Rgb) => Rgb,
1025 (Yuyv, Rgba) => Rgba,
1026 (Yuyv, Grey) => Grey,
1027 (Yuyv, Yuyv) => Rgba,
1028 (Yuyv, PlanarRgb) => Rgb,
1029 (Yuyv, PlanarRgba) => Rgba,
1030 (Yuyv, Nv16) => Rgba,
1031 (Vyuy, Rgb) => Rgb,
1032 (Vyuy, Rgba) => Rgba,
1033 (Vyuy, Grey) => Grey,
1034 (Vyuy, Vyuy) => Rgba,
1035 (Vyuy, PlanarRgb) => Rgb,
1036 (Vyuy, PlanarRgba) => Rgba,
1037 (Vyuy, Nv16) => Rgba,
1038 (Rgba, Rgb) => Rgba,
1039 (Rgba, Rgba) => Rgba,
1040 (Rgba, Grey) => Grey,
1041 (Rgba, Yuyv) => Rgba,
1042 (Rgba, PlanarRgb) => Rgba,
1043 (Rgba, PlanarRgba) => Rgba,
1044 (Rgba, Nv16) => Rgba,
1045 (Rgb, Rgb) => Rgb,
1046 (Rgb, Rgba) => Rgb,
1047 (Rgb, Grey) => Grey,
1048 (Rgb, Yuyv) => Rgb,
1049 (Rgb, PlanarRgb) => Rgb,
1050 (Rgb, PlanarRgba) => Rgb,
1051 (Rgb, Nv16) => Rgb,
1052 (Grey, Rgb) => Rgb,
1053 (Grey, Rgba) => Rgba,
1054 (Grey, Grey) => Grey,
1055 (Grey, Yuyv) => Grey,
1056 (Grey, PlanarRgb) => Grey,
1057 (Grey, PlanarRgba) => Grey,
1058 (Grey, Nv16) => Grey,
1059 (Nv12, Bgra) => Rgba,
1060 (Yuyv, Bgra) => Rgba,
1061 (Vyuy, Bgra) => Rgba,
1062 (Rgba, Bgra) => Rgba,
1063 (Rgb, Bgra) => Rgb,
1064 (Grey, Bgra) => Grey,
1065 (Bgra, Bgra) => Bgra,
1066 (Nv16, Rgb) => Rgb,
1067 (Nv16, Rgba) => Rgba,
1068 (Nv16, Bgra) => Rgba,
1069 (Nv24, Rgb) => Rgb,
1070 (Nv24, Rgba) => Rgba,
1071 (Nv24, Grey) => Grey,
1072 (Nv24, Bgra) => Rgba,
1073 (PlanarRgb, Rgb) => Rgb,
1074 (PlanarRgb, Rgba) => Rgb,
1075 (PlanarRgb, Bgra) => Rgb,
1076 (PlanarRgba, Rgb) => Rgba,
1077 (PlanarRgba, Rgba) => Rgba,
1078 (PlanarRgba, Bgra) => Rgba,
1079 (s, d) => {
1080 return Err(Error::NotSupported(format!("Conversion from {s} to {d}",)));
1081 }
1082 };
1083
1084 let need_resize_flip_rotation = rotation != Rotation::None
1085 || flip != Flip::None
1086 || src_w != dst_w
1087 || src_h != dst_h
1088 || crop.src_rect.is_some_and(|c| {
1089 c != Rect {
1090 left: 0,
1091 top: 0,
1092 width: src_w,
1093 height: src_h,
1094 }
1095 })
1096 || crop.dst_rect.is_some_and(|c| {
1097 c != Rect {
1098 left: 0,
1099 top: 0,
1100 width: dst_w,
1101 height: dst_h,
1102 }
1103 });
1104
1105 let direct_is_yuv_src = matches!(src_fmt, Nv12 | Nv16 | Nv24 | Yuyv | Vyuy);
1108 let direct_params = if direct_is_yuv_src {
1109 src_params
1110 } else {
1111 dst_params
1112 };
1113
1114 let full_dst_rect = Rect {
1131 left: 0,
1132 top: 0,
1133 width: dst_w,
1134 height: dst_h,
1135 };
1136 let fused_region = if rotation != Rotation::None || flip != Flip::None {
1137 None
1138 } else {
1139 match crop.src_rect {
1140 None if src_w == dst_w && src_h == dst_h => Some(None),
1141 None => None,
1142 Some(r)
1143 if r.width == dst_w
1144 && r.height == dst_h
1145 && crop.dst_rect.is_none_or(|d| d == full_dst_rect)
1146 && chroma_alignment_ok(src_fmt, r) =>
1147 {
1148 Some(Some(r))
1149 }
1150 Some(_) => None,
1151 }
1152 };
1153 if let Some(region) = fused_region {
1154 if matches!(src_fmt, Nv12 | Nv16 | Nv24) && matches!(dst_fmt, PlanarRgb | PlanarRgba) {
1155 #[cfg(test)]
1156 {
1157 self.fused_hits += 1;
1158 }
1159 return self.convert_nv_to_planar_fused(
1160 src,
1161 dst,
1162 src_fmt,
1163 dst_fmt,
1164 direct_params,
1165 region,
1166 );
1167 }
1168 }
1169
1170 if !need_resize_flip_rotation && Self::support_conversion_pf(src_fmt, dst_fmt) {
1172 return Self::convert_format_pf(src, dst, src_fmt, dst_fmt, direct_params);
1173 }
1174
1175 if dst_fmt == Yuyv && !dst_w.is_multiple_of(2) {
1177 return Err(Error::NotSupported(format!(
1178 "{} destination must have width divisible by 2",
1179 dst_fmt,
1180 )));
1181 }
1182
1183 let mut cached_tmp = self.convert_tmp.take();
1191 let mut cached_tmp2 = self.convert_tmp2.take();
1192
1193 let pre_region = if intermediate != src_fmt {
1198 self.pre_resize_region(src_fmt, (src_w, src_h), (dst_w, dst_h), rotation, crop)
1199 } else {
1200 None
1201 };
1202
1203 let tmp_holder: Option<Tensor<u8>> = if intermediate != src_fmt {
1205 let _s = tracing::trace_span!(
1206 "image.convert.cpu.format_convert",
1207 from = ?src_fmt,
1208 to = ?intermediate,
1209 pass = "pre_resize",
1210 )
1211 .entered();
1212 let (tmp_w, tmp_h) = pre_region.map_or((src_w, src_h), |g| (g.width, g.height));
1213 let mut t = Self::reuse_or_alloc_image(cached_tmp.take(), tmp_w, tmp_h, intermediate)?;
1214 #[cfg(test)]
1215 {
1216 self.last_tmp_dims = Some((tmp_w, tmp_h));
1217 }
1218 match pre_region {
1219 Some(g) => {
1220 let mut sub = Self::reuse_or_alloc_image(
1221 self.convert_src_sub.take(),
1222 g.width,
1223 g.height,
1224 src_fmt,
1225 )?;
1226 {
1227 let _s = tracing::trace_span!(
1228 "image.convert.cpu.extract_region",
1229 region_w = g.width,
1230 region_h = g.height,
1231 )
1232 .entered();
1233 Self::extract_nv_region(src, &mut sub, src_fmt, g)?;
1234 }
1235 Self::convert_format_pf(&sub, &mut t, src_fmt, intermediate, src_params)?;
1236 self.convert_src_sub = Some(sub);
1237 }
1238 None => Self::convert_format_pf(src, &mut t, src_fmt, intermediate, src_params)?,
1239 }
1240 Some(t)
1241 } else {
1242 None
1243 };
1244
1245 let crop = match (pre_region, crop.src_rect) {
1250 (Some(g), Some(r)) => ResolvedCrop {
1251 src_rect: Some(Rect {
1252 left: r.left - g.left,
1253 top: r.top - g.top,
1254 ..r
1255 }),
1256 ..crop
1257 },
1258 _ => crop,
1259 };
1260 let (tmp, tmp_fmt): (&Tensor<u8>, PixelFormat) = match &tmp_holder {
1261 Some(t) => (t, intermediate),
1262 None => (src, src_fmt),
1263 };
1264
1265 debug_assert!(matches!(tmp_fmt, Rgb | Rgba | Grey));
1267 if tmp_fmt == dst_fmt {
1268 let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
1269 self.resize_flip_rotate_pf(tmp, dst, dst_fmt, rotation, flip, crop)?;
1270 } else if !need_resize_flip_rotation {
1271 let _s = tracing::trace_span!(
1272 "image.convert.cpu.format_convert",
1273 from = ?tmp_fmt,
1274 to = ?dst_fmt,
1275 pass = "direct",
1276 )
1277 .entered();
1278 Self::convert_format_pf(tmp, dst, tmp_fmt, dst_fmt, dst_params)?;
1279 } else {
1280 let mut tmp2 = Self::reuse_or_alloc_image(cached_tmp2.take(), dst_w, dst_h, tmp_fmt)?;
1281 if crop.dst_rect.is_some_and(|c| {
1282 c != Rect {
1283 left: 0,
1284 top: 0,
1285 width: dst_w,
1286 height: dst_h,
1287 }
1288 }) && crop.dst_color.is_none()
1289 {
1290 Self::convert_format_pf(dst, &mut tmp2, dst_fmt, tmp_fmt, dst_params)?;
1291 }
1292 {
1293 let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
1294 self.resize_flip_rotate_pf(tmp, &mut tmp2, tmp_fmt, rotation, flip, crop)?;
1295 }
1296 {
1297 let _s = tracing::trace_span!(
1298 "image.convert.cpu.format_convert",
1299 from = ?tmp_fmt,
1300 to = ?dst_fmt,
1301 pass = "post_resize",
1302 )
1303 .entered();
1304 Self::convert_format_pf(&tmp2, dst, tmp_fmt, dst_fmt, dst_params)?;
1305 }
1306 cached_tmp2 = Some(tmp2);
1307 }
1308 if let Some(t) = tmp_holder {
1311 cached_tmp = Some(t);
1312 }
1313 self.convert_tmp = cached_tmp;
1314 self.convert_tmp2 = cached_tmp2;
1315
1316 if let (Some(dst_rect), Some(dst_color)) = (crop.dst_rect, crop.dst_color) {
1317 let full_rect = Rect {
1318 left: 0,
1319 top: 0,
1320 width: dst_w,
1321 height: dst_h,
1322 };
1323 if dst_rect != full_rect {
1324 Self::fill_image_outside_crop_u8(dst, dst_color, dst_rect)?;
1325 }
1326 }
1327
1328 Ok(())
1329 }
1330
1331 fn draw_decoded_masks_impl(
1332 &mut self,
1333 dst: &mut Tensor<u8>,
1334 detect: &[DetectBox],
1335 segmentation: &[Segmentation],
1336 opacity: f32,
1337 color_mode: crate::ColorMode,
1338 ) -> Result<()> {
1339 let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
1340 if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
1341 return Err(crate::Error::NotSupported(
1342 "CPU image rendering only supports RGBA or RGB images".to_string(),
1343 ));
1344 }
1345
1346 let _timer = FunctionTimer::new("CPUProcessor::draw_decoded_masks");
1347
1348 let dst_w = dst.width().unwrap();
1349 let dst_h = dst.height().unwrap();
1350 let dst_rs = tensor_row_stride(dst);
1351 let dst_c = dst_fmt.channels();
1352
1353 let mut map = dst.map_mut()?;
1354 let dst_slice = map.as_mut_slice();
1355
1356 self.render_box(dst_w, dst_h, dst_rs, dst_c, dst_slice, detect, color_mode)?;
1357
1358 if segmentation.is_empty() {
1359 return Ok(());
1360 }
1361
1362 let is_semantic = segmentation[0].segmentation.shape()[2] > 1;
1365
1366 if is_semantic {
1367 self.render_modelpack_segmentation(
1368 dst_w,
1369 dst_h,
1370 dst_rs,
1371 dst_c,
1372 dst_slice,
1373 &segmentation[0],
1374 opacity,
1375 )?;
1376 } else {
1377 for (idx, (seg, det)) in segmentation.iter().zip(detect).enumerate() {
1378 let color_index = color_mode.index(idx, det.label);
1379 self.render_yolo_segmentation(
1380 dst_w,
1381 dst_h,
1382 dst_rs,
1383 dst_c,
1384 dst_slice,
1385 seg,
1386 color_index,
1387 opacity,
1388 )?;
1389 }
1390 }
1391
1392 Ok(())
1393 }
1394
1395 fn draw_proto_masks_impl(
1396 &mut self,
1397 dst: &mut Tensor<u8>,
1398 detect: &[DetectBox],
1399 proto_data: &ProtoData,
1400 opacity: f32,
1401 letterbox: Option<[f32; 4]>,
1402 color_mode: crate::ColorMode,
1403 ) -> Result<()> {
1404 let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
1405 if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
1406 return Err(crate::Error::NotSupported(
1407 "CPU image rendering only supports RGBA or RGB images".to_string(),
1408 ));
1409 }
1410
1411 let _timer = FunctionTimer::new("CPUProcessor::draw_proto_masks");
1412
1413 let dst_w = dst.width().unwrap();
1414 let dst_h = dst.height().unwrap();
1415 let dst_rs = tensor_row_stride(dst);
1416 let channels = dst_fmt.channels();
1417
1418 let mut map = dst.map_mut()?;
1419 let dst_slice = map.as_mut_slice();
1420
1421 self.render_box(
1422 dst_w, dst_h, dst_rs, channels, dst_slice, detect, color_mode,
1423 )?;
1424
1425 if detect.is_empty() {
1426 return Ok(());
1427 }
1428 let proto_shape = proto_data.protos.shape();
1429 if proto_shape.len() != 3 {
1430 return Err(Error::InvalidShape(format!(
1431 "protos tensor must be rank-3, got {proto_shape:?}"
1432 )));
1433 }
1434 let proto_h = proto_shape[0];
1435 let proto_w = proto_shape[1];
1436 let num_protos = proto_shape[2];
1437 let coeff_shape = proto_data.mask_coefficients.shape();
1438 if coeff_shape.len() != 2 {
1439 return Err(Error::InvalidShape(format!(
1440 "mask_coefficients tensor must be rank-2, got {coeff_shape:?}"
1441 )));
1442 }
1443 if coeff_shape[0] == 0 {
1445 return Ok(());
1446 }
1447 if coeff_shape[1] != num_protos {
1448 return Err(Error::InvalidShape(format!(
1449 "mask_coefficients second dimension must match num_protos \
1450 ({num_protos}), got {coeff_shape:?}"
1451 )));
1452 }
1453
1454 let coeff_f32: Vec<f32> = match proto_data.mask_coefficients.dtype() {
1456 DType::F32 => {
1457 let t = proto_data.mask_coefficients.as_f32().expect("F32");
1458 let m = t.map_read()?;
1459 m.as_slice().to_vec()
1460 }
1461 DType::F16 => {
1462 let t = proto_data.mask_coefficients.as_f16().expect("F16");
1463 let m = t.map_read()?;
1464 m.as_slice().iter().map(|v| v.to_f32()).collect()
1465 }
1466 DType::I8 => {
1467 let t = proto_data.mask_coefficients.as_i8().expect("I8");
1468 let m = t.map_read()?;
1469 if let Some(q) = t.quantization() {
1470 use edgefirst_tensor::QuantMode;
1471 let (scale, zp) = match q.mode() {
1472 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1473 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1474 other => {
1475 return Err(Error::NotSupported(format!(
1476 "I8 mask_coefficients quantization mode {other:?} not supported"
1477 )));
1478 }
1479 };
1480 m.as_slice()
1481 .iter()
1482 .map(|&v| (v as f32 - zp) * scale)
1483 .collect()
1484 } else {
1485 m.as_slice().iter().map(|&v| v as f32).collect()
1486 }
1487 }
1488 DType::I16 => {
1489 let t = proto_data.mask_coefficients.as_i16().expect("I16");
1490 let m = t.map_read()?;
1491 if let Some(q) = t.quantization() {
1492 use edgefirst_tensor::QuantMode;
1493 let (scale, zp) = match q.mode() {
1494 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1495 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1496 other => {
1497 return Err(Error::NotSupported(format!(
1498 "I16 mask_coefficients quantization mode {other:?} not supported"
1499 )));
1500 }
1501 };
1502 m.as_slice()
1503 .iter()
1504 .map(|&v| (v as f32 - zp) * scale)
1505 .collect()
1506 } else {
1507 m.as_slice().iter().map(|&v| v as f32).collect()
1508 }
1509 }
1510 other => {
1511 return Err(Error::InvalidShape(format!(
1512 "mask_coefficients dtype {other:?} not supported"
1513 )));
1514 }
1515 };
1516
1517 let (lx0, lx_range, ly0, ly_range) = match letterbox {
1519 Some([lx0, ly0, lx1, ly1]) => (lx0, lx1 - lx0, ly0, ly1 - ly0),
1520 None => (0.0_f32, 1.0_f32, 0.0_f32, 1.0_f32),
1521 };
1522
1523 match proto_data.protos.dtype() {
1526 DType::F32 => {
1527 let t = proto_data.protos.as_f32().expect("F32");
1528 let m = t.map_read()?;
1529 self.draw_proto_masks_inner(
1530 dst_slice,
1531 dst_w,
1532 dst_h,
1533 dst_rs,
1534 channels,
1535 detect,
1536 m.as_slice(),
1537 &coeff_f32,
1538 proto_h,
1539 proto_w,
1540 num_protos,
1541 opacity,
1542 (lx0, lx_range, ly0, ly_range),
1543 color_mode,
1544 0.0_f32,
1545 |p: &f32, _| *p,
1546 );
1547 }
1548 DType::F16 => {
1549 let t = proto_data.protos.as_f16().expect("F16");
1550 let m = t.map_read()?;
1551 self.draw_proto_masks_inner(
1552 dst_slice,
1553 dst_w,
1554 dst_h,
1555 dst_rs,
1556 channels,
1557 detect,
1558 m.as_slice(),
1559 &coeff_f32,
1560 proto_h,
1561 proto_w,
1562 num_protos,
1563 opacity,
1564 (lx0, lx_range, ly0, ly_range),
1565 color_mode,
1566 0.0_f32,
1567 |p: &half::f16, _| p.to_f32(),
1568 );
1569 }
1570 DType::I8 => {
1571 use edgefirst_tensor::QuantMode;
1572 let t = proto_data.protos.as_i8().expect("I8");
1573 let m = t.map_read()?;
1574 let quant = t.quantization().ok_or_else(|| {
1575 Error::InvalidShape("I8 protos require quantization metadata".into())
1576 })?;
1577 let (scale, zp) = match quant.mode() {
1578 QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
1579 QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
1580 QuantMode::PerChannel { axis, .. }
1581 | QuantMode::PerChannelSymmetric { axis, .. } => {
1582 return Err(Error::NotSupported(format!(
1583 "per-channel quantization (axis={axis}) in draw_proto_masks \
1584 CPU path not yet supported"
1585 )));
1586 }
1587 };
1588 self.draw_proto_masks_inner(
1589 dst_slice,
1590 dst_w,
1591 dst_h,
1592 dst_rs,
1593 channels,
1594 detect,
1595 m.as_slice(),
1596 &coeff_f32,
1597 proto_h,
1598 proto_w,
1599 num_protos,
1600 opacity,
1601 (lx0, lx_range, ly0, ly_range),
1602 color_mode,
1603 scale,
1604 move |p: &i8, _| (*p as f32) - zp,
1605 );
1606 }
1607 other => {
1608 return Err(Error::InvalidShape(format!(
1609 "proto tensor dtype {other:?} not supported"
1610 )));
1611 }
1612 }
1613
1614 Ok(())
1615 }
1616
1617 #[allow(clippy::too_many_arguments)]
1618 fn draw_proto_masks_inner<P: Copy>(
1619 &self,
1620 dst_slice: &mut [u8],
1621 dst_w: usize,
1622 dst_h: usize,
1623 dst_rs: usize,
1624 channels: usize,
1625 detect: &[DetectBox],
1626 protos: &[P],
1627 coeff_all_f32: &[f32],
1628 proto_h: usize,
1629 proto_w: usize,
1630 num_protos: usize,
1631 opacity: f32,
1632 letterbox_xy: (f32, f32, f32, f32),
1633 color_mode: crate::ColorMode,
1634 acc_scale: f32,
1635 load_f32: impl Fn(&P, f32) -> f32 + Copy,
1636 ) {
1637 let (lx0, lx_range, ly0, ly_range) = letterbox_xy;
1638 let stride_y = proto_w * num_protos;
1639 for (idx, det) in detect.iter().enumerate() {
1640 let coeff = &coeff_all_f32[idx * num_protos..(idx + 1) * num_protos];
1641 let color_index = color_mode.index(idx, det.label);
1642 let color = self.colors[color_index % self.colors.len()];
1643 let alpha = if opacity == 1.0 {
1644 color[3] as u16
1645 } else {
1646 (color[3] as f32 * opacity).round() as u16
1647 };
1648
1649 let start_x = (dst_w as f32 * det.bbox.xmin).round() as usize;
1650 let start_y = (dst_h as f32 * det.bbox.ymin).round() as usize;
1651 let end_x = ((dst_w as f32 * det.bbox.xmax).round() as usize).min(dst_w);
1652 let end_y = ((dst_h as f32 * det.bbox.ymax).round() as usize).min(dst_h);
1653
1654 for y in start_y..end_y {
1655 for x in start_x..end_x {
1656 let px = (lx0 + (x as f32 / dst_w as f32) * lx_range) * proto_w as f32 - 0.5;
1657 let py = (ly0 + (y as f32 / dst_h as f32) * ly_range) * proto_h as f32 - 0.5;
1658
1659 let x0 = (px.floor() as isize).clamp(0, proto_w as isize - 1) as usize;
1663 let y0 = (py.floor() as isize).clamp(0, proto_h as isize - 1) as usize;
1664 let x1 = (x0 + 1).min(proto_w - 1);
1665 let y1 = (y0 + 1).min(proto_h - 1);
1666 let fx = px - px.floor();
1667 let fy = py - py.floor();
1668 let w00 = (1.0 - fx) * (1.0 - fy);
1669 let w10 = fx * (1.0 - fy);
1670 let w01 = (1.0 - fx) * fy;
1671 let w11 = fx * fy;
1672 let b00 = y0 * stride_y + x0 * num_protos;
1673 let b10 = y0 * stride_y + x1 * num_protos;
1674 let b01 = y1 * stride_y + x0 * num_protos;
1675 let b11 = y1 * stride_y + x1 * num_protos;
1676 let mut acc = 0.0_f32;
1677 for p in 0..num_protos {
1678 let v00 = load_f32(&protos[b00 + p], 0.0);
1679 let v10 = load_f32(&protos[b10 + p], 0.0);
1680 let v01 = load_f32(&protos[b01 + p], 0.0);
1681 let v11 = load_f32(&protos[b11 + p], 0.0);
1682 let val = w00 * v00 + w10 * v10 + w01 * v01 + w11 * v11;
1683 acc += coeff[p] * val;
1684 }
1685 let final_acc = if acc_scale == 0.0 {
1686 acc
1687 } else {
1688 acc_scale * acc
1689 };
1690 let mask = 1.0 / (1.0 + (-final_acc).exp());
1694 if mask < 0.5 {
1695 continue;
1696 }
1697 let dst_index = y * dst_rs + x * channels;
1698 for c in 0..3 {
1699 dst_slice[dst_index + c] = ((color[c] as u16 * alpha
1700 + dst_slice[dst_index + c] as u16 * (255 - alpha))
1701 / 255) as u8;
1702 }
1703 }
1704 }
1705 }
1706 }
1707}