cu_sensor_payloads/
image.rs

1use alloc::vec::Vec;
2use bincode::de::Decoder;
3use bincode::error::DecodeError;
4use bincode::{Decode, Encode};
5use core::fmt::Debug;
6use cu29::prelude::{ArrayLike, CuHandle};
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#[cfg(feature = "kornia")]
15use kornia_image::allocator::ImageAllocator;
16use serde::{Deserialize, Serialize, Serializer};
17
18#[derive(Default, Debug, Encode, Decode, Clone, Copy, Serialize, Deserialize)]
19pub struct CuImageBufferFormat {
20    pub width: u32,
21    pub height: u32,
22    pub stride: u32,
23    pub pixel_format: [u8; 4],
24}
25
26impl CuImageBufferFormat {
27    pub fn byte_size(&self) -> usize {
28        self.stride as usize * self.height as usize
29    }
30}
31
32#[derive(Debug, Default, Clone, Encode)]
33pub struct CuImage<A>
34where
35    A: ArrayLike<Element = u8>,
36{
37    pub seq: u64,
38    pub format: CuImageBufferFormat,
39    pub buffer_handle: CuHandle<A>,
40}
41
42impl Decode<()> for CuImage<Vec<u8>> {
43    fn decode<D: Decoder>(decoder: &mut D) -> Result<Self, DecodeError> {
44        let seq = u64::decode(decoder)?;
45        let format = CuImageBufferFormat::decode(decoder)?;
46        let buffer = Vec::decode(decoder)?;
47        let buffer_handle = CuHandle::new_detached(buffer);
48
49        Ok(Self {
50            seq,
51            format,
52            buffer_handle,
53        })
54    }
55}
56
57impl<'de> Deserialize<'de> for CuImage<Vec<u8>> {
58    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
59    where
60        D: serde::Deserializer<'de>,
61    {
62        #[derive(Deserialize)]
63        struct CuImageWire {
64            seq: u64,
65            format: CuImageBufferFormat,
66            handle: Vec<u8>,
67        }
68
69        let wire = CuImageWire::deserialize(deserializer)?;
70        Ok(Self {
71            seq: wire.seq,
72            format: wire.format,
73            buffer_handle: CuHandle::new_detached(wire.handle),
74        })
75    }
76}
77
78impl<A> Serialize for CuImage<A>
79where
80    A: ArrayLike<Element = u8>,
81{
82    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
83    where
84        S: Serializer,
85    {
86        use serde::ser::SerializeStruct;
87        let mut struct_ = serializer.serialize_struct("CuImage", 3)?;
88        struct_.serialize_field("seq", &self.seq)?;
89        struct_.serialize_field("format", &self.format)?;
90        // Use empty Vec as placeholder for buffer_handle to keep it opaque
91        let placeholder_buffer: Vec<u8> = Vec::new();
92        struct_.serialize_field("handle", &placeholder_buffer)?;
93        struct_.end()
94    }
95}
96
97impl<A> CuImage<A>
98where
99    A: ArrayLike<Element = u8>,
100{
101    pub fn new(format: CuImageBufferFormat, buffer_handle: CuHandle<A>) -> Self {
102        assert!(
103            format.byte_size() <= buffer_handle.with_inner(|i| i.len()),
104            "Buffer size must at least match the format."
105        );
106        CuImage {
107            seq: 0,
108            format,
109            buffer_handle,
110        }
111    }
112}
113
114impl<A> CuImage<A>
115where
116    A: ArrayLike<Element = u8>,
117{
118    /// Builds an ImageBuffer from the image crate backed by the CuImage's pixel data.
119    #[cfg(feature = "image")]
120    pub fn as_image_buffer<P: Pixel>(&self) -> CuResult<ImageBuffer<P, &[P::Subpixel]>> {
121        let width = self.format.width;
122        let height = self.format.height;
123        assert_eq!(
124            width, self.format.stride,
125            "STRIDE must equal WIDTH for ImageBuffer compatibility."
126        );
127
128        let raw_pixels: &[P::Subpixel] = self.buffer_handle.with_inner(|inner| unsafe {
129            let data: &[u8] = inner;
130            core::slice::from_raw_parts(data.as_ptr() as *const P::Subpixel, data.len())
131        });
132        ImageBuffer::from_raw(width, height, raw_pixels)
133            .ok_or("Could not create the image:: buffer".into())
134    }
135
136    #[cfg(feature = "kornia")]
137    pub fn as_kornia_image<T: Clone, const C: usize, K: ImageAllocator>(
138        &self,
139        k: K,
140    ) -> CuResult<Image<T, C, K>> {
141        let width = self.format.width as usize;
142        let height = self.format.height as usize;
143
144        assert_eq!(
145            width, self.format.stride as usize,
146            "stride must equal width for Kornia compatibility."
147        );
148
149        let size = width * height * C;
150        let raw_pixels: &[T] = self.buffer_handle.with_inner(|inner| unsafe {
151            let data: &[u8] = inner;
152            core::slice::from_raw_parts(
153                data.as_ptr() as *const T,
154                data.len() / core::mem::size_of::<T>(),
155            )
156        });
157
158        unsafe { Image::from_raw_parts([height, width].into(), raw_pixels.as_ptr(), size, k) }
159            .map_err(|e| CuError::new_with_cause("Could not create a Kornia Image", e))
160    }
161}