use crate::{
Crop, Error, Flip, FunctionTimer, ImageProcessorTrait, Rect, ResolvedCrop, Result, Rotation,
};
use edgefirst_decoder::{DetectBox, ProtoData, Segmentation};
use edgefirst_tensor::{
DType, PixelFormat, Tensor, TensorDyn, TensorMapTrait, TensorMemory, TensorTrait,
};
mod convert;
mod masks;
mod resize;
mod simd;
mod tests;
#[derive(Debug, Clone, Copy)]
pub(crate) struct ColorParams {
pub matrix: yuv::YuvStandardMatrix,
pub range: yuv::YuvRange,
pub encoding: edgefirst_tensor::ColorEncoding,
pub range_kind: edgefirst_tensor::ColorRange,
pub src_full_range: bool,
pub dst_full_range: bool,
}
#[derive(Debug)]
pub struct CPUProcessor {
resizer: fast_image_resize::Resizer,
options: fast_image_resize::ResizeOptions,
colors: [[u8; 4]; 20],
widen_scratch: Option<TensorDyn>,
resize_destride_scratch: Vec<u8>,
resize_dst_destride_scratch: Vec<u8>,
nv_strip_scratch: Vec<u8>,
nv_strip_y_pack: Vec<u8>,
nv_strip_uv_pack: Vec<u8>,
#[cfg(test)]
fused_hits: u64,
#[cfg(test)]
last_tmp_dims: Option<(usize, usize)>,
convert_tmp: Option<Tensor<u8>>,
convert_tmp2: Option<Tensor<u8>>,
convert_src_sub: Option<Tensor<u8>>,
}
impl Clone for CPUProcessor {
fn clone(&self) -> Self {
Self {
resizer: self.resizer.clone(),
options: self.options,
colors: self.colors,
widen_scratch: None,
resize_destride_scratch: Vec::new(),
resize_dst_destride_scratch: Vec::new(),
nv_strip_scratch: Vec::new(),
nv_strip_y_pack: Vec::new(),
nv_strip_uv_pack: Vec::new(),
#[cfg(test)]
fused_hits: 0,
#[cfg(test)]
last_tmp_dims: None,
convert_tmp: None,
convert_tmp2: None,
convert_src_sub: None,
}
}
}
unsafe impl Send for CPUProcessor {}
unsafe impl Sync for CPUProcessor {}
impl Default for CPUProcessor {
fn default() -> Self {
Self::new_bilinear()
}
}
fn prepare_dst_base_cpu(dst: &mut TensorDyn, background: Option<&TensorDyn>) -> Result<()> {
match background {
Some(bg) => {
if bg.shape() != dst.shape() {
return Err(Error::InvalidShape(
"background shape does not match dst".into(),
));
}
if bg.format() != dst.format() {
return Err(Error::InvalidShape(
"background pixel format does not match dst".into(),
));
}
let bg_u8 = bg.as_u8().ok_or(Error::NotAnImage)?;
let dst_u8 = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
let bg_map = bg_u8.map_read()?;
let mut dst_map = dst_u8.map_mut()?;
let bg_slice = bg_map.as_slice();
let dst_slice = dst_map.as_mut_slice();
if bg_slice.len() != dst_slice.len() {
return Err(Error::InvalidShape(
"background buffer size does not match dst".into(),
));
}
dst_slice.copy_from_slice(bg_slice);
}
None => {
let dst_u8 = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
let mut dst_map = dst_u8.map_mut()?;
dst_map.as_mut_slice().fill(0);
}
}
Ok(())
}
fn chroma_alignment_ok(fmt: PixelFormat, r: Rect) -> bool {
match fmt {
PixelFormat::Nv12 => r.left.is_multiple_of(2) && r.top.is_multiple_of(2),
PixelFormat::Nv16 => r.left.is_multiple_of(2),
PixelFormat::Nv24 => true,
_ => false,
}
}
fn row_stride_for(width: usize, fmt: PixelFormat) -> usize {
use edgefirst_tensor::PixelLayout;
match fmt.layout() {
PixelLayout::Packed => width * fmt.channels(),
PixelLayout::Planar | PixelLayout::SemiPlanar => width,
_ => width, }
}
fn tensor_row_stride(tensor: &Tensor<u8>) -> usize {
tensor.effective_row_stride().unwrap_or_else(|| {
let w = tensor.width().unwrap_or(0);
let fmt = tensor.format().unwrap_or(PixelFormat::Rgb);
row_stride_for(w, fmt)
})
}
fn logical_surface(tensor: &Tensor<u8>) -> Result<(usize, usize)> {
use edgefirst_tensor::PixelLayout;
let fmt = tensor.format().ok_or(Error::NotAnImage)?;
let w = tensor.width().ok_or(Error::NotAnImage)?;
let h = tensor.height().ok_or(Error::NotAnImage)?;
Ok(match fmt.layout() {
PixelLayout::Packed => (h, w * fmt.channels()),
PixelLayout::Planar => (fmt.channels() * h, w),
PixelLayout::SemiPlanar => (fmt.combined_plane_height(h).unwrap_or(h), w),
_ => (h, row_stride_for(w, fmt)),
})
}
fn packed_row_pairs<'s, 'd>(
src: &'s [u8],
src_stride: usize,
src_row_bytes: usize,
dst: &'d mut [u8],
dst_stride: usize,
dst_row_bytes: usize,
rows: usize,
) -> impl Iterator<Item = (&'s [u8], &'d mut [u8])> {
src.chunks(src_stride)
.zip(dst.chunks_mut(dst_stride))
.take(rows)
.map(move |(s, d)| (&s[..src_row_bytes], &mut d[..dst_row_bytes]))
}
fn split_semi_planar(
bytes: &[u8],
stride: usize,
src_h: usize,
fmt: PixelFormat,
) -> Result<(&[u8], &[u8])> {
let total_h = fmt.combined_plane_height(src_h).unwrap_or(src_h);
let need = stride.checked_mul(total_h).ok_or_else(|| {
Error::InvalidShape(format!(
"{fmt:?} plane size overflow (stride={stride}, h={src_h})"
))
})?;
if bytes.len() < need {
return Err(Error::InvalidShape(format!(
"{fmt:?} source has {} bytes but needs {need} (stride={stride}, h={src_h})",
bytes.len()
)));
}
Ok(bytes.split_at(stride * src_h))
}
fn split_semi_planar_mut(
bytes: &mut [u8],
stride: usize,
dst_h: usize,
fmt: PixelFormat,
) -> Result<(&mut [u8], &mut [u8])> {
let total_h = fmt.combined_plane_height(dst_h).unwrap_or(dst_h);
let need = stride.checked_mul(total_h).ok_or_else(|| {
Error::InvalidShape(format!(
"{fmt:?} plane size overflow (stride={stride}, combined_h={total_h})"
))
})?;
if bytes.len() < need {
return Err(Error::InvalidShape(format!(
"{fmt:?} destination has {} bytes but needs {need} (stride={stride}, combined_h={total_h})",
bytes.len()
)));
}
Ok(bytes.split_at_mut(stride * dst_h))
}
fn guard_plane(
buf_len: usize,
stride: usize,
rows: usize,
row_bytes: usize,
what: &str,
) -> Result<()> {
let need = stride.checked_mul(rows).ok_or_else(|| {
Error::InvalidShape(format!(
"{what} plane size overflow (stride={stride}, rows={rows})"
))
})?;
if row_bytes > stride || buf_len < need {
return Err(Error::InvalidShape(format!(
"{what} buffer too small: {buf_len} bytes, need {need} (stride={stride}, rows={rows}, row_bytes={row_bytes})"
)));
}
Ok(())
}
pub(crate) fn apply_int8_xor_bias(data: &mut [u8], fmt: PixelFormat) {
use edgefirst_tensor::PixelLayout;
if !fmt.has_alpha() {
for b in data.iter_mut() {
*b ^= 0x80;
}
} else if fmt.layout() == PixelLayout::Planar {
let channels = fmt.channels();
let plane_size = data.len() / channels;
for b in data[..plane_size * (channels - 1)].iter_mut() {
*b ^= 0x80;
}
} else {
let channels = fmt.channels();
for pixel in data.chunks_exact_mut(channels) {
for b in &mut pixel[..channels - 1] {
*b ^= 0x80;
}
}
}
}
fn apply_int8_xor_bias_rows(tensor: &mut Tensor<u8>, fmt: PixelFormat) -> Result<()> {
use edgefirst_tensor::PixelLayout;
let (rows, row_bytes) = logical_surface(tensor)?;
let stride = tensor_row_stride(tensor);
let color_rows = if fmt.has_alpha() && fmt.layout() == PixelLayout::Planar {
rows / fmt.channels() * (fmt.channels() - 1)
} else {
rows
};
let row_fmt = if fmt.layout() == PixelLayout::Planar {
PixelFormat::Grey
} else {
fmt
};
let mut map = tensor.map_mut()?;
let buf = map.as_mut_slice();
guard_plane(buf.len(), stride, rows, row_bytes, "int8 bias dst")?;
for row in buf.chunks_mut(stride).take(color_rows) {
apply_int8_xor_bias(&mut row[..row_bytes], row_fmt);
}
Ok(())
}
impl CPUProcessor {
pub fn new() -> Self {
Self::new_bilinear()
}
fn new_bilinear() -> Self {
let resizer = fast_image_resize::Resizer::new();
let options = fast_image_resize::ResizeOptions::new()
.resize_alg(fast_image_resize::ResizeAlg::Convolution(
fast_image_resize::FilterType::Bilinear,
))
.use_alpha(false);
log::debug!("CPUConverter created");
Self {
resizer,
options,
colors: crate::DEFAULT_COLORS_U8,
widen_scratch: None,
resize_destride_scratch: Vec::new(),
resize_dst_destride_scratch: Vec::new(),
nv_strip_scratch: Vec::new(),
nv_strip_y_pack: Vec::new(),
nv_strip_uv_pack: Vec::new(),
#[cfg(test)]
fused_hits: 0,
#[cfg(test)]
last_tmp_dims: None,
convert_tmp: None,
convert_tmp2: None,
convert_src_sub: None,
}
}
pub fn new_nearest() -> Self {
let resizer = fast_image_resize::Resizer::new();
let options = fast_image_resize::ResizeOptions::new()
.resize_alg(fast_image_resize::ResizeAlg::Nearest)
.use_alpha(false);
log::debug!("CPUConverter created");
Self {
resizer,
options,
colors: crate::DEFAULT_COLORS_U8,
widen_scratch: None,
resize_destride_scratch: Vec::new(),
resize_dst_destride_scratch: Vec::new(),
nv_strip_scratch: Vec::new(),
nv_strip_y_pack: Vec::new(),
nv_strip_uv_pack: Vec::new(),
#[cfg(test)]
fused_hits: 0,
#[cfg(test)]
last_tmp_dims: None,
convert_tmp: None,
convert_tmp2: None,
convert_src_sub: None,
}
}
#[cfg(test)]
pub(super) fn fused_hits(&self) -> u64 {
self.fused_hits
}
#[cfg(test)]
pub(super) fn last_tmp_dims(&self) -> Option<(usize, usize)> {
self.last_tmp_dims
}
pub(crate) fn support_conversion_pf(src: PixelFormat, dst: PixelFormat) -> bool {
use PixelFormat::*;
matches!(
(src, dst),
(Nv12, Rgb)
| (Nv12, Rgba)
| (Nv12, Grey)
| (Nv16, Rgb)
| (Nv16, Rgba)
| (Nv16, Bgra)
| (Nv24, Rgb)
| (Nv24, Rgba)
| (Nv24, Grey)
| (Nv24, Bgra)
| (Yuyv, Rgb)
| (Yuyv, Rgba)
| (Yuyv, Grey)
| (Yuyv, Yuyv)
| (Yuyv, PlanarRgb)
| (Yuyv, PlanarRgba)
| (Yuyv, Nv16)
| (Vyuy, Rgb)
| (Vyuy, Rgba)
| (Vyuy, Grey)
| (Vyuy, Vyuy)
| (Vyuy, PlanarRgb)
| (Vyuy, PlanarRgba)
| (Vyuy, Nv16)
| (Rgba, Rgb)
| (Rgba, Rgba)
| (Rgba, Grey)
| (Rgba, Yuyv)
| (Rgba, PlanarRgb)
| (Rgba, PlanarRgba)
| (Rgba, Nv16)
| (Rgb, Rgb)
| (Rgb, Rgba)
| (Rgb, Grey)
| (Rgb, Yuyv)
| (Rgb, PlanarRgb)
| (Rgb, PlanarRgba)
| (Rgb, Nv16)
| (Grey, Rgb)
| (Grey, Rgba)
| (Grey, Grey)
| (Grey, Yuyv)
| (Grey, PlanarRgb)
| (Grey, PlanarRgba)
| (Grey, Nv16)
| (Nv12, Bgra)
| (Yuyv, Bgra)
| (Vyuy, Bgra)
| (Rgba, Bgra)
| (Rgb, Bgra)
| (Grey, Bgra)
| (Bgra, Bgra)
| (PlanarRgb, Rgb)
| (PlanarRgb, Rgba)
| (PlanarRgba, Rgb)
| (PlanarRgba, Rgba)
| (PlanarRgb, Bgra)
| (PlanarRgba, Bgra)
)
}
pub(crate) fn convert_format_pf(
src: &Tensor<u8>,
dst: &mut Tensor<u8>,
src_fmt: PixelFormat,
dst_fmt: PixelFormat,
cp: ColorParams,
) -> Result<()> {
let _timer = FunctionTimer::new(format!(
"ImageProcessor::convert_format {} to {}",
src_fmt, dst_fmt,
));
use PixelFormat::*;
match (src_fmt, dst_fmt) {
(Nv12, Rgb) => Self::convert_nv12_to_rgb(src, dst, cp),
(Nv12, Rgba) => Self::convert_nv12_to_rgba(src, dst, cp),
(Nv12, Grey) => Self::convert_nv12_to_grey(src, dst, cp),
(Yuyv, Rgb) => Self::convert_yuyv_to_rgb(src, dst, cp),
(Yuyv, Rgba) => Self::convert_yuyv_to_rgba(src, dst, cp),
(Yuyv, Grey) => Self::convert_yuyv_to_grey(src, dst, cp),
(Yuyv, Yuyv) => Self::copy_image(src, dst),
(Yuyv, PlanarRgb) => Self::convert_yuyv_to_8bps(src, dst, cp),
(Yuyv, PlanarRgba) => Self::convert_yuyv_to_prgba(src, dst, cp),
(Yuyv, Nv16) => Self::convert_yuyv_to_nv16(src, dst),
(Vyuy, Rgb) => Self::convert_vyuy_to_rgb(src, dst, cp),
(Vyuy, Rgba) => Self::convert_vyuy_to_rgba(src, dst, cp),
(Vyuy, Grey) => Self::convert_vyuy_to_grey(src, dst, cp),
(Vyuy, Vyuy) => Self::copy_image(src, dst),
(Vyuy, PlanarRgb) => Self::convert_vyuy_to_8bps(src, dst, cp),
(Vyuy, PlanarRgba) => Self::convert_vyuy_to_prgba(src, dst, cp),
(Vyuy, Nv16) => Self::convert_vyuy_to_nv16(src, dst),
(Rgba, Rgb) => Self::convert_rgba_to_rgb(src, dst),
(Rgba, Rgba) => Self::copy_image(src, dst),
(Rgba, Grey) => Self::convert_rgba_to_grey(src, dst),
(Rgba, Yuyv) => Self::convert_rgba_to_yuyv(src, dst, cp),
(Rgba, PlanarRgb) => Self::convert_rgba_to_8bps(src, dst),
(Rgba, PlanarRgba) => Self::convert_rgba_to_prgba(src, dst),
(Rgba, Nv16) => Self::convert_rgba_to_nv16(src, dst, cp),
(Rgb, Rgb) => Self::copy_image(src, dst),
(Rgb, Rgba) => Self::convert_rgb_to_rgba(src, dst),
(Rgb, Grey) => Self::convert_rgb_to_grey(src, dst),
(Rgb, Yuyv) => Self::convert_rgb_to_yuyv(src, dst, cp),
(Rgb, PlanarRgb) => Self::convert_rgb_to_8bps(src, dst),
(Rgb, PlanarRgba) => Self::convert_rgb_to_prgba(src, dst),
(Rgb, Nv16) => Self::convert_rgb_to_nv16(src, dst, cp),
(Grey, Rgb) => Self::convert_grey_to_rgb(src, dst),
(Grey, Rgba) => Self::convert_grey_to_rgba(src, dst),
(Grey, Grey) => Self::copy_image(src, dst),
(Grey, Yuyv) => Self::convert_grey_to_yuyv(src, dst, cp),
(Grey, PlanarRgb) => Self::convert_grey_to_8bps(src, dst),
(Grey, PlanarRgba) => Self::convert_grey_to_prgba(src, dst),
(Grey, Nv16) => Self::convert_grey_to_nv16(src, dst, cp),
(Nv16, Rgb) => Self::convert_nv16_to_rgb(src, dst, cp),
(Nv16, Rgba) => Self::convert_nv16_to_rgba(src, dst, cp),
(Nv24, Rgb) => Self::convert_nv24_to_rgb(src, dst, cp),
(Nv24, Rgba) => Self::convert_nv24_to_rgba(src, dst, cp),
(Nv24, Grey) => Self::convert_nv24_to_grey(src, dst, cp),
(PlanarRgb, Rgb) => Self::convert_8bps_to_rgb(src, dst),
(PlanarRgb, Rgba) => Self::convert_8bps_to_rgba(src, dst),
(PlanarRgba, Rgb) => Self::convert_prgba_to_rgb(src, dst),
(PlanarRgba, Rgba) => Self::convert_prgba_to_rgba(src, dst),
(Bgra, Bgra) => Self::copy_image(src, dst),
(Nv12, Bgra) => {
Self::convert_nv12_to_rgba(src, dst, cp)?;
Self::swizzle_rb_4chan(dst)
}
(Nv16, Bgra) => {
Self::convert_nv16_to_rgba(src, dst, cp)?;
Self::swizzle_rb_4chan(dst)
}
(Nv24, Bgra) => {
Self::convert_nv24_to_rgba(src, dst, cp)?;
Self::swizzle_rb_4chan(dst)
}
(Yuyv, Bgra) => {
Self::convert_yuyv_to_rgba(src, dst, cp)?;
Self::swizzle_rb_4chan(dst)
}
(Vyuy, Bgra) => {
Self::convert_vyuy_to_rgba(src, dst, cp)?;
Self::swizzle_rb_4chan(dst)
}
(Rgba, Bgra) => {
Self::copy_image(src, dst)?;
Self::swizzle_rb_4chan(dst)
}
(Rgb, Bgra) => {
Self::convert_rgb_to_rgba(src, dst)?;
Self::swizzle_rb_4chan(dst)
}
(Grey, Bgra) => {
Self::convert_grey_to_rgba(src, dst)?;
Self::swizzle_rb_4chan(dst)
}
(PlanarRgb, Bgra) => {
Self::convert_8bps_to_rgba(src, dst)?;
Self::swizzle_rb_4chan(dst)
}
(PlanarRgba, Bgra) => {
Self::convert_prgba_to_rgba(src, dst)?;
Self::swizzle_rb_4chan(dst)
}
(s, d) => Err(Error::NotSupported(format!("Conversion from {s} to {d}",))),
}
}
pub(crate) fn fill_image_outside_crop_u8(
dst: &mut Tensor<u8>,
rgba: [u8; 4],
crop: Rect,
) -> Result<()> {
let dst_fmt = dst.format().unwrap();
let dst_w = dst.width().unwrap();
let dst_h = dst.height().unwrap();
let cm = crate::colorimetry::resolve_colorimetry(dst.colorimetry(), dst.height());
let cp = ColorParams {
matrix: crate::colorimetry::yuv_matrix(cm.encoding.unwrap()),
range: crate::colorimetry::yuv_range(cm.range.unwrap()),
encoding: cm.encoding.unwrap(),
range_kind: cm.range.unwrap(),
src_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
dst_full_range: cm.range == Some(edgefirst_tensor::ColorRange::Full),
};
let dst_stride = tensor_row_stride(dst);
let mut dst_map = dst.map_mut()?;
let dst_tup = (dst_map.as_mut_slice(), dst_w, dst_h, dst_stride);
Self::fill_outside_crop_dispatch(dst_tup, dst_fmt, rgba, crop, cp)
}
fn fill_outside_crop_dispatch(
dst: (&mut [u8], usize, usize, usize),
fmt: PixelFormat,
rgba: [u8; 4],
crop: Rect,
cp: ColorParams,
) -> Result<()> {
use PixelFormat::*;
match fmt {
Rgba | Bgra => Self::fill_image_outside_crop_(dst, rgba, crop),
Rgb => Self::fill_image_outside_crop_(dst, Self::rgba_to_rgb(rgba), crop),
Grey => Self::fill_image_outside_crop_(dst, Self::rgba_to_grey(rgba), crop),
Yuyv => {
let (bytes, w, h, stride) = dst;
let yuyv = Self::rgba_to_yuyv(rgba, cp);
Self::fill_image_outside_crop_(
(&mut *bytes, w / 2, h, stride),
yuyv,
Rect::new(crop.left / 2, crop.top, crop.width.div_ceil(2), crop.height),
)?;
if w % 2 == 1 {
let last = w - 1;
let col_outside = last < crop.left || last >= crop.left + crop.width;
for y in 0..h {
let row_outside = y < crop.top || y >= crop.top + crop.height;
if row_outside || col_outside {
let off = y * stride + last * 2;
bytes[off] = yuyv[0];
bytes[off + 1] = yuyv[1];
}
}
}
Ok(())
}
PlanarRgb => Self::fill_image_outside_crop_planar(dst, Self::rgba_to_rgb(rgba), crop),
PlanarRgba => Self::fill_image_outside_crop_planar(dst, rgba, crop),
Nv16 => {
let yuyv = Self::rgba_to_yuyv(rgba, cp);
Self::fill_image_outside_crop_yuv_semiplanar(dst, yuyv[0], [yuyv[1], yuyv[3]], crop)
}
_ => Err(Error::Internal(format!(
"Found unexpected destination {fmt}",
))),
}
}
}
impl ImageProcessorTrait for CPUProcessor {
fn convert(
&mut self,
src: &TensorDyn,
dst: &mut TensorDyn,
rotation: Rotation,
flip: Flip,
crop: Crop,
) -> Result<()> {
let crop = crop.resolve(
src.width().unwrap_or(0),
src.height().unwrap_or(0),
dst.width().unwrap_or(0),
dst.height().unwrap_or(0),
)?;
self.convert_impl(src, dst, rotation, flip, crop)
}
fn draw_decoded_masks(
&mut self,
dst: &mut TensorDyn,
detect: &[DetectBox],
segmentation: &[Segmentation],
overlay: crate::MaskOverlay<'_>,
) -> Result<()> {
prepare_dst_base_cpu(dst, overlay.background)?;
let dst = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
self.draw_decoded_masks_impl(
dst,
detect,
segmentation,
overlay.opacity,
overlay.color_mode,
)
}
fn draw_proto_masks(
&mut self,
dst: &mut TensorDyn,
detect: &[DetectBox],
proto_data: &ProtoData,
overlay: crate::MaskOverlay<'_>,
) -> Result<()> {
prepare_dst_base_cpu(dst, overlay.background)?;
let dst = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
self.draw_proto_masks_impl(
dst,
detect,
proto_data,
overlay.opacity,
overlay.letterbox,
overlay.color_mode,
)
}
fn set_class_colors(&mut self, colors: &[[u8; 4]]) -> Result<()> {
for (c, new_c) in self.colors.iter_mut().zip(colors.iter()) {
*c = *new_c;
}
Ok(())
}
}
impl CPUProcessor {
pub(crate) fn convert_impl(
&mut self,
src: &TensorDyn,
dst: &mut TensorDyn,
rotation: Rotation,
flip: Flip,
crop: ResolvedCrop,
) -> Result<()> {
let src_fmt = src.format().ok_or(Error::NotAnImage)?;
let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
let src_cm = crate::colorimetry::effective_colorimetry(src);
let dst_cm = crate::colorimetry::effective_colorimetry(dst);
let src_full = src_cm.range == Some(edgefirst_tensor::ColorRange::Full);
let dst_full = dst_cm.range == Some(edgefirst_tensor::ColorRange::Full);
let src_params = ColorParams {
matrix: crate::colorimetry::yuv_matrix(src_cm.encoding.unwrap()),
range: crate::colorimetry::yuv_range(src_cm.range.unwrap()),
encoding: src_cm.encoding.unwrap(),
range_kind: src_cm.range.unwrap(),
src_full_range: src_full,
dst_full_range: dst_full,
};
let dst_params = ColorParams {
matrix: crate::colorimetry::yuv_matrix(dst_cm.encoding.unwrap()),
range: crate::colorimetry::yuv_range(dst_cm.range.unwrap()),
encoding: dst_cm.encoding.unwrap(),
range_kind: dst_cm.range.unwrap(),
src_full_range: src_full,
dst_full_range: dst_full,
};
match (src.dtype(), dst.dtype()) {
(DType::U8, DType::U8) => {
let src = src.as_u8().unwrap();
let dst = dst.as_u8_mut().unwrap();
self.convert_u8(
src, dst, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
)
}
(DType::U8, DType::I8) => {
let src_u8 = src.as_u8().unwrap();
let dst_i8 = dst.as_i8_mut().unwrap();
let dst_u8 = unsafe { &mut *(dst_i8 as *mut Tensor<i8> as *mut Tensor<u8>) };
self.convert_u8(
src_u8, dst_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params, dst_params,
)?;
apply_int8_xor_bias_rows(dst_u8, dst_fmt)
}
(DType::U8, d @ (DType::F32 | DType::F16)) => {
let src_u8 = src.as_u8().unwrap();
let dw = dst.width().ok_or(Error::NotAnImage)?;
let dh = dst.height().ok_or(Error::NotAnImage)?;
let scratch_matches = self.widen_scratch.as_ref().is_some_and(|t| {
t.width() == Some(dw) && t.height() == Some(dh) && t.format() == Some(dst_fmt)
});
let mut tmp = if scratch_matches {
self.widen_scratch.take().unwrap()
} else {
TensorDyn::image(
dw,
dh,
dst_fmt,
DType::U8,
Some(TensorMemory::Mem),
edgefirst_tensor::CpuAccess::ReadWrite,
)?
};
{
let tmp_u8 = tmp.as_u8_mut().unwrap();
self.convert_u8(
src_u8, tmp_u8, src_fmt, dst_fmt, rotation, flip, crop, src_params,
dst_params,
)?;
}
{
let tmp_u8 = tmp.as_u8().unwrap();
let (rows, row_len) = logical_surface(tmp_u8)?;
let src_stride = tensor_row_stride(tmp_u8);
let dst_stride_bytes = dst.effective_row_stride().ok_or(Error::NotAnImage)?;
let src_map = tmp_u8.map_read()?;
guard_plane(
src_map.as_slice().len(),
src_stride,
rows,
row_len,
"widen src",
)?;
let src_rows = src_map.as_slice().chunks(src_stride).take(rows);
let elem = d.size();
if !dst_stride_bytes.is_multiple_of(elem) {
return Err(Error::InvalidShape(format!(
"{d} destination row stride {dst_stride_bytes} is not a multiple of \
the element size {elem}"
)));
}
let dst_stride = dst_stride_bytes / elem;
match d {
DType::F32 => {
let dst_t = dst.as_f32_mut().unwrap();
let mut dst_map = dst_t.map_mut()?;
guard_plane(
dst_map.as_slice().len(),
dst_stride,
rows,
row_len,
"widen f32 dst",
)?;
for (s, dr) in
src_rows.zip(dst_map.as_mut_slice().chunks_mut(dst_stride))
{
simd::widen_u8_to_f32_norm(s, &mut dr[..row_len]);
}
}
DType::F16 => {
let dst_t = dst.as_f16_mut().unwrap();
let mut dst_map = dst_t.map_mut()?;
guard_plane(
dst_map.as_slice().len(),
dst_stride,
rows,
row_len,
"widen f16 dst",
)?;
for (s, dr) in
src_rows.zip(dst_map.as_mut_slice().chunks_mut(dst_stride))
{
simd::widen_u8_to_f16_norm(s, &mut dr[..row_len]);
}
}
_ => unreachable!(),
}
}
self.widen_scratch = Some(tmp);
Ok(())
}
(s, d) => Err(Error::NotSupported(format!("dtype {s} -> {d}",))),
}
}
fn reuse_or_alloc_image(
cached: Option<Tensor<u8>>,
w: usize,
h: usize,
fmt: PixelFormat,
) -> Result<Tensor<u8>> {
if let Some(t) = cached {
if t.width() == Some(w) && t.height() == Some(h) && t.format() == Some(fmt) {
return Ok(t);
}
}
Ok(Tensor::<u8>::image(
w,
h,
fmt,
Some(TensorMemory::Mem),
edgefirst_tensor::CpuAccess::ReadWrite,
)?)
}
fn pre_resize_region(
&self,
src_fmt: PixelFormat,
(src_w, src_h): (usize, usize),
(dst_w, dst_h): (usize, usize),
rotation: Rotation,
crop: ResolvedCrop,
) -> Option<Rect> {
use PixelFormat::{Nv12, Nv16, Nv24};
if !matches!(src_fmt, Nv12 | Nv16 | Nv24) {
return None;
}
let r = crop.src_rect?;
let full_src = Rect {
left: 0,
top: 0,
width: src_w,
height: src_h,
};
if r == full_src {
return None;
}
let d = crop.dst_rect.unwrap_or(Rect {
left: 0,
top: 0,
width: dst_w,
height: dst_h,
});
let (dst_x, dst_y) = match rotation {
Rotation::None | Rotation::Rotate180 => (d.width, d.height),
Rotation::Clockwise90 | Rotation::CounterClockwise90 => (d.height, d.width),
};
let halo_x = self.filter_halo(r.width, dst_x)?;
let halo_y = self.filter_halo(r.height, dst_y)?;
let mut left = r.left.saturating_sub(halo_x);
let mut top = r.top.saturating_sub(halo_y);
let mut right = (r.left + r.width + halo_x).min(src_w);
let mut bottom = (r.top + r.height + halo_y).min(src_h);
let (align_x, align_y) = match src_fmt {
Nv12 => (2, 2),
Nv16 => (2, 1),
_ => (1, 1),
};
left -= left % align_x;
top -= top % align_y;
right = right.next_multiple_of(align_x).min(src_w);
bottom = bottom.next_multiple_of(align_y).min(src_h);
let grown = Rect {
left,
top,
width: right - left,
height: bottom - top,
};
(grown != full_src).then_some(grown)
}
#[allow(clippy::too_many_arguments)]
fn convert_u8(
&mut self,
src: &Tensor<u8>,
dst: &mut Tensor<u8>,
src_fmt: PixelFormat,
dst_fmt: PixelFormat,
rotation: Rotation,
flip: Flip,
crop: ResolvedCrop,
src_params: ColorParams,
dst_params: ColorParams,
) -> Result<()> {
use PixelFormat::*;
#[cfg(test)]
{
self.last_tmp_dims = None;
}
let src_w = src.width().unwrap();
let src_h = src.height().unwrap();
let dst_w = dst.width().unwrap();
let dst_h = dst.height().unwrap();
crop.check_crop_dims(src_w, src_h, dst_w, dst_h)?;
let intermediate = match (src_fmt, dst_fmt) {
(Nv12, Rgb) => Rgb,
(Nv12, Rgba) => Rgba,
(Nv12, Grey) => Grey,
(Nv12, Yuyv) => Rgba,
(Nv12, Nv16) => Rgba,
(Nv12, PlanarRgb) => Rgb,
(Nv12, PlanarRgba) => Rgba,
(Nv16, PlanarRgb) => Rgb,
(Nv16, PlanarRgba) => Rgba,
(Nv24, PlanarRgb) => Rgb,
(Nv24, PlanarRgba) => Rgba,
(Yuyv, Rgb) => Rgb,
(Yuyv, Rgba) => Rgba,
(Yuyv, Grey) => Grey,
(Yuyv, Yuyv) => Rgba,
(Yuyv, PlanarRgb) => Rgb,
(Yuyv, PlanarRgba) => Rgba,
(Yuyv, Nv16) => Rgba,
(Vyuy, Rgb) => Rgb,
(Vyuy, Rgba) => Rgba,
(Vyuy, Grey) => Grey,
(Vyuy, Vyuy) => Rgba,
(Vyuy, PlanarRgb) => Rgb,
(Vyuy, PlanarRgba) => Rgba,
(Vyuy, Nv16) => Rgba,
(Rgba, Rgb) => Rgba,
(Rgba, Rgba) => Rgba,
(Rgba, Grey) => Grey,
(Rgba, Yuyv) => Rgba,
(Rgba, PlanarRgb) => Rgba,
(Rgba, PlanarRgba) => Rgba,
(Rgba, Nv16) => Rgba,
(Rgb, Rgb) => Rgb,
(Rgb, Rgba) => Rgb,
(Rgb, Grey) => Grey,
(Rgb, Yuyv) => Rgb,
(Rgb, PlanarRgb) => Rgb,
(Rgb, PlanarRgba) => Rgb,
(Rgb, Nv16) => Rgb,
(Grey, Rgb) => Rgb,
(Grey, Rgba) => Rgba,
(Grey, Grey) => Grey,
(Grey, Yuyv) => Grey,
(Grey, PlanarRgb) => Grey,
(Grey, PlanarRgba) => Grey,
(Grey, Nv16) => Grey,
(Nv12, Bgra) => Rgba,
(Yuyv, Bgra) => Rgba,
(Vyuy, Bgra) => Rgba,
(Rgba, Bgra) => Rgba,
(Rgb, Bgra) => Rgb,
(Grey, Bgra) => Grey,
(Bgra, Bgra) => Bgra,
(Nv16, Rgb) => Rgb,
(Nv16, Rgba) => Rgba,
(Nv16, Bgra) => Rgba,
(Nv24, Rgb) => Rgb,
(Nv24, Rgba) => Rgba,
(Nv24, Grey) => Grey,
(Nv24, Bgra) => Rgba,
(PlanarRgb, Rgb) => Rgb,
(PlanarRgb, Rgba) => Rgb,
(PlanarRgb, Bgra) => Rgb,
(PlanarRgba, Rgb) => Rgba,
(PlanarRgba, Rgba) => Rgba,
(PlanarRgba, Bgra) => Rgba,
(s, d) => {
return Err(Error::NotSupported(format!("Conversion from {s} to {d}",)));
}
};
let need_resize_flip_rotation = rotation != Rotation::None
|| flip != Flip::None
|| src_w != dst_w
|| src_h != dst_h
|| crop.src_rect.is_some_and(|c| {
c != Rect {
left: 0,
top: 0,
width: src_w,
height: src_h,
}
})
|| crop.dst_rect.is_some_and(|c| {
c != Rect {
left: 0,
top: 0,
width: dst_w,
height: dst_h,
}
});
let direct_is_yuv_src = matches!(src_fmt, Nv12 | Nv16 | Nv24 | Yuyv | Vyuy);
let direct_params = if direct_is_yuv_src {
src_params
} else {
dst_params
};
let full_dst_rect = Rect {
left: 0,
top: 0,
width: dst_w,
height: dst_h,
};
let fused_region = if rotation != Rotation::None || flip != Flip::None {
None
} else {
match crop.src_rect {
None if src_w == dst_w && src_h == dst_h => Some(None),
None => None,
Some(r)
if r.width == dst_w
&& r.height == dst_h
&& crop.dst_rect.is_none_or(|d| d == full_dst_rect)
&& chroma_alignment_ok(src_fmt, r) =>
{
Some(Some(r))
}
Some(_) => None,
}
};
if let Some(region) = fused_region {
if matches!(src_fmt, Nv12 | Nv16 | Nv24) && matches!(dst_fmt, PlanarRgb | PlanarRgba) {
#[cfg(test)]
{
self.fused_hits += 1;
}
return self.convert_nv_to_planar_fused(
src,
dst,
src_fmt,
dst_fmt,
direct_params,
region,
);
}
}
if !need_resize_flip_rotation && Self::support_conversion_pf(src_fmt, dst_fmt) {
return Self::convert_format_pf(src, dst, src_fmt, dst_fmt, direct_params);
}
if dst_fmt == Yuyv && !dst_w.is_multiple_of(2) {
return Err(Error::NotSupported(format!(
"{} destination must have width divisible by 2",
dst_fmt,
)));
}
let mut cached_tmp = self.convert_tmp.take();
let mut cached_tmp2 = self.convert_tmp2.take();
let pre_region = if intermediate != src_fmt {
self.pre_resize_region(src_fmt, (src_w, src_h), (dst_w, dst_h), rotation, crop)
} else {
None
};
let tmp_holder: Option<Tensor<u8>> = if intermediate != src_fmt {
let _s = tracing::trace_span!(
"image.convert.cpu.format_convert",
from = ?src_fmt,
to = ?intermediate,
pass = "pre_resize",
)
.entered();
let (tmp_w, tmp_h) = pre_region.map_or((src_w, src_h), |g| (g.width, g.height));
let mut t = Self::reuse_or_alloc_image(cached_tmp.take(), tmp_w, tmp_h, intermediate)?;
#[cfg(test)]
{
self.last_tmp_dims = Some((tmp_w, tmp_h));
}
match pre_region {
Some(g) => {
let mut sub = Self::reuse_or_alloc_image(
self.convert_src_sub.take(),
g.width,
g.height,
src_fmt,
)?;
{
let _s = tracing::trace_span!(
"image.convert.cpu.extract_region",
region_w = g.width,
region_h = g.height,
)
.entered();
Self::extract_nv_region(src, &mut sub, src_fmt, g)?;
}
Self::convert_format_pf(&sub, &mut t, src_fmt, intermediate, src_params)?;
self.convert_src_sub = Some(sub);
}
None => Self::convert_format_pf(src, &mut t, src_fmt, intermediate, src_params)?,
}
Some(t)
} else {
None
};
let crop = match (pre_region, crop.src_rect) {
(Some(g), Some(r)) => ResolvedCrop {
src_rect: Some(Rect {
left: r.left - g.left,
top: r.top - g.top,
..r
}),
..crop
},
_ => crop,
};
let (tmp, tmp_fmt): (&Tensor<u8>, PixelFormat) = match &tmp_holder {
Some(t) => (t, intermediate),
None => (src, src_fmt),
};
debug_assert!(matches!(tmp_fmt, Rgb | Rgba | Grey));
if tmp_fmt == dst_fmt {
let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
self.resize_flip_rotate_pf(tmp, dst, dst_fmt, rotation, flip, crop)?;
} else if !need_resize_flip_rotation {
let _s = tracing::trace_span!(
"image.convert.cpu.format_convert",
from = ?tmp_fmt,
to = ?dst_fmt,
pass = "direct",
)
.entered();
Self::convert_format_pf(tmp, dst, tmp_fmt, dst_fmt, dst_params)?;
} else {
let mut tmp2 = Self::reuse_or_alloc_image(cached_tmp2.take(), dst_w, dst_h, tmp_fmt)?;
if crop.dst_rect.is_some_and(|c| {
c != Rect {
left: 0,
top: 0,
width: dst_w,
height: dst_h,
}
}) && crop.dst_color.is_none()
{
Self::convert_format_pf(dst, &mut tmp2, dst_fmt, tmp_fmt, dst_params)?;
}
{
let _s = tracing::trace_span!("image.convert.cpu.resize_flip_rotate").entered();
self.resize_flip_rotate_pf(tmp, &mut tmp2, tmp_fmt, rotation, flip, crop)?;
}
{
let _s = tracing::trace_span!(
"image.convert.cpu.format_convert",
from = ?tmp_fmt,
to = ?dst_fmt,
pass = "post_resize",
)
.entered();
Self::convert_format_pf(&tmp2, dst, tmp_fmt, dst_fmt, dst_params)?;
}
cached_tmp2 = Some(tmp2);
}
if let Some(t) = tmp_holder {
cached_tmp = Some(t);
}
self.convert_tmp = cached_tmp;
self.convert_tmp2 = cached_tmp2;
if let (Some(dst_rect), Some(dst_color)) = (crop.dst_rect, crop.dst_color) {
let full_rect = Rect {
left: 0,
top: 0,
width: dst_w,
height: dst_h,
};
if dst_rect != full_rect {
Self::fill_image_outside_crop_u8(dst, dst_color, dst_rect)?;
}
}
Ok(())
}
fn draw_decoded_masks_impl(
&mut self,
dst: &mut Tensor<u8>,
detect: &[DetectBox],
segmentation: &[Segmentation],
opacity: f32,
color_mode: crate::ColorMode,
) -> Result<()> {
let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
return Err(crate::Error::NotSupported(
"CPU image rendering only supports RGBA or RGB images".to_string(),
));
}
let _timer = FunctionTimer::new("CPUProcessor::draw_decoded_masks");
let dst_w = dst.width().unwrap();
let dst_h = dst.height().unwrap();
let dst_rs = tensor_row_stride(dst);
let dst_c = dst_fmt.channels();
let mut map = dst.map_mut()?;
let dst_slice = map.as_mut_slice();
self.render_box(dst_w, dst_h, dst_rs, dst_c, dst_slice, detect, color_mode)?;
if segmentation.is_empty() {
return Ok(());
}
let is_semantic = segmentation[0].segmentation.shape()[2] > 1;
if is_semantic {
self.render_modelpack_segmentation(
dst_w,
dst_h,
dst_rs,
dst_c,
dst_slice,
&segmentation[0],
opacity,
)?;
} else {
for (idx, (seg, det)) in segmentation.iter().zip(detect).enumerate() {
let color_index = color_mode.index(idx, det.label);
self.render_yolo_segmentation(
dst_w,
dst_h,
dst_rs,
dst_c,
dst_slice,
seg,
color_index,
opacity,
)?;
}
}
Ok(())
}
fn draw_proto_masks_impl(
&mut self,
dst: &mut Tensor<u8>,
detect: &[DetectBox],
proto_data: &ProtoData,
opacity: f32,
letterbox: Option<[f32; 4]>,
color_mode: crate::ColorMode,
) -> Result<()> {
let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
return Err(crate::Error::NotSupported(
"CPU image rendering only supports RGBA or RGB images".to_string(),
));
}
let _timer = FunctionTimer::new("CPUProcessor::draw_proto_masks");
let dst_w = dst.width().unwrap();
let dst_h = dst.height().unwrap();
let dst_rs = tensor_row_stride(dst);
let channels = dst_fmt.channels();
let mut map = dst.map_mut()?;
let dst_slice = map.as_mut_slice();
self.render_box(
dst_w, dst_h, dst_rs, channels, dst_slice, detect, color_mode,
)?;
if detect.is_empty() {
return Ok(());
}
let proto_shape = proto_data.protos.shape();
if proto_shape.len() != 3 {
return Err(Error::InvalidShape(format!(
"protos tensor must be rank-3, got {proto_shape:?}"
)));
}
let proto_h = proto_shape[0];
let proto_w = proto_shape[1];
let num_protos = proto_shape[2];
let coeff_shape = proto_data.mask_coefficients.shape();
if coeff_shape.len() != 2 {
return Err(Error::InvalidShape(format!(
"mask_coefficients tensor must be rank-2, got {coeff_shape:?}"
)));
}
if coeff_shape[0] == 0 {
return Ok(());
}
if coeff_shape[1] != num_protos {
return Err(Error::InvalidShape(format!(
"mask_coefficients second dimension must match num_protos \
({num_protos}), got {coeff_shape:?}"
)));
}
let coeff_f32: Vec<f32> = match proto_data.mask_coefficients.dtype() {
DType::F32 => {
let t = proto_data.mask_coefficients.as_f32().expect("F32");
let m = t.map_read()?;
m.as_slice().to_vec()
}
DType::F16 => {
let t = proto_data.mask_coefficients.as_f16().expect("F16");
let m = t.map_read()?;
m.as_slice().iter().map(|v| v.to_f32()).collect()
}
DType::I8 => {
let t = proto_data.mask_coefficients.as_i8().expect("I8");
let m = t.map_read()?;
if let Some(q) = t.quantization() {
use edgefirst_tensor::QuantMode;
let (scale, zp) = match q.mode() {
QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
other => {
return Err(Error::NotSupported(format!(
"I8 mask_coefficients quantization mode {other:?} not supported"
)));
}
};
m.as_slice()
.iter()
.map(|&v| (v as f32 - zp) * scale)
.collect()
} else {
m.as_slice().iter().map(|&v| v as f32).collect()
}
}
DType::I16 => {
let t = proto_data.mask_coefficients.as_i16().expect("I16");
let m = t.map_read()?;
if let Some(q) = t.quantization() {
use edgefirst_tensor::QuantMode;
let (scale, zp) = match q.mode() {
QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
other => {
return Err(Error::NotSupported(format!(
"I16 mask_coefficients quantization mode {other:?} not supported"
)));
}
};
m.as_slice()
.iter()
.map(|&v| (v as f32 - zp) * scale)
.collect()
} else {
m.as_slice().iter().map(|&v| v as f32).collect()
}
}
other => {
return Err(Error::InvalidShape(format!(
"mask_coefficients dtype {other:?} not supported"
)));
}
};
let (lx0, lx_range, ly0, ly_range) = match letterbox {
Some([lx0, ly0, lx1, ly1]) => (lx0, lx1 - lx0, ly0, ly1 - ly0),
None => (0.0_f32, 1.0_f32, 0.0_f32, 1.0_f32),
};
match proto_data.protos.dtype() {
DType::F32 => {
let t = proto_data.protos.as_f32().expect("F32");
let m = t.map_read()?;
self.draw_proto_masks_inner(
dst_slice,
dst_w,
dst_h,
dst_rs,
channels,
detect,
m.as_slice(),
&coeff_f32,
proto_h,
proto_w,
num_protos,
opacity,
(lx0, lx_range, ly0, ly_range),
color_mode,
0.0_f32,
|p: &f32, _| *p,
);
}
DType::F16 => {
let t = proto_data.protos.as_f16().expect("F16");
let m = t.map_read()?;
self.draw_proto_masks_inner(
dst_slice,
dst_w,
dst_h,
dst_rs,
channels,
detect,
m.as_slice(),
&coeff_f32,
proto_h,
proto_w,
num_protos,
opacity,
(lx0, lx_range, ly0, ly_range),
color_mode,
0.0_f32,
|p: &half::f16, _| p.to_f32(),
);
}
DType::I8 => {
use edgefirst_tensor::QuantMode;
let t = proto_data.protos.as_i8().expect("I8");
let m = t.map_read()?;
let quant = t.quantization().ok_or_else(|| {
Error::InvalidShape("I8 protos require quantization metadata".into())
})?;
let (scale, zp) = match quant.mode() {
QuantMode::PerTensor { scale, zero_point } => (scale, zero_point as f32),
QuantMode::PerTensorSymmetric { scale } => (scale, 0.0),
QuantMode::PerChannel { axis, .. }
| QuantMode::PerChannelSymmetric { axis, .. } => {
return Err(Error::NotSupported(format!(
"per-channel quantization (axis={axis}) in draw_proto_masks \
CPU path not yet supported"
)));
}
};
self.draw_proto_masks_inner(
dst_slice,
dst_w,
dst_h,
dst_rs,
channels,
detect,
m.as_slice(),
&coeff_f32,
proto_h,
proto_w,
num_protos,
opacity,
(lx0, lx_range, ly0, ly_range),
color_mode,
scale,
move |p: &i8, _| (*p as f32) - zp,
);
}
other => {
return Err(Error::InvalidShape(format!(
"proto tensor dtype {other:?} not supported"
)));
}
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn draw_proto_masks_inner<P: Copy>(
&self,
dst_slice: &mut [u8],
dst_w: usize,
dst_h: usize,
dst_rs: usize,
channels: usize,
detect: &[DetectBox],
protos: &[P],
coeff_all_f32: &[f32],
proto_h: usize,
proto_w: usize,
num_protos: usize,
opacity: f32,
letterbox_xy: (f32, f32, f32, f32),
color_mode: crate::ColorMode,
acc_scale: f32,
load_f32: impl Fn(&P, f32) -> f32 + Copy,
) {
let (lx0, lx_range, ly0, ly_range) = letterbox_xy;
let stride_y = proto_w * num_protos;
for (idx, det) in detect.iter().enumerate() {
let coeff = &coeff_all_f32[idx * num_protos..(idx + 1) * num_protos];
let color_index = color_mode.index(idx, det.label);
let color = self.colors[color_index % self.colors.len()];
let alpha = if opacity == 1.0 {
color[3] as u16
} else {
(color[3] as f32 * opacity).round() as u16
};
let start_x = (dst_w as f32 * det.bbox.xmin).round() as usize;
let start_y = (dst_h as f32 * det.bbox.ymin).round() as usize;
let end_x = ((dst_w as f32 * det.bbox.xmax).round() as usize).min(dst_w);
let end_y = ((dst_h as f32 * det.bbox.ymax).round() as usize).min(dst_h);
for y in start_y..end_y {
for x in start_x..end_x {
let px = (lx0 + (x as f32 / dst_w as f32) * lx_range) * proto_w as f32 - 0.5;
let py = (ly0 + (y as f32 / dst_h as f32) * ly_range) * proto_h as f32 - 0.5;
let x0 = (px.floor() as isize).clamp(0, proto_w as isize - 1) as usize;
let y0 = (py.floor() as isize).clamp(0, proto_h as isize - 1) as usize;
let x1 = (x0 + 1).min(proto_w - 1);
let y1 = (y0 + 1).min(proto_h - 1);
let fx = px - px.floor();
let fy = py - py.floor();
let w00 = (1.0 - fx) * (1.0 - fy);
let w10 = fx * (1.0 - fy);
let w01 = (1.0 - fx) * fy;
let w11 = fx * fy;
let b00 = y0 * stride_y + x0 * num_protos;
let b10 = y0 * stride_y + x1 * num_protos;
let b01 = y1 * stride_y + x0 * num_protos;
let b11 = y1 * stride_y + x1 * num_protos;
let mut acc = 0.0_f32;
for p in 0..num_protos {
let v00 = load_f32(&protos[b00 + p], 0.0);
let v10 = load_f32(&protos[b10 + p], 0.0);
let v01 = load_f32(&protos[b01 + p], 0.0);
let v11 = load_f32(&protos[b11 + p], 0.0);
let val = w00 * v00 + w10 * v10 + w01 * v01 + w11 * v11;
acc += coeff[p] * val;
}
let final_acc = if acc_scale == 0.0 {
acc
} else {
acc_scale * acc
};
let mask = 1.0 / (1.0 + (-final_acc).exp());
if mask < 0.5 {
continue;
}
let dst_index = y * dst_rs + x * channels;
for c in 0..3 {
dst_slice[dst_index + c] = ((color[c] as u16 * alpha
+ dst_slice[dst_index + c] as u16 * (255 - alpha))
/ 255) as u8;
}
}
}
}
}
}