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;
13use serde::{Serialize, Serializer};
14
15#[derive(Default, Debug, Encode, Decode, Clone, Copy, Serialize)]
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> Serialize for CuImage<A>
55where
56    A: ArrayLike<Element = u8>,
57{
58    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
59    where
60        S: Serializer,
61    {
62        use serde::ser::SerializeStruct;
63        let mut struct_ = serializer.serialize_struct("CuImage", 3)?;
64        struct_.serialize_field("seq", &self.seq)?;
65        struct_.serialize_field("format", &self.format)?;
66        // Use empty Vec as placeholder for buffer_handle to keep it opaque
67        let placeholder_buffer: Vec<u8> = Vec::new();
68        struct_.serialize_field("handle", &placeholder_buffer)?;
69        struct_.end()
70    }
71}
72
73impl<A> CuImage<A>
74where
75    A: ArrayLike<Element = u8>,
76{
77    pub fn new(format: CuImageBufferFormat, buffer_handle: CuHandle<A>) -> Self {
78        assert!(
79            format.byte_size() <= buffer_handle.with_inner(|i| i.len()),
80            "Buffer size must at least match the format."
81        );
82        CuImage {
83            seq: 0,
84            format,
85            buffer_handle,
86        }
87    }
88}
89
90impl<A> CuImage<A>
91where
92    A: ArrayLike<Element = u8>,
93{
94    /// Builds an ImageBuffer from the image crate backed by the CuImage's pixel data.
95    #[cfg(feature = "image")]
96    pub fn as_image_buffer<P: Pixel>(&self) -> CuResult<ImageBuffer<P, &[P::Subpixel]>> {
97        let width = self.format.width;
98        let height = self.format.height;
99        assert_eq!(
100            width, self.format.stride,
101            "STRIDE must equal WIDTH for ImageBuffer compatibility."
102        );
103
104        let raw_pixels: &[P::Subpixel] = self.buffer_handle.with_inner(|inner| unsafe {
105            let data: &[u8] = inner;
106            core::slice::from_raw_parts(data.as_ptr() as *const P::Subpixel, data.len())
107        });
108        ImageBuffer::from_raw(width, height, raw_pixels)
109            .ok_or("Could not create the image:: buffer".into())
110    }
111
112    #[cfg(feature = "kornia")]
113    pub fn as_kornia_image<T: Clone, const C: usize>(&self) -> CuResult<Image<T, C>> {
114        let width = self.format.width as usize;
115        let height = self.format.height as usize;
116
117        assert_eq!(
118            width, self.format.stride as usize,
119            "stride must equal width for Kornia compatibility."
120        );
121
122        let size = width * height * C;
123        let raw_pixels: &[T] = self.buffer_handle.with_inner(|inner| unsafe {
124            let data: &[u8] = inner;
125            core::slice::from_raw_parts(
126                data.as_ptr() as *const T,
127                data.len() / std::mem::size_of::<T>(),
128            )
129        });
130
131        unsafe { Image::from_raw_parts([height, width].into(), raw_pixels.as_ptr(), size) }
132            .map_err(|e| CuError::new_with_cause("Could not create a Kornia Image", e))
133    }
134}