cu_sensor_payloads/
image.rs

1use bincode::de::Decoder;
2use bincode::error::DecodeError;
3use bincode::{Decode, Encode};
4use cu29::prelude::{ArrayLike, CuHandle};
5use std::fmt::Debug;
6
7#[allow(unused_imports)]
8use cu29::{CuError, CuResult};
9
10#[cfg(feature = "image")]
11use image::{ImageBuffer, Pixel};
12#[cfg(feature = "kornia")]
13use kornia::image::Image;
14
15#[derive(Default, Debug, Encode, Decode, Clone, Copy)]
16pub struct CuImageBufferFormat {
17    pub width: u32,
18    pub height: u32,
19    pub stride: u32,
20    pub pixel_format: [u8; 4],
21}
22
23impl CuImageBufferFormat {
24    pub fn byte_size(&self) -> usize {
25        self.stride as usize * self.height as usize
26    }
27}
28
29#[derive(Debug, Default, Clone, Encode)]
30pub struct CuImage<A>
31where
32    A: ArrayLike<Element = u8>,
33{
34    pub seq: u64,
35    pub format: CuImageBufferFormat,
36    pub buffer_handle: CuHandle<A>,
37}
38
39impl Decode for CuImage<Vec<u8>> {
40    fn decode<D: Decoder>(decoder: &mut D) -> Result<Self, DecodeError> {
41        let seq = u64::decode(decoder)?;
42        let format = CuImageBufferFormat::decode(decoder)?;
43        let buffer = Vec::decode(decoder)?;
44        let buffer_handle = CuHandle::new_detached(buffer);
45
46        Ok(Self {
47            seq,
48            format,
49            buffer_handle,
50        })
51    }
52}
53
54impl<A> CuImage<A>
55where
56    A: ArrayLike<Element = u8>,
57{
58    pub fn new(format: CuImageBufferFormat, buffer_handle: CuHandle<A>) -> Self {
59        assert!(
60            format.byte_size() < buffer_handle.with_inner(|i| i.len()),
61            "Buffer size must at least match the format."
62        );
63        CuImage {
64            seq: 0,
65            format,
66            buffer_handle,
67        }
68    }
69}
70
71impl<A> CuImage<A>
72where
73    A: ArrayLike<Element = u8>,
74{
75    /// Builds an ImageBuffer from the image crate backed by the CuImage's pixel data.
76    #[cfg(feature = "image")]
77    pub fn as_image_buffer<P: Pixel>(&self) -> CuResult<ImageBuffer<P, &[P::Subpixel]>> {
78        let width = self.format.width;
79        let height = self.format.height;
80        assert_eq!(
81            width, self.format.stride,
82            "STRIDE must equal WIDTH for ImageBuffer compatibility."
83        );
84
85        let raw_pixels: &[P::Subpixel] = self.buffer_handle.with_inner(|inner| unsafe {
86            let data: &[u8] = inner;
87            core::slice::from_raw_parts(data.as_ptr() as *const P::Subpixel, data.len())
88        });
89        ImageBuffer::from_raw(width, height, raw_pixels)
90            .ok_or("Could not create the image:: buffer".into())
91    }
92
93    #[cfg(feature = "kornia")]
94    pub fn as_kornia_image<T: Clone, const C: usize>(&self) -> CuResult<Image<T, C>> {
95        let width = self.format.width as usize;
96        let height = self.format.height as usize;
97
98        assert_eq!(
99            width, self.format.stride as usize,
100            "stride must equal width for Kornia compatibility."
101        );
102
103        let size = width * height * C;
104        let raw_pixels: &[T] = self.buffer_handle.with_inner(|inner| unsafe {
105            let data: &[u8] = inner;
106            core::slice::from_raw_parts(
107                data.as_ptr() as *const T,
108                data.len() / std::mem::size_of::<T>(),
109            )
110        });
111
112        unsafe { Image::from_raw_parts([height, width].into(), raw_pixels.as_ptr(), size) }
113            .map_err(|e| CuError::new_with_cause("Could not create a Kornia Image", e))
114    }
115}