1use bincode::de::Decoder;
2use bincode::error::DecodeError;
3use bincode::{Decode, Encode};
4use core::fmt::Debug;
5use core::ops::Range;
6use cu29::prelude::*;
7
8#[cfg(feature = "image")]
9use image::{ImageBuffer, Pixel};
10#[cfg(feature = "kornia")]
11use kornia_image::Image;
12#[cfg(feature = "kornia")]
13use kornia_image::allocator::ImageAllocator;
14use serde::{Deserialize, Serialize, Serializer};
15
16#[derive(Default, Debug, Encode, Decode, Clone, Copy, Serialize, Deserialize, Reflect)]
17pub struct CuImageBufferFormat {
18 pub width: u32,
19 pub height: u32,
20 pub stride: u32,
21 pub pixel_format: [u8; 4],
22}
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub struct CuImagePlaneLayout {
26 pub offset_bytes: usize,
27 pub row_bytes: u32,
28 pub stride_bytes: u32,
29 pub height: u32,
30}
31
32impl CuImagePlaneLayout {
33 pub fn byte_len(&self) -> usize {
34 self.stride_bytes as usize * self.height as usize
35 }
36
37 pub fn byte_range(&self) -> Range<usize> {
38 self.offset_bytes..self.offset_bytes + self.byte_len()
39 }
40}
41
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43enum CuImageMemoryLayout {
44 Packed { bytes_per_pixel: u32 },
45 SemiPlanar420,
46 Planar420,
47 SinglePlane,
48}
49
50impl CuImageBufferFormat {
51 fn memory_layout(&self) -> CuImageMemoryLayout {
52 match &self.pixel_format {
53 b"GRAY" | b"Y800" => CuImageMemoryLayout::Packed { bytes_per_pixel: 1 },
54 b"YUYV" | b"UYVY" => CuImageMemoryLayout::Packed { bytes_per_pixel: 2 },
55 b"RGB3" | b"BGR3" | b"RGB " | b"BGR " => {
56 CuImageMemoryLayout::Packed { bytes_per_pixel: 3 }
57 }
58 b"RGBA" | b"BGRA" => CuImageMemoryLayout::Packed { bytes_per_pixel: 4 },
59 b"NV12" | b"NV21" => CuImageMemoryLayout::SemiPlanar420,
60 b"I420" | b"YV12" => CuImageMemoryLayout::Planar420,
61 _ => CuImageMemoryLayout::SinglePlane,
62 }
63 }
64
65 pub fn plane_count(&self) -> usize {
66 match self.memory_layout() {
67 CuImageMemoryLayout::Planar420 => 3,
68 CuImageMemoryLayout::SemiPlanar420 => 2,
69 CuImageMemoryLayout::Packed { .. } | CuImageMemoryLayout::SinglePlane => 1,
70 }
71 }
72
73 pub fn is_packed(&self) -> bool {
74 matches!(self.memory_layout(), CuImageMemoryLayout::Packed { .. })
75 }
76
77 pub fn packed_row_bytes(&self) -> Option<u32> {
78 match self.memory_layout() {
79 CuImageMemoryLayout::Packed { bytes_per_pixel } => Some(self.width * bytes_per_pixel),
80 CuImageMemoryLayout::SemiPlanar420
81 | CuImageMemoryLayout::Planar420
82 | CuImageMemoryLayout::SinglePlane => None,
83 }
84 }
85
86 pub fn plane(&self, index: usize) -> Option<CuImagePlaneLayout> {
87 let y_plane_bytes = self.stride as usize * self.height as usize;
88 let chroma_height = self.height.div_ceil(2);
89
90 match self.memory_layout() {
91 CuImageMemoryLayout::Packed { bytes_per_pixel } if index == 0 => {
92 Some(CuImagePlaneLayout {
93 offset_bytes: 0,
94 row_bytes: self.width * bytes_per_pixel,
95 stride_bytes: self.stride,
96 height: self.height,
97 })
98 }
99 CuImageMemoryLayout::SinglePlane if index == 0 => Some(CuImagePlaneLayout {
100 offset_bytes: 0,
101 row_bytes: self.stride,
102 stride_bytes: self.stride,
103 height: self.height,
104 }),
105 CuImageMemoryLayout::SemiPlanar420 if index == 0 => Some(CuImagePlaneLayout {
106 offset_bytes: 0,
107 row_bytes: self.width,
108 stride_bytes: self.stride,
109 height: self.height,
110 }),
111 CuImageMemoryLayout::SemiPlanar420 if index == 1 => Some(CuImagePlaneLayout {
112 offset_bytes: y_plane_bytes,
113 row_bytes: self.width.div_ceil(2) * 2,
114 stride_bytes: self.stride,
115 height: chroma_height,
116 }),
117 CuImageMemoryLayout::Planar420 if index == 0 => Some(CuImagePlaneLayout {
118 offset_bytes: 0,
119 row_bytes: self.width,
120 stride_bytes: self.stride,
121 height: self.height,
122 }),
123 CuImageMemoryLayout::Planar420 if index == 1 => Some({
124 let chroma_stride = self.stride.div_ceil(2);
125 CuImagePlaneLayout {
126 offset_bytes: y_plane_bytes,
127 row_bytes: self.width.div_ceil(2),
128 stride_bytes: chroma_stride,
129 height: chroma_height,
130 }
131 }),
132 CuImageMemoryLayout::Planar420 if index == 2 => Some({
133 let chroma_stride = self.stride.div_ceil(2);
134 let chroma_plane_bytes = chroma_stride as usize * chroma_height as usize;
135 CuImagePlaneLayout {
136 offset_bytes: y_plane_bytes + chroma_plane_bytes,
137 row_bytes: self.width.div_ceil(2),
138 stride_bytes: chroma_stride,
139 height: chroma_height,
140 }
141 }),
142 _ => None,
143 }
144 }
145
146 pub fn is_valid(&self) -> bool {
147 (0..self.plane_count()).all(|index| {
148 self.plane(index)
149 .map(|plane| plane.row_bytes <= plane.stride_bytes)
150 .unwrap_or(false)
151 })
152 }
153
154 pub fn required_bytes(&self) -> usize {
155 self.plane(self.plane_count().saturating_sub(1))
156 .map(|plane| plane.offset_bytes + plane.byte_len())
157 .unwrap_or(0)
158 }
159
160 pub fn byte_size(&self) -> usize {
161 self.required_bytes()
162 }
163}
164
165#[derive(Debug, Default, Clone, Encode, Reflect)]
166#[reflect(from_reflect = false, no_field_bounds, type_path = false)]
167pub struct CuImage<A>
168where
169 A: ArrayLike<Element = u8> + Send + Sync + 'static,
170{
171 pub seq: u64,
172 pub format: CuImageBufferFormat,
173 #[reflect(ignore)]
174 pub buffer_handle: CuHandle<A>,
175}
176
177impl<A> TypePath for CuImage<A>
178where
179 A: ArrayLike<Element = u8> + Send + Sync + 'static,
180{
181 fn type_path() -> &'static str {
182 "cu_sensor_payloads::CuImage"
183 }
184
185 fn short_type_path() -> &'static str {
186 "CuImage"
187 }
188
189 fn type_ident() -> Option<&'static str> {
190 Some("CuImage")
191 }
192
193 fn crate_name() -> Option<&'static str> {
194 Some("cu_sensor_payloads")
195 }
196
197 fn module_path() -> Option<&'static str> {
198 Some("cu_sensor_payloads")
199 }
200}
201
202impl<A> Decode<()> for CuImage<A>
203where
204 A: ArrayLike<Element = u8> + Send + Sync + 'static,
205 CuHandle<A>: Decode<()>,
206{
207 fn decode<D: Decoder<Context = ()>>(decoder: &mut D) -> Result<Self, DecodeError> {
208 let seq: u64 = Decode::decode(decoder)?;
209 let format: CuImageBufferFormat = Decode::decode(decoder)?;
210 let buffer_handle: CuHandle<A> = Decode::decode(decoder)?;
211
212 Ok(Self {
213 seq,
214 format,
215 buffer_handle,
216 })
217 }
218}
219
220impl<'de, A> Deserialize<'de> for CuImage<A>
221where
222 A: ArrayLike<Element = u8> + Send + Sync + 'static,
223 CuHandle<A>: Deserialize<'de>,
224{
225 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
226 where
227 D: serde::Deserializer<'de>,
228 {
229 #[derive(Deserialize)]
230 struct CuImageWire<H> {
231 seq: u64,
232 format: CuImageBufferFormat,
233 handle: H,
234 }
235
236 let wire = CuImageWire::<CuHandle<A>>::deserialize(deserializer)?;
237 Ok(Self {
238 seq: wire.seq,
239 format: wire.format,
240 buffer_handle: wire.handle,
241 })
242 }
243}
244
245impl<A> Serialize for CuImage<A>
246where
247 A: ArrayLike<Element = u8> + Send + Sync + 'static,
248 CuHandle<A>: Serialize,
249{
250 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
251 where
252 S: Serializer,
253 {
254 use serde::ser::SerializeStruct;
255 let mut struct_ = serializer.serialize_struct("CuImage", 3)?;
256 struct_.serialize_field("seq", &self.seq)?;
257 struct_.serialize_field("format", &self.format)?;
258 struct_.serialize_field("handle", &self.buffer_handle)?;
259 struct_.end()
260 }
261}
262
263impl<A> CuImage<A>
264where
265 A: ArrayLike<Element = u8> + Send + Sync + 'static,
266{
267 pub fn new(format: CuImageBufferFormat, buffer_handle: CuHandle<A>) -> Self {
268 assert!(
269 format.is_valid(),
270 "Image format layout is invalid for the declared stride."
271 );
272 assert!(
273 format.required_bytes() <= buffer_handle.with_inner(|i| i.len()),
274 "Buffer size must at least match the format."
275 );
276 CuImage {
277 seq: 0,
278 format,
279 buffer_handle,
280 }
281 }
282}
283
284impl<A> CuImage<A>
285where
286 A: ArrayLike<Element = u8> + Send + Sync + 'static,
287{
288 pub fn with_plane_bytes<R>(
289 &self,
290 plane_index: usize,
291 f: impl FnOnce(&[u8], CuImagePlaneLayout) -> R,
292 ) -> CuResult<R> {
293 let plane = self
294 .format
295 .plane(plane_index)
296 .ok_or_else(|| CuError::from(format!("Invalid image plane index {plane_index}")))?;
297 Ok(self.buffer_handle.with_inner(|inner| {
298 let range = plane.byte_range();
299 f(&inner[range], plane)
300 }))
301 }
302
303 pub fn with_plane_bytes_mut<R>(
304 &mut self,
305 plane_index: usize,
306 f: impl FnOnce(&mut [u8], CuImagePlaneLayout) -> R,
307 ) -> CuResult<R> {
308 let plane = self
309 .format
310 .plane(plane_index)
311 .ok_or_else(|| CuError::from(format!("Invalid image plane index {plane_index}")))?;
312 Ok(self.buffer_handle.with_inner_mut(|inner| {
313 let range = plane.byte_range();
314 f(&mut inner[range], plane)
315 }))
316 }
317
318 #[cfg(feature = "image")]
320 pub fn as_image_buffer<P: Pixel>(&self) -> CuResult<ImageBuffer<P, &[P::Subpixel]>> {
321 let width = self.format.width;
322 let height = self.format.height;
323 let plane = self
324 .format
325 .plane(0)
326 .ok_or_else(|| CuError::from("Image format has no addressable planes"))?;
327 if self.format.plane_count() != 1 {
328 return Err(CuError::from(
329 "ImageBuffer compatibility requires a single-plane packed image.",
330 ));
331 }
332 if plane.row_bytes != plane.stride_bytes {
333 return Err(CuError::from(
334 "ImageBuffer compatibility requires tightly packed rows without padding.",
335 ));
336 }
337
338 self.with_plane_bytes(0, |data, _| {
339 let raw_pixels: &[P::Subpixel] = unsafe {
340 core::slice::from_raw_parts(
341 data.as_ptr() as *const P::Subpixel,
342 data.len() / core::mem::size_of::<P::Subpixel>(),
343 )
344 };
345 ImageBuffer::from_raw(width, height, raw_pixels)
346 .ok_or("Could not create the image:: buffer".into())
347 })?
348 }
349
350 #[cfg(feature = "kornia")]
351 pub fn as_kornia_image<T: Clone, const C: usize, K: ImageAllocator>(
352 &self,
353 k: K,
354 ) -> CuResult<Image<T, C, K>> {
355 let width = self.format.width as usize;
356 let height = self.format.height as usize;
357 let plane = self
358 .format
359 .plane(0)
360 .ok_or_else(|| CuError::from("Image format has no addressable planes"))?;
361 if self.format.plane_count() != 1 {
362 return Err(CuError::from(
363 "Kornia compatibility requires a single-plane packed image.",
364 ));
365 }
366 if plane.row_bytes != plane.stride_bytes {
367 return Err(CuError::from(
368 "Kornia compatibility requires tightly packed rows without padding.",
369 ));
370 }
371
372 let size = width * height * C;
373 self.with_plane_bytes(0, |data, _| {
374 let raw_pixels: &[T] = unsafe {
375 core::slice::from_raw_parts(
376 data.as_ptr() as *const T,
377 data.len() / core::mem::size_of::<T>(),
378 )
379 };
380
381 unsafe { Image::from_raw_parts([height, width].into(), raw_pixels.as_ptr(), size, k) }
382 .map_err(|e| CuError::new_with_cause("Could not create a Kornia Image", e))
383 })?
384 }
385}
386
387#[cfg(test)]
388mod tests {
389 use super::{CuImageBufferFormat, CuImagePlaneLayout};
390
391 fn assert_plane(
392 plane: Option<CuImagePlaneLayout>,
393 offset_bytes: usize,
394 row_bytes: u32,
395 stride_bytes: u32,
396 height: u32,
397 ) {
398 assert_eq!(
399 plane,
400 Some(CuImagePlaneLayout {
401 offset_bytes,
402 row_bytes,
403 stride_bytes,
404 height,
405 })
406 );
407 }
408
409 #[test]
410 fn packed_rgb3_layout_uses_single_plane() {
411 let format = CuImageBufferFormat {
412 width: 4,
413 height: 2,
414 stride: 12,
415 pixel_format: *b"RGB3",
416 };
417
418 assert!(format.is_packed());
419 assert_eq!(format.plane_count(), 1);
420 assert_eq!(format.packed_row_bytes(), Some(12));
421 assert_plane(format.plane(0), 0, 12, 12, 2);
422 assert_eq!(format.required_bytes(), 24);
423 assert!(format.is_valid());
424 }
425
426 #[test]
427 fn nv12_layout_exposes_two_planes() {
428 let format = CuImageBufferFormat {
429 width: 640,
430 height: 360,
431 stride: 640,
432 pixel_format: *b"NV12",
433 };
434
435 assert!(!format.is_packed());
436 assert_eq!(format.plane_count(), 2);
437 assert_plane(format.plane(0), 0, 640, 640, 360);
438 assert_plane(format.plane(1), 230_400, 640, 640, 180);
439 assert_eq!(format.required_bytes(), 345_600);
440 assert!(format.is_valid());
441 }
442
443 #[test]
444 fn i420_layout_exposes_three_planes() {
445 let format = CuImageBufferFormat {
446 width: 640,
447 height: 360,
448 stride: 640,
449 pixel_format: *b"I420",
450 };
451
452 assert_eq!(format.plane_count(), 3);
453 assert_plane(format.plane(0), 0, 640, 640, 360);
454 assert_plane(format.plane(1), 230_400, 320, 320, 180);
455 assert_plane(format.plane(2), 288_000, 320, 320, 180);
456 assert_eq!(format.required_bytes(), 345_600);
457 assert!(format.is_valid());
458 }
459
460 #[test]
461 fn invalid_stride_is_detected_for_packed_formats() {
462 let format = CuImageBufferFormat {
463 width: 4,
464 height: 2,
465 stride: 4,
466 pixel_format: *b"RGB3",
467 };
468
469 assert!(!format.is_valid());
470 assert_eq!(format.packed_row_bytes(), Some(12));
471 }
472
473 #[test]
474 fn byte_size_for_packed_formats_is_stride_times_height() {
475 let format = CuImageBufferFormat {
476 width: 4,
477 height: 3,
478 stride: 16,
479 pixel_format: *b"BGRA",
480 };
481
482 assert_eq!(format.byte_size(), 48);
483 }
484
485 #[test]
486 fn byte_size_for_nv12_includes_uv_plane() {
487 let format = CuImageBufferFormat {
488 width: 1280,
489 height: 720,
490 stride: 1280,
491 pixel_format: *b"NV12",
492 };
493
494 assert_eq!(format.byte_size(), 1_382_400);
495 }
496
497 #[test]
498 fn byte_size_for_i420_includes_chroma_planes() {
499 let format = CuImageBufferFormat {
500 width: 640,
501 height: 480,
502 stride: 640,
503 pixel_format: *b"I420",
504 };
505
506 assert_eq!(format.byte_size(), 460_800);
507 }
508}