cu_sensor_payloads/
image.rs1use 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 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 #[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}