edgefirst_image/cpu/
mod.rs

1// SPDX-FileCopyrightText: Copyright 2025 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::{Crop, Error, Flip, FunctionTimer, ImageProcessorTrait, Rect, Result, Rotation};
5use edgefirst_decoder::{DetectBox, ProtoData, Segmentation};
6use edgefirst_tensor::{
7    DType, PixelFormat, Tensor, TensorDyn, TensorMapTrait, TensorMemory, TensorTrait,
8};
9
10mod convert;
11mod masks;
12mod resize;
13mod tests;
14
15use masks::bilinear_dot;
16
17/// CPUConverter implements the ImageProcessor trait using the fallback CPU
18/// implementation for image processing.
19#[derive(Debug, Clone)]
20pub struct CPUProcessor {
21    resizer: fast_image_resize::Resizer,
22    options: fast_image_resize::ResizeOptions,
23    colors: [[u8; 4]; 20],
24}
25
26unsafe impl Send for CPUProcessor {}
27unsafe impl Sync for CPUProcessor {}
28
29impl Default for CPUProcessor {
30    fn default() -> Self {
31        Self::new_bilinear()
32    }
33}
34
35/// Compute row stride for a packed-format Tensor<u8> image given its format.
36fn row_stride_for(width: usize, fmt: PixelFormat) -> usize {
37    use edgefirst_tensor::PixelLayout;
38    match fmt.layout() {
39        PixelLayout::Packed => width * fmt.channels(),
40        PixelLayout::Planar | PixelLayout::SemiPlanar => width,
41        _ => width, // fallback for non-exhaustive
42    }
43}
44
45/// Apply XOR 0x80 bias to color channels only, preserving alpha.
46///
47/// Matches GL int8 shader behavior: `vec4(int8_bias(c.rgb), c.a)`.
48/// For formats without alpha, XORs every byte (fast path).
49pub(crate) fn apply_int8_xor_bias(data: &mut [u8], fmt: PixelFormat) {
50    use edgefirst_tensor::PixelLayout;
51    if !fmt.has_alpha() {
52        for b in data.iter_mut() {
53            *b ^= 0x80;
54        }
55    } else if fmt.layout() == PixelLayout::Planar {
56        // Planar with alpha (e.g. PlanarRgba): XOR color planes, skip alpha plane.
57        let channels = fmt.channels();
58        let plane_size = data.len() / channels;
59        for b in data[..plane_size * (channels - 1)].iter_mut() {
60            *b ^= 0x80;
61        }
62    } else {
63        // Packed with alpha (Rgba, Bgra): XOR color bytes, skip alpha byte.
64        let channels = fmt.channels();
65        for pixel in data.chunks_exact_mut(channels) {
66            for b in &mut pixel[..channels - 1] {
67                *b ^= 0x80;
68            }
69        }
70    }
71}
72
73impl CPUProcessor {
74    /// Creates a new CPUConverter with bilinear resizing.
75    pub fn new() -> Self {
76        Self::new_bilinear()
77    }
78
79    /// Creates a new CPUConverter with bilinear resizing.
80    fn new_bilinear() -> Self {
81        let resizer = fast_image_resize::Resizer::new();
82        let options = fast_image_resize::ResizeOptions::new()
83            .resize_alg(fast_image_resize::ResizeAlg::Convolution(
84                fast_image_resize::FilterType::Bilinear,
85            ))
86            .use_alpha(false);
87
88        log::debug!("CPUConverter created");
89        Self {
90            resizer,
91            options,
92            colors: crate::DEFAULT_COLORS_U8,
93        }
94    }
95
96    /// Creates a new CPUConverter with nearest neighbor resizing.
97    pub fn new_nearest() -> Self {
98        let resizer = fast_image_resize::Resizer::new();
99        let options = fast_image_resize::ResizeOptions::new()
100            .resize_alg(fast_image_resize::ResizeAlg::Nearest)
101            .use_alpha(false);
102        log::debug!("CPUConverter created");
103        Self {
104            resizer,
105            options,
106            colors: crate::DEFAULT_COLORS_U8,
107        }
108    }
109
110    pub(crate) fn support_conversion_pf(src: PixelFormat, dst: PixelFormat) -> bool {
111        use PixelFormat::*;
112        matches!(
113            (src, dst),
114            (Nv12, Rgb)
115                | (Nv12, Rgba)
116                | (Nv12, Grey)
117                | (Nv16, Rgb)
118                | (Nv16, Rgba)
119                | (Nv16, Bgra)
120                | (Yuyv, Rgb)
121                | (Yuyv, Rgba)
122                | (Yuyv, Grey)
123                | (Yuyv, Yuyv)
124                | (Yuyv, PlanarRgb)
125                | (Yuyv, PlanarRgba)
126                | (Yuyv, Nv16)
127                | (Vyuy, Rgb)
128                | (Vyuy, Rgba)
129                | (Vyuy, Grey)
130                | (Vyuy, Vyuy)
131                | (Vyuy, PlanarRgb)
132                | (Vyuy, PlanarRgba)
133                | (Vyuy, Nv16)
134                | (Rgba, Rgb)
135                | (Rgba, Rgba)
136                | (Rgba, Grey)
137                | (Rgba, Yuyv)
138                | (Rgba, PlanarRgb)
139                | (Rgba, PlanarRgba)
140                | (Rgba, Nv16)
141                | (Rgb, Rgb)
142                | (Rgb, Rgba)
143                | (Rgb, Grey)
144                | (Rgb, Yuyv)
145                | (Rgb, PlanarRgb)
146                | (Rgb, PlanarRgba)
147                | (Rgb, Nv16)
148                | (Grey, Rgb)
149                | (Grey, Rgba)
150                | (Grey, Grey)
151                | (Grey, Yuyv)
152                | (Grey, PlanarRgb)
153                | (Grey, PlanarRgba)
154                | (Grey, Nv16)
155                | (Nv12, Bgra)
156                | (Yuyv, Bgra)
157                | (Vyuy, Bgra)
158                | (Rgba, Bgra)
159                | (Rgb, Bgra)
160                | (Grey, Bgra)
161                | (Bgra, Bgra)
162                | (PlanarRgb, Rgb)
163                | (PlanarRgb, Rgba)
164                | (PlanarRgba, Rgb)
165                | (PlanarRgba, Rgba)
166                | (PlanarRgb, Bgra)
167                | (PlanarRgba, Bgra)
168        )
169    }
170
171    /// Format conversion dispatch for Tensor<u8> with PixelFormat metadata.
172    pub(crate) fn convert_format_pf(
173        src: &Tensor<u8>,
174        dst: &mut Tensor<u8>,
175        src_fmt: PixelFormat,
176        dst_fmt: PixelFormat,
177    ) -> Result<()> {
178        let _timer = FunctionTimer::new(format!(
179            "ImageProcessor::convert_format {} to {}",
180            src_fmt, dst_fmt,
181        ));
182
183        use PixelFormat::*;
184        match (src_fmt, dst_fmt) {
185            (Nv12, Rgb) => Self::convert_nv12_to_rgb(src, dst),
186            (Nv12, Rgba) => Self::convert_nv12_to_rgba(src, dst),
187            (Nv12, Grey) => Self::convert_nv12_to_grey(src, dst),
188            (Yuyv, Rgb) => Self::convert_yuyv_to_rgb(src, dst),
189            (Yuyv, Rgba) => Self::convert_yuyv_to_rgba(src, dst),
190            (Yuyv, Grey) => Self::convert_yuyv_to_grey(src, dst),
191            (Yuyv, Yuyv) => Self::copy_image(src, dst),
192            (Yuyv, PlanarRgb) => Self::convert_yuyv_to_8bps(src, dst),
193            (Yuyv, PlanarRgba) => Self::convert_yuyv_to_prgba(src, dst),
194            (Yuyv, Nv16) => Self::convert_yuyv_to_nv16(src, dst),
195            (Vyuy, Rgb) => Self::convert_vyuy_to_rgb(src, dst),
196            (Vyuy, Rgba) => Self::convert_vyuy_to_rgba(src, dst),
197            (Vyuy, Grey) => Self::convert_vyuy_to_grey(src, dst),
198            (Vyuy, Vyuy) => Self::copy_image(src, dst),
199            (Vyuy, PlanarRgb) => Self::convert_vyuy_to_8bps(src, dst),
200            (Vyuy, PlanarRgba) => Self::convert_vyuy_to_prgba(src, dst),
201            (Vyuy, Nv16) => Self::convert_vyuy_to_nv16(src, dst),
202            (Rgba, Rgb) => Self::convert_rgba_to_rgb(src, dst),
203            (Rgba, Rgba) => Self::copy_image(src, dst),
204            (Rgba, Grey) => Self::convert_rgba_to_grey(src, dst),
205            (Rgba, Yuyv) => Self::convert_rgba_to_yuyv(src, dst),
206            (Rgba, PlanarRgb) => Self::convert_rgba_to_8bps(src, dst),
207            (Rgba, PlanarRgba) => Self::convert_rgba_to_prgba(src, dst),
208            (Rgba, Nv16) => Self::convert_rgba_to_nv16(src, dst),
209            (Rgb, Rgb) => Self::copy_image(src, dst),
210            (Rgb, Rgba) => Self::convert_rgb_to_rgba(src, dst),
211            (Rgb, Grey) => Self::convert_rgb_to_grey(src, dst),
212            (Rgb, Yuyv) => Self::convert_rgb_to_yuyv(src, dst),
213            (Rgb, PlanarRgb) => Self::convert_rgb_to_8bps(src, dst),
214            (Rgb, PlanarRgba) => Self::convert_rgb_to_prgba(src, dst),
215            (Rgb, Nv16) => Self::convert_rgb_to_nv16(src, dst),
216            (Grey, Rgb) => Self::convert_grey_to_rgb(src, dst),
217            (Grey, Rgba) => Self::convert_grey_to_rgba(src, dst),
218            (Grey, Grey) => Self::copy_image(src, dst),
219            (Grey, Yuyv) => Self::convert_grey_to_yuyv(src, dst),
220            (Grey, PlanarRgb) => Self::convert_grey_to_8bps(src, dst),
221            (Grey, PlanarRgba) => Self::convert_grey_to_prgba(src, dst),
222            (Grey, Nv16) => Self::convert_grey_to_nv16(src, dst),
223
224            // the following converts are added for use in testing
225            (Nv16, Rgb) => Self::convert_nv16_to_rgb(src, dst),
226            (Nv16, Rgba) => Self::convert_nv16_to_rgba(src, dst),
227            (PlanarRgb, Rgb) => Self::convert_8bps_to_rgb(src, dst),
228            (PlanarRgb, Rgba) => Self::convert_8bps_to_rgba(src, dst),
229            (PlanarRgba, Rgb) => Self::convert_prgba_to_rgb(src, dst),
230            (PlanarRgba, Rgba) => Self::convert_prgba_to_rgba(src, dst),
231
232            // BGRA destination: convert to RGBA layout, then swap R and B
233            (Bgra, Bgra) => Self::copy_image(src, dst),
234            (Nv12, Bgra) => {
235                Self::convert_nv12_to_rgba(src, dst)?;
236                Self::swizzle_rb_4chan(dst)
237            }
238            (Nv16, Bgra) => {
239                Self::convert_nv16_to_rgba(src, dst)?;
240                Self::swizzle_rb_4chan(dst)
241            }
242            (Yuyv, Bgra) => {
243                Self::convert_yuyv_to_rgba(src, dst)?;
244                Self::swizzle_rb_4chan(dst)
245            }
246            (Vyuy, Bgra) => {
247                Self::convert_vyuy_to_rgba(src, dst)?;
248                Self::swizzle_rb_4chan(dst)
249            }
250            (Rgba, Bgra) => {
251                dst.map()?.copy_from_slice(&src.map()?);
252                Self::swizzle_rb_4chan(dst)
253            }
254            (Rgb, Bgra) => {
255                Self::convert_rgb_to_rgba(src, dst)?;
256                Self::swizzle_rb_4chan(dst)
257            }
258            (Grey, Bgra) => {
259                Self::convert_grey_to_rgba(src, dst)?;
260                Self::swizzle_rb_4chan(dst)
261            }
262            (PlanarRgb, Bgra) => {
263                Self::convert_8bps_to_rgba(src, dst)?;
264                Self::swizzle_rb_4chan(dst)
265            }
266            (PlanarRgba, Bgra) => {
267                Self::convert_prgba_to_rgba(src, dst)?;
268                Self::swizzle_rb_4chan(dst)
269            }
270
271            (s, d) => Err(Error::NotSupported(format!("Conversion from {s} to {d}",))),
272        }
273    }
274
275    /// Tensor<u8>-based fill_image_outside_crop.
276    pub(crate) fn fill_image_outside_crop_u8(
277        dst: &mut Tensor<u8>,
278        rgba: [u8; 4],
279        crop: Rect,
280    ) -> Result<()> {
281        let dst_fmt = dst.format().unwrap();
282        let dst_w = dst.width().unwrap();
283        let dst_h = dst.height().unwrap();
284        let mut dst_map = dst.map()?;
285        let dst_tup = (dst_map.as_mut_slice(), dst_w, dst_h);
286        Self::fill_outside_crop_dispatch(dst_tup, dst_fmt, rgba, crop)
287    }
288
289    /// Common fill dispatch by format.
290    fn fill_outside_crop_dispatch(
291        dst: (&mut [u8], usize, usize),
292        fmt: PixelFormat,
293        rgba: [u8; 4],
294        crop: Rect,
295    ) -> Result<()> {
296        use PixelFormat::*;
297        match fmt {
298            Rgba | Bgra => Self::fill_image_outside_crop_(dst, rgba, crop),
299            Rgb => Self::fill_image_outside_crop_(dst, Self::rgba_to_rgb(rgba), crop),
300            Grey => Self::fill_image_outside_crop_(dst, Self::rgba_to_grey(rgba), crop),
301            Yuyv => Self::fill_image_outside_crop_(
302                (dst.0, dst.1 / 2, dst.2),
303                Self::rgba_to_yuyv(rgba),
304                Rect::new(crop.left / 2, crop.top, crop.width.div_ceil(2), crop.height),
305            ),
306            PlanarRgb => Self::fill_image_outside_crop_planar(dst, Self::rgba_to_rgb(rgba), crop),
307            PlanarRgba => Self::fill_image_outside_crop_planar(dst, rgba, crop),
308            Nv16 => {
309                let yuyv = Self::rgba_to_yuyv(rgba);
310                Self::fill_image_outside_crop_yuv_semiplanar(dst, yuyv[0], [yuyv[1], yuyv[3]], crop)
311            }
312            _ => Err(Error::Internal(format!(
313                "Found unexpected destination {fmt}",
314            ))),
315        }
316    }
317}
318
319impl ImageProcessorTrait for CPUProcessor {
320    fn convert(
321        &mut self,
322        src: &TensorDyn,
323        dst: &mut TensorDyn,
324        rotation: Rotation,
325        flip: Flip,
326        crop: Crop,
327    ) -> Result<()> {
328        self.convert_impl(src, dst, rotation, flip, crop)
329    }
330
331    fn draw_decoded_masks(
332        &mut self,
333        dst: &mut TensorDyn,
334        detect: &[DetectBox],
335        segmentation: &[Segmentation],
336        overlay: crate::MaskOverlay<'_>,
337    ) -> Result<()> {
338        let dst = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
339        self.draw_decoded_masks_impl(dst, detect, segmentation, overlay.opacity)
340    }
341
342    fn draw_proto_masks(
343        &mut self,
344        dst: &mut TensorDyn,
345        detect: &[DetectBox],
346        proto_data: &ProtoData,
347        overlay: crate::MaskOverlay<'_>,
348    ) -> Result<()> {
349        let dst = dst.as_u8_mut().ok_or(Error::NotAnImage)?;
350        self.draw_proto_masks_impl(dst, detect, proto_data, overlay.opacity)
351    }
352
353    fn set_class_colors(&mut self, colors: &[[u8; 4]]) -> Result<()> {
354        for (c, new_c) in self.colors.iter_mut().zip(colors.iter()) {
355            *c = *new_c;
356        }
357        Ok(())
358    }
359}
360
361// Internal methods — dtype-aware dispatch layer.
362impl CPUProcessor {
363    /// Top-level conversion dispatcher: handles dtype combinations.
364    pub(crate) fn convert_impl(
365        &mut self,
366        src: &TensorDyn,
367        dst: &mut TensorDyn,
368        rotation: Rotation,
369        flip: Flip,
370        crop: Crop,
371    ) -> Result<()> {
372        let src_fmt = src.format().ok_or(Error::NotAnImage)?;
373        let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
374
375        match (src.dtype(), dst.dtype()) {
376            (DType::U8, DType::U8) => {
377                let src = src.as_u8().unwrap();
378                let dst = dst.as_u8_mut().unwrap();
379                self.convert_u8(src, dst, src_fmt, dst_fmt, rotation, flip, crop)
380            }
381            (DType::U8, DType::I8) => {
382                // Int8 output: reinterpret the i8 destination as u8 (layout-
383                // identical), convert directly into it, then XOR 0x80 in-place.
384                let src_u8 = src.as_u8().unwrap();
385                let dst_i8 = dst.as_i8_mut().unwrap();
386                // SAFETY: Tensor<i8> and Tensor<u8> are layout-identical
387                // (same element size, no T-dependent drop glue). Same
388                // rationale as gl::processor::tensor_i8_as_u8_mut.
389                let dst_u8 = unsafe { &mut *(dst_i8 as *mut Tensor<i8> as *mut Tensor<u8>) };
390                self.convert_u8(src_u8, dst_u8, src_fmt, dst_fmt, rotation, flip, crop)?;
391                // Apply XOR 0x80 bias in-place (u8 → i8 conversion)
392                let mut map = dst_u8.map()?;
393                apply_int8_xor_bias(map.as_mut_slice(), dst_fmt);
394                Ok(())
395            }
396            (s, d) => Err(Error::NotSupported(format!("dtype {s} -> {d}",))),
397        }
398    }
399
400    /// U8-to-U8 conversion: the full format conversion + resize pipeline.
401    #[allow(clippy::too_many_arguments)]
402    fn convert_u8(
403        &mut self,
404        src: &Tensor<u8>,
405        dst: &mut Tensor<u8>,
406        src_fmt: PixelFormat,
407        dst_fmt: PixelFormat,
408        rotation: Rotation,
409        flip: Flip,
410        crop: Crop,
411    ) -> Result<()> {
412        use PixelFormat::*;
413
414        let src_w = src.width().unwrap();
415        let src_h = src.height().unwrap();
416        let dst_w = dst.width().unwrap();
417        let dst_h = dst.height().unwrap();
418
419        crop.check_crop_dims(src_w, src_h, dst_w, dst_h)?;
420
421        // Determine intermediate format for the resize step
422        let intermediate = match (src_fmt, dst_fmt) {
423            (Nv12, Rgb) => Rgb,
424            (Nv12, Rgba) => Rgba,
425            (Nv12, Grey) => Grey,
426            (Nv12, Yuyv) => Rgba,
427            (Nv12, Nv16) => Rgba,
428            (Nv12, PlanarRgb) => Rgb,
429            (Nv12, PlanarRgba) => Rgba,
430            (Yuyv, Rgb) => Rgb,
431            (Yuyv, Rgba) => Rgba,
432            (Yuyv, Grey) => Grey,
433            (Yuyv, Yuyv) => Rgba,
434            (Yuyv, PlanarRgb) => Rgb,
435            (Yuyv, PlanarRgba) => Rgba,
436            (Yuyv, Nv16) => Rgba,
437            (Vyuy, Rgb) => Rgb,
438            (Vyuy, Rgba) => Rgba,
439            (Vyuy, Grey) => Grey,
440            (Vyuy, Vyuy) => Rgba,
441            (Vyuy, PlanarRgb) => Rgb,
442            (Vyuy, PlanarRgba) => Rgba,
443            (Vyuy, Nv16) => Rgba,
444            (Rgba, Rgb) => Rgba,
445            (Rgba, Rgba) => Rgba,
446            (Rgba, Grey) => Grey,
447            (Rgba, Yuyv) => Rgba,
448            (Rgba, PlanarRgb) => Rgba,
449            (Rgba, PlanarRgba) => Rgba,
450            (Rgba, Nv16) => Rgba,
451            (Rgb, Rgb) => Rgb,
452            (Rgb, Rgba) => Rgb,
453            (Rgb, Grey) => Grey,
454            (Rgb, Yuyv) => Rgb,
455            (Rgb, PlanarRgb) => Rgb,
456            (Rgb, PlanarRgba) => Rgb,
457            (Rgb, Nv16) => Rgb,
458            (Grey, Rgb) => Rgb,
459            (Grey, Rgba) => Rgba,
460            (Grey, Grey) => Grey,
461            (Grey, Yuyv) => Grey,
462            (Grey, PlanarRgb) => Grey,
463            (Grey, PlanarRgba) => Grey,
464            (Grey, Nv16) => Grey,
465            (Nv12, Bgra) => Rgba,
466            (Yuyv, Bgra) => Rgba,
467            (Vyuy, Bgra) => Rgba,
468            (Rgba, Bgra) => Rgba,
469            (Rgb, Bgra) => Rgb,
470            (Grey, Bgra) => Grey,
471            (Bgra, Bgra) => Bgra,
472            (Nv16, Rgb) => Rgb,
473            (Nv16, Rgba) => Rgba,
474            (Nv16, Bgra) => Rgba,
475            (PlanarRgb, Rgb) => Rgb,
476            (PlanarRgb, Rgba) => Rgb,
477            (PlanarRgb, Bgra) => Rgb,
478            (PlanarRgba, Rgb) => Rgba,
479            (PlanarRgba, Rgba) => Rgba,
480            (PlanarRgba, Bgra) => Rgba,
481            (s, d) => {
482                return Err(Error::NotSupported(format!("Conversion from {s} to {d}",)));
483            }
484        };
485
486        let need_resize_flip_rotation = rotation != Rotation::None
487            || flip != Flip::None
488            || src_w != dst_w
489            || src_h != dst_h
490            || crop.src_rect.is_some_and(|c| {
491                c != Rect {
492                    left: 0,
493                    top: 0,
494                    width: src_w,
495                    height: src_h,
496                }
497            })
498            || crop.dst_rect.is_some_and(|c| {
499                c != Rect {
500                    left: 0,
501                    top: 0,
502                    width: dst_w,
503                    height: dst_h,
504                }
505            });
506
507        // check if a direct conversion can be done
508        if !need_resize_flip_rotation && Self::support_conversion_pf(src_fmt, dst_fmt) {
509            return Self::convert_format_pf(src, dst, src_fmt, dst_fmt);
510        }
511
512        // any extra checks
513        if dst_fmt == Yuyv && !dst_w.is_multiple_of(2) {
514            return Err(Error::NotSupported(format!(
515                "{} destination must have width divisible by 2",
516                dst_fmt,
517            )));
518        }
519
520        // create tmp buffer
521        let mut tmp_buffer;
522        let tmp;
523        let tmp_fmt;
524        if intermediate != src_fmt {
525            tmp_buffer = Tensor::<u8>::image(src_w, src_h, intermediate, Some(TensorMemory::Mem))?;
526
527            Self::convert_format_pf(src, &mut tmp_buffer, src_fmt, intermediate)?;
528            tmp = &tmp_buffer;
529            tmp_fmt = intermediate;
530        } else {
531            tmp = src;
532            tmp_fmt = src_fmt;
533        }
534
535        // format must be RGB/RGBA/GREY
536        debug_assert!(matches!(tmp_fmt, Rgb | Rgba | Grey));
537        if tmp_fmt == dst_fmt {
538            self.resize_flip_rotate_pf(tmp, dst, dst_fmt, rotation, flip, crop)?;
539        } else if !need_resize_flip_rotation {
540            Self::convert_format_pf(tmp, dst, tmp_fmt, dst_fmt)?;
541        } else {
542            let mut tmp2 = Tensor::<u8>::image(dst_w, dst_h, tmp_fmt, Some(TensorMemory::Mem))?;
543            if crop.dst_rect.is_some_and(|c| {
544                c != Rect {
545                    left: 0,
546                    top: 0,
547                    width: dst_w,
548                    height: dst_h,
549                }
550            }) && crop.dst_color.is_none()
551            {
552                Self::convert_format_pf(dst, &mut tmp2, dst_fmt, tmp_fmt)?;
553            }
554            self.resize_flip_rotate_pf(tmp, &mut tmp2, tmp_fmt, rotation, flip, crop)?;
555            Self::convert_format_pf(&tmp2, dst, tmp_fmt, dst_fmt)?;
556        }
557        if let (Some(dst_rect), Some(dst_color)) = (crop.dst_rect, crop.dst_color) {
558            let full_rect = Rect {
559                left: 0,
560                top: 0,
561                width: dst_w,
562                height: dst_h,
563            };
564            if dst_rect != full_rect {
565                Self::fill_image_outside_crop_u8(dst, dst_color, dst_rect)?;
566            }
567        }
568
569        Ok(())
570    }
571
572    fn draw_decoded_masks_impl(
573        &mut self,
574        dst: &mut Tensor<u8>,
575        detect: &[DetectBox],
576        segmentation: &[Segmentation],
577        opacity: f32,
578    ) -> Result<()> {
579        let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
580        if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
581            return Err(crate::Error::NotSupported(
582                "CPU image rendering only supports RGBA or RGB images".to_string(),
583            ));
584        }
585
586        let _timer = FunctionTimer::new("CPUProcessor::draw_decoded_masks");
587
588        let dst_w = dst.width().unwrap();
589        let dst_h = dst.height().unwrap();
590        let dst_rs = row_stride_for(dst_w, dst_fmt);
591        let dst_c = dst_fmt.channels();
592
593        let mut map = dst.map()?;
594        let dst_slice = map.as_mut_slice();
595
596        self.render_box(dst_w, dst_h, dst_rs, dst_c, dst_slice, detect)?;
597
598        if segmentation.is_empty() {
599            return Ok(());
600        }
601
602        // Semantic segmentation (e.g. ModelPack) has C > 1 (multi-class),
603        // instance segmentation (e.g. YOLO) has C = 1 (binary per-instance).
604        let is_semantic = segmentation[0].segmentation.shape()[2] > 1;
605
606        if is_semantic {
607            self.render_modelpack_segmentation(
608                dst_w,
609                dst_h,
610                dst_rs,
611                dst_c,
612                dst_slice,
613                &segmentation[0],
614                opacity,
615            )?;
616        } else {
617            for (seg, detect) in segmentation.iter().zip(detect) {
618                self.render_yolo_segmentation(
619                    dst_w,
620                    dst_h,
621                    dst_rs,
622                    dst_c,
623                    dst_slice,
624                    seg,
625                    detect.label,
626                    opacity,
627                )?;
628            }
629        }
630
631        Ok(())
632    }
633
634    fn draw_proto_masks_impl(
635        &mut self,
636        dst: &mut Tensor<u8>,
637        detect: &[DetectBox],
638        proto_data: &ProtoData,
639        opacity: f32,
640    ) -> Result<()> {
641        let dst_fmt = dst.format().ok_or(Error::NotAnImage)?;
642        if !matches!(dst_fmt, PixelFormat::Rgba | PixelFormat::Rgb) {
643            return Err(crate::Error::NotSupported(
644                "CPU image rendering only supports RGBA or RGB images".to_string(),
645            ));
646        }
647
648        let _timer = FunctionTimer::new("CPUProcessor::draw_proto_masks");
649
650        let dst_w = dst.width().unwrap();
651        let dst_h = dst.height().unwrap();
652        let dst_rs = row_stride_for(dst_w, dst_fmt);
653        let channels = dst_fmt.channels();
654
655        let mut map = dst.map()?;
656        let dst_slice = map.as_mut_slice();
657
658        self.render_box(dst_w, dst_h, dst_rs, channels, dst_slice, detect)?;
659
660        if detect.is_empty() || proto_data.mask_coefficients.is_empty() {
661            return Ok(());
662        }
663
664        let protos_cow = proto_data.protos.as_f32();
665        let protos = protos_cow.as_ref();
666        let proto_h = protos.shape()[0];
667        let proto_w = protos.shape()[1];
668        let num_protos = protos.shape()[2];
669
670        for (det, coeff) in detect.iter().zip(proto_data.mask_coefficients.iter()) {
671            let color = self.colors[det.label % self.colors.len()];
672            let alpha = if opacity == 1.0 {
673                color[3] as u16
674            } else {
675                (color[3] as f32 * opacity).round() as u16
676            };
677
678            // Pixel bounds of the detection in dst image space
679            let start_x = (dst_w as f32 * det.bbox.xmin).round() as usize;
680            let start_y = (dst_h as f32 * det.bbox.ymin).round() as usize;
681            let end_x = ((dst_w as f32 * det.bbox.xmax).round() as usize).min(dst_w);
682            let end_y = ((dst_h as f32 * det.bbox.ymax).round() as usize).min(dst_h);
683
684            for y in start_y..end_y {
685                for x in start_x..end_x {
686                    // Map pixel (x, y) to proto space
687                    let px = (x as f32 / dst_w as f32) * proto_w as f32 - 0.5;
688                    let py = (y as f32 / dst_h as f32) * proto_h as f32 - 0.5;
689
690                    // Bilinear interpolation + dot product
691                    let acc = bilinear_dot(protos, coeff, num_protos, px, py, proto_w, proto_h);
692
693                    // Sigmoid threshold
694                    let mask = 1.0 / (1.0 + (-acc).exp());
695                    if mask < 0.5 {
696                        continue;
697                    }
698
699                    // Alpha blend
700                    let dst_index = y * dst_rs + x * channels;
701                    for c in 0..3 {
702                        dst_slice[dst_index + c] = ((color[c] as u16 * alpha
703                            + dst_slice[dst_index + c] as u16 * (255 - alpha))
704                            / 255) as u8;
705                    }
706                }
707            }
708        }
709
710        Ok(())
711    }
712}
edgefirst_image/cpu/mod.rs

edgefirst_image/cpu/
mod.rs
