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, 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> Clone for CuImage<A>
180where
181 A: ArrayLike<Element = u8> + Send + Sync + 'static,
182{
183 fn clone(&self) -> Self {
184 Self {
185 seq: self.seq,
186 format: self.format,
187 buffer_handle: self.buffer_handle.clone(),
188 }
189 }
190}
191
192impl<A> Default for CuImage<A>
193where
194 A: ArrayLike<Element = u8> + Send + Sync + 'static,
195{
196 fn default() -> Self {
197 Self {
198 seq: 0,
199 format: CuImageBufferFormat::default(),
200 buffer_handle: CuHandle::default(),
201 }
202 }
203}
204
205impl<A> TypePath for CuImage<A>
206where
207 A: ArrayLike<Element = u8> + Send + Sync + 'static,
208{
209 fn type_path() -> &'static str {
210 "cu_sensor_payloads::CuImage"
211 }
212
213 fn short_type_path() -> &'static str {
214 "CuImage"
215 }
216
217 fn type_ident() -> Option<&'static str> {
218 Some("CuImage")
219 }
220
221 fn crate_name() -> Option<&'static str> {
222 Some("cu_sensor_payloads")
223 }
224
225 fn module_path() -> Option<&'static str> {
226 Some("cu_sensor_payloads")
227 }
228}
229
230impl<A> Decode<()> for CuImage<A>
231where
232 A: ArrayLike<Element = u8> + Send + Sync + 'static,
233 CuHandle<A>: Decode<()>,
234{
235 fn decode<D: Decoder<Context = ()>>(decoder: &mut D) -> Result<Self, DecodeError> {
236 let seq: u64 = Decode::decode(decoder)?;
237 let format: CuImageBufferFormat = Decode::decode(decoder)?;
238 let buffer_handle: CuHandle<A> = Decode::decode(decoder)?;
239
240 Ok(Self {
241 seq,
242 format,
243 buffer_handle,
244 })
245 }
246}
247
248impl<'de, A> Deserialize<'de> for CuImage<A>
249where
250 A: ArrayLike<Element = u8> + Send + Sync + 'static,
251 CuHandle<A>: Deserialize<'de>,
252{
253 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
254 where
255 D: serde::Deserializer<'de>,
256 {
257 #[derive(Deserialize)]
258 struct CuImageWire<H> {
259 seq: u64,
260 format: CuImageBufferFormat,
261 handle: H,
262 }
263
264 let wire = CuImageWire::<CuHandle<A>>::deserialize(deserializer)?;
265 Ok(Self {
266 seq: wire.seq,
267 format: wire.format,
268 buffer_handle: wire.handle,
269 })
270 }
271}
272
273impl<A> Serialize for CuImage<A>
274where
275 A: ArrayLike<Element = u8> + Send + Sync + 'static,
276 CuHandle<A>: Serialize,
277{
278 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
279 where
280 S: Serializer,
281 {
282 use serde::ser::SerializeStruct;
283 let mut struct_ = serializer.serialize_struct("CuImage", 3)?;
284 struct_.serialize_field("seq", &self.seq)?;
285 struct_.serialize_field("format", &self.format)?;
286 struct_.serialize_field("handle", &self.buffer_handle)?;
287 struct_.end()
288 }
289}
290
291impl<A> CuImage<A>
292where
293 A: ArrayLike<Element = u8> + Send + Sync + 'static,
294{
295 pub fn new(format: CuImageBufferFormat, buffer_handle: CuHandle<A>) -> Self {
296 assert!(
297 format.is_valid(),
298 "Image format layout is invalid for the declared stride."
299 );
300 assert!(
301 format.required_bytes() <= buffer_handle.with_inner(|i| i.len()),
302 "Buffer size must at least match the format."
303 );
304 CuImage {
305 seq: 0,
306 format,
307 buffer_handle,
308 }
309 }
310}
311
312impl<A> cu29::pool::HandleContentAware for CuImage<A> where
314 A: ArrayLike<Element = u8> + Send + Sync + 'static
315{
316}
317
318impl<A> CuImage<A>
319where
320 A: ArrayLike<Element = u8> + Send + Sync + 'static,
321{
322 pub fn payload_should_log(&self) -> bool {
327 self.buffer_handle.payload_should_log()
328 }
329
330 pub fn apply_handle_content_policy(&self, mode: cu29::pool::HandleContent) {
335 self.buffer_handle.apply_handle_content_policy(mode);
336 }
337
338 pub fn mark_touched(&self) {
342 self.buffer_handle.mark_touched();
343 }
344
345 pub fn with_plane_bytes<R>(
346 &self,
347 plane_index: usize,
348 f: impl FnOnce(&[u8], CuImagePlaneLayout) -> R,
349 ) -> CuResult<R> {
350 let plane = self
351 .format
352 .plane(plane_index)
353 .ok_or_else(|| CuError::from(format!("Invalid image plane index {plane_index}")))?;
354 Ok(self.buffer_handle.with_inner(|inner| {
355 let range = plane.byte_range();
356 f(&inner[range], plane)
357 }))
358 }
359
360 pub fn with_plane_bytes_mut<R>(
361 &mut self,
362 plane_index: usize,
363 f: impl FnOnce(&mut [u8], CuImagePlaneLayout) -> R,
364 ) -> CuResult<R> {
365 let plane = self
366 .format
367 .plane(plane_index)
368 .ok_or_else(|| CuError::from(format!("Invalid image plane index {plane_index}")))?;
369 Ok(self.buffer_handle.with_inner_mut(|inner| {
370 let range = plane.byte_range();
371 f(&mut inner[range], plane)
372 }))
373 }
374
375 #[cfg(feature = "image")]
377 pub fn as_image_buffer<P: Pixel>(&self) -> CuResult<ImageBuffer<P, &[P::Subpixel]>> {
378 let width = self.format.width;
379 let height = self.format.height;
380 let plane = self
381 .format
382 .plane(0)
383 .ok_or_else(|| CuError::from("Image format has no addressable planes"))?;
384 if self.format.plane_count() != 1 {
385 return Err(CuError::from(
386 "ImageBuffer compatibility requires a single-plane packed image.",
387 ));
388 }
389 if plane.row_bytes != plane.stride_bytes {
390 return Err(CuError::from(
391 "ImageBuffer compatibility requires tightly packed rows without padding.",
392 ));
393 }
394
395 self.with_plane_bytes(0, |data, _| {
396 let raw_pixels: &[P::Subpixel] = unsafe {
397 core::slice::from_raw_parts(
398 data.as_ptr() as *const P::Subpixel,
399 data.len() / core::mem::size_of::<P::Subpixel>(),
400 )
401 };
402 ImageBuffer::from_raw(width, height, raw_pixels)
403 .ok_or("Could not create the image:: buffer".into())
404 })?
405 }
406
407 #[cfg(feature = "kornia")]
408 pub fn as_kornia_image<T: Clone, const C: usize, K: ImageAllocator>(
409 &self,
410 k: K,
411 ) -> CuResult<Image<T, C, K>> {
412 let width = self.format.width as usize;
413 let height = self.format.height as usize;
414 let plane = self
415 .format
416 .plane(0)
417 .ok_or_else(|| CuError::from("Image format has no addressable planes"))?;
418 if self.format.plane_count() != 1 {
419 return Err(CuError::from(
420 "Kornia compatibility requires a single-plane packed image.",
421 ));
422 }
423 if plane.row_bytes != plane.stride_bytes {
424 return Err(CuError::from(
425 "Kornia compatibility requires tightly packed rows without padding.",
426 ));
427 }
428
429 let size = width * height * C;
430 self.with_plane_bytes(0, |data, _| {
431 let raw_pixels: &[T] = unsafe {
432 core::slice::from_raw_parts(
433 data.as_ptr() as *const T,
434 data.len() / core::mem::size_of::<T>(),
435 )
436 };
437
438 unsafe { Image::from_raw_parts([height, width].into(), raw_pixels.as_ptr(), size, k) }
439 .map_err(|e| CuError::new_with_cause("Could not create a Kornia Image", e))
440 })?
441 }
442}
443
444#[cfg(test)]
445mod tests {
446 #[cfg(feature = "std")]
447 #[test]
448 fn test_shared_memory_image_clones_its_handle_and_satisfies_payload_contract() {
449 use super::CuImage;
450 use cu29::pool::CuPool;
451 use cu29::pool::CuSharedMemoryBuffer;
452 use cu29::pool::CuSharedMemoryPool;
453 use cu29::prelude::*;
454
455 fn assert_payload<T: CuMsgPayload>() {}
456 assert_payload::<CuImage<CuSharedMemoryBuffer<u8>>>();
457
458 let pool = CuSharedMemoryPool::new("image-clone-test", 1, 3).unwrap();
459 let handle = pool.acquire().unwrap();
460 handle.with_inner_mut(|bytes| bytes.copy_from_slice(&[12, 34, 56]));
461 let mut image = CuImage::new(
462 CuImageBufferFormat {
463 width: 1,
464 height: 1,
465 stride: 3,
466 pixel_format: *b"RGB3",
467 },
468 handle,
469 );
470 image.seq = 42;
471 let cloned = image.clone();
472 assert_eq!(cloned.seq, 42);
473 assert_eq!(cloned.format.pixel_format, *b"RGB3");
474 assert_eq!(
475 cloned.buffer_handle.storage_id(),
476 image.buffer_handle.storage_id()
477 );
478 drop(image);
479 assert!(pool.acquire().is_none());
480 assert_eq!(
481 cloned
482 .with_plane_bytes(0, |bytes, _| bytes.to_vec())
483 .unwrap(),
484 [12, 34, 56]
485 );
486
487 let config = bincode::config::standard();
488 let encoded = bincode::encode_to_vec(&cloned, config).unwrap();
489 drop(cloned);
490 assert!(pool.acquire().is_some());
491 let (decoded, used): (CuImage<CuSharedMemoryBuffer<u8>>, usize) =
492 bincode::decode_from_slice(&encoded, config).unwrap();
493 assert_eq!(used, encoded.len());
494 assert_eq!(decoded.seq, 42);
495 assert_eq!(
496 decoded
497 .with_plane_bytes(0, |bytes, _| bytes.to_vec())
498 .unwrap(),
499 [12, 34, 56]
500 );
501 }
502
503 use super::{CuImageBufferFormat, CuImagePlaneLayout};
504
505 fn assert_plane(
506 plane: Option<CuImagePlaneLayout>,
507 offset_bytes: usize,
508 row_bytes: u32,
509 stride_bytes: u32,
510 height: u32,
511 ) {
512 assert_eq!(
513 plane,
514 Some(CuImagePlaneLayout {
515 offset_bytes,
516 row_bytes,
517 stride_bytes,
518 height,
519 })
520 );
521 }
522
523 #[test]
524 fn packed_rgb3_layout_uses_single_plane() {
525 let format = CuImageBufferFormat {
526 width: 4,
527 height: 2,
528 stride: 12,
529 pixel_format: *b"RGB3",
530 };
531
532 assert!(format.is_packed());
533 assert_eq!(format.plane_count(), 1);
534 assert_eq!(format.packed_row_bytes(), Some(12));
535 assert_plane(format.plane(0), 0, 12, 12, 2);
536 assert_eq!(format.required_bytes(), 24);
537 assert!(format.is_valid());
538 }
539
540 #[test]
541 fn nv12_layout_exposes_two_planes() {
542 let format = CuImageBufferFormat {
543 width: 640,
544 height: 360,
545 stride: 640,
546 pixel_format: *b"NV12",
547 };
548
549 assert!(!format.is_packed());
550 assert_eq!(format.plane_count(), 2);
551 assert_plane(format.plane(0), 0, 640, 640, 360);
552 assert_plane(format.plane(1), 230_400, 640, 640, 180);
553 assert_eq!(format.required_bytes(), 345_600);
554 assert!(format.is_valid());
555 }
556
557 #[test]
558 fn i420_layout_exposes_three_planes() {
559 let format = CuImageBufferFormat {
560 width: 640,
561 height: 360,
562 stride: 640,
563 pixel_format: *b"I420",
564 };
565
566 assert_eq!(format.plane_count(), 3);
567 assert_plane(format.plane(0), 0, 640, 640, 360);
568 assert_plane(format.plane(1), 230_400, 320, 320, 180);
569 assert_plane(format.plane(2), 288_000, 320, 320, 180);
570 assert_eq!(format.required_bytes(), 345_600);
571 assert!(format.is_valid());
572 }
573
574 #[test]
575 fn invalid_stride_is_detected_for_packed_formats() {
576 let format = CuImageBufferFormat {
577 width: 4,
578 height: 2,
579 stride: 4,
580 pixel_format: *b"RGB3",
581 };
582
583 assert!(!format.is_valid());
584 assert_eq!(format.packed_row_bytes(), Some(12));
585 }
586
587 #[test]
588 fn byte_size_for_packed_formats_is_stride_times_height() {
589 let format = CuImageBufferFormat {
590 width: 4,
591 height: 3,
592 stride: 16,
593 pixel_format: *b"BGRA",
594 };
595
596 assert_eq!(format.byte_size(), 48);
597 }
598
599 #[test]
600 fn byte_size_for_nv12_includes_uv_plane() {
601 let format = CuImageBufferFormat {
602 width: 1280,
603 height: 720,
604 stride: 1280,
605 pixel_format: *b"NV12",
606 };
607
608 assert_eq!(format.byte_size(), 1_382_400);
609 }
610
611 #[test]
612 fn byte_size_for_i420_includes_chroma_planes() {
613 let format = CuImageBufferFormat {
614 width: 640,
615 height: 480,
616 stride: 640,
617 pixel_format: *b"I420",
618 };
619
620 assert_eq!(format.byte_size(), 460_800);
621 }
622
623 mod only_log_what_you_use {
642 use crate::CuImage;
643 use bincode::config;
644 use cu29::config::HandleContent;
645 use cu29::cutask::{CuMsg, encode_metadata_only};
646 use cu29::pool::CuHandle;
647 const FORMAT: super::super::CuImageBufferFormat = super::super::CuImageBufferFormat {
654 width: 2,
655 height: 2,
656 stride: 2,
657 pixel_format: *b"GRAY",
658 };
659
660 fn make_image() -> CuImage<Vec<u8>> {
661 let handle = CuHandle::new_detached(vec![0xDE, 0xAD, 0xBE, 0xEF]);
664 CuImage::new(FORMAT, handle)
665 }
666
667 fn encode_with_policy(msg: &CuMsg<CuImage<Vec<u8>>>, mode: HandleContent) -> Vec<u8> {
676 let should_log = match msg.payload() {
677 Some(p) => {
678 p.apply_handle_content_policy(mode);
679 p.payload_should_log()
680 }
681 None => false,
682 };
683
684 use bincode::enc::write::Writer;
686 struct VecWriter(Vec<u8>);
687 impl Writer for VecWriter {
688 fn write(&mut self, bytes: &[u8]) -> Result<(), bincode::error::EncodeError> {
689 self.0.extend_from_slice(bytes);
690 Ok(())
691 }
692 }
693 let mut encoder =
694 bincode::enc::EncoderImpl::new(VecWriter(Vec::new()), config::standard());
695 use bincode::enc::Encode as _;
696 if should_log {
697 msg.encode(&mut encoder).expect("encode");
698 } else {
699 encode_metadata_only(msg, &mut encoder).expect("encode metadata-only");
700 }
701 encoder.into_writer().0
702 }
703
704 #[test]
706 fn touched_only_untouched_skips_payload() {
707 let image = make_image();
708 let msg: CuMsg<CuImage<Vec<u8>>> = CuMsg::new(Some(image));
709 let bytes = encode_with_policy(&msg, HandleContent::TouchedOnly);
710 assert_eq!(
711 bytes.first().copied(),
712 Some(0u8),
713 "untouched TouchedOnly must emit no-payload tag"
714 );
715 }
716
717 #[test]
719 fn touched_only_touched_keeps_payload() {
720 let image = make_image();
721 image.mark_touched();
722 let msg: CuMsg<CuImage<Vec<u8>>> = CuMsg::new(Some(image));
723 let bytes = encode_with_policy(&msg, HandleContent::TouchedOnly);
724 assert_eq!(bytes.first().copied(), Some(1u8));
725 assert!(bytes.len() > 8, "payload bytes must follow the present tag");
726 }
727
728 #[test]
730 fn none_always_skips_payload() {
731 let image = make_image();
732 image.mark_touched();
733 let msg: CuMsg<CuImage<Vec<u8>>> = CuMsg::new(Some(image));
734 let bytes = encode_with_policy(&msg, HandleContent::None);
735 assert_eq!(bytes.first().copied(), Some(0u8));
736 }
737
738 #[test]
742 fn all_mode_keeps_payload_with_or_without_touch() {
743 for touched in [false, true] {
744 let image = make_image();
745 if touched {
746 image.mark_touched();
747 }
748 let msg: CuMsg<CuImage<Vec<u8>>> = CuMsg::new(Some(image));
749 let bytes = encode_with_policy(&msg, HandleContent::All);
750 assert_eq!(
751 bytes.first().copied(),
752 Some(1u8),
753 "All mode must keep payload (touched={touched})"
754 );
755 }
756 }
757
758 #[test]
761 fn touched_flag_is_shared_across_clones() {
762 let image = make_image();
763 image.apply_handle_content_policy(HandleContent::TouchedOnly);
764 let consumer_view = image.buffer_handle.clone();
765 let _ = consumer_view.with_touched_inner(|inner| inner.as_ref()[0]);
766 assert!(image.payload_should_log());
767 }
768 }
769}