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