cu_sensor_payloads/
image.rs

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