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>,
nv_strip_scratch: Vec<u8>,
convert_tmp: Option<Tensor<u8>>,
convert_tmp2: 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(),
nv_strip_scratch: Vec::new(),
convert_tmp: None,
convert_tmp2: 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 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 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;
}
}
}
}
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(),
nv_strip_scratch: Vec::new(),
convert_tmp: None,
convert_tmp2: 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(),
nv_strip_scratch: Vec::new(),
convert_tmp: None,
convert_tmp2: None,
}
}
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) => {
dst.map_mut()?.copy_from_slice(&src.map_read()?);
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 mut dst_map = dst.map_mut()?;
let dst_tup = (dst_map.as_mut_slice(), dst_w, dst_h);
Self::fill_outside_crop_dispatch(dst_tup, dst_fmt, rgba, crop, cp)
}
fn fill_outside_crop_dispatch(
dst: (&mut [u8], 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 => Self::fill_image_outside_crop_(
(dst.0, dst.1 / 2, dst.2),
Self::rgba_to_yuyv(rgba, cp),
Rect::new(crop.left / 2, crop.top, crop.width.div_ceil(2), crop.height),
),
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,
)?;
let mut map = dst_u8.map_mut()?;
apply_int8_xor_bias(map.as_mut_slice(), dst_fmt);
Ok(())
}
(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 src_map = tmp_u8.map_read()?;
match d {
DType::F32 => {
let dst_t = dst.as_f32_mut().unwrap();
let mut dst_map = dst_t.map_mut()?;
debug_assert_eq!(src_map.as_slice().len(), dst_map.as_slice().len());
simd::widen_u8_to_f32_norm(src_map.as_slice(), dst_map.as_mut_slice());
}
DType::F16 => {
let dst_t = dst.as_f16_mut().unwrap();
let mut dst_map = dst_t.map_mut()?;
debug_assert_eq!(src_map.as_slice().len(), dst_map.as_slice().len());
simd::widen_u8_to_f16_norm(src_map.as_slice(), dst_map.as_mut_slice());
}
_ => 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,
)?)
}
#[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::*;
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
};
if !need_resize_flip_rotation
&& matches!(src_fmt, Nv12 | Nv16 | Nv24)
&& matches!(dst_fmt, PlanarRgb | PlanarRgba)
{
return self.convert_nv_to_planar_fused(src, dst, src_fmt, dst_fmt, direct_params);
}
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 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 mut t = Self::reuse_or_alloc_image(cached_tmp.take(), src_w, src_h, intermediate)?;
Self::convert_format_pf(src, &mut t, src_fmt, intermediate, src_params)?;
Some(t)
} else {
None
};
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;
}
}
}
}
}
}