spatialrust_tensor/
image.rs

1//! Explicit bridges between typed CPU images and generic tensors.
2
3use bytemuck::{cast_slice, Pod};
4use spatialrust_image::{Image, ImageView, PlanarImage, PlanarImageView};
5
6use crate::{DataType, Device, TensorBuffer, TensorDescriptor, TensorError, TensorView};
7
8/// A native-endian scalar that has a stable tensor dtype and no invalid bit patterns.
9pub trait TensorElement: Pod {
10    /// Tensor dtype corresponding to this Rust scalar.
11    const DTYPE: DataType;
12}
13
14macro_rules! tensor_elements {
15    ($($type:ty => $dtype:expr),+ $(,)?) => {
16        $(impl TensorElement for $type {
17            const DTYPE: DataType = $dtype;
18        })+
19    };
20}
21
22tensor_elements! {
23    u8 => DataType::U8,
24    u16 => DataType::U16,
25    i8 => DataType::I8,
26    i16 => DataType::I16,
27    i32 => DataType::I32,
28    i64 => DataType::I64,
29    f32 => DataType::F32,
30    f64 => DataType::F64,
31}
32
33/// Borrows a packed interleaved image as a zero-copy `[height, width, channels]` tensor.
34pub fn interleaved_image_view<T: TensorElement, const CHANNELS: usize>(
35    image: &Image<T, CHANNELS>,
36) -> Result<TensorView<'_>, TensorError> {
37    let descriptor = TensorDescriptor::contiguous(
38        T::DTYPE,
39        vec![image.height(), image.width(), CHANNELS],
40        Device::CPU,
41    );
42    TensorView::try_new(cast_slice(image.as_slice()), descriptor)
43}
44
45/// Borrows a packed planar image as a zero-copy `[channels, height, width]` tensor.
46pub fn planar_image_view<T: TensorElement, const CHANNELS: usize>(
47    image: &PlanarImage<T, CHANNELS>,
48) -> Result<TensorView<'_>, TensorError> {
49    let descriptor = TensorDescriptor::contiguous(
50        T::DTYPE,
51        vec![CHANNELS, image.height(), image.width()],
52        Device::CPU,
53    );
54    TensorView::try_new(cast_slice(image.as_slice()), descriptor)
55}
56
57/// Explicitly packs a possibly strided interleaved image view into an owned HWC tensor.
58pub fn pack_interleaved_image<T: TensorElement, const CHANNELS: usize>(
59    image: ImageView<'_, T, CHANNELS>,
60) -> Result<TensorBuffer, TensorError> {
61    let scalar_count = image
62        .width()
63        .checked_mul(image.height())
64        .and_then(|value| value.checked_mul(CHANNELS))
65        .ok_or(TensorError::LayoutOverflow)?;
66    let byte_count =
67        scalar_count.checked_mul(T::DTYPE.element_size()).ok_or(TensorError::LayoutOverflow)?;
68    let mut bytes = Vec::with_capacity(byte_count);
69    for y in 0..image.height() {
70        let row = image.row(y).expect("row index is within the image height");
71        bytes.extend_from_slice(cast_slice(row));
72    }
73    let descriptor = TensorDescriptor::contiguous(
74        T::DTYPE,
75        vec![image.height(), image.width(), CHANNELS],
76        Device::CPU,
77    );
78    TensorBuffer::try_new(bytes, descriptor)
79}
80
81/// Explicitly packs a possibly strided planar image view into an owned CHW tensor.
82pub fn pack_planar_image<T: TensorElement, const CHANNELS: usize>(
83    image: PlanarImageView<'_, T, CHANNELS>,
84) -> Result<TensorBuffer, TensorError> {
85    let scalar_count = image
86        .width()
87        .checked_mul(image.height())
88        .and_then(|value| value.checked_mul(CHANNELS))
89        .ok_or(TensorError::LayoutOverflow)?;
90    let byte_count =
91        scalar_count.checked_mul(T::DTYPE.element_size()).ok_or(TensorError::LayoutOverflow)?;
92    let mut bytes = Vec::with_capacity(byte_count);
93    for channel in 0..CHANNELS {
94        for y in 0..image.height() {
95            for x in 0..image.width() {
96                let value = image
97                    .get(channel, x, y)
98                    .expect("channel and pixel coordinates are within the image");
99                bytes.extend_from_slice(bytemuck::bytes_of(value));
100            }
101        }
102    }
103    let descriptor = TensorDescriptor::contiguous(
104        T::DTYPE,
105        vec![CHANNELS, image.height(), image.width()],
106        Device::CPU,
107    );
108    TensorBuffer::try_new(bytes, descriptor)
109}
110
111#[cfg(test)]
112mod tests {
113    use super::{
114        interleaved_image_view, pack_interleaved_image, pack_planar_image, planar_image_view,
115    };
116    use spatialrust_image::{Image, ImageRegion, ImageView, PlanarImage, PlanarImageView};
117
118    #[test]
119    fn packed_images_are_zero_copy_hwc_and_chw() {
120        let interleaved = Image::<u16, 3>::try_new(2, 2, (0..12).collect()).unwrap();
121        let tensor = interleaved_image_view(&interleaved).unwrap();
122        assert_eq!(tensor.descriptor().shape(), &[2, 2, 3]);
123        assert_eq!(tensor.allocation_bytes().as_ptr(), interleaved.as_slice().as_ptr().cast());
124
125        let planar =
126            PlanarImage::<f32, 3>::try_new(2, 2, (0..12).map(|v| v as f32).collect()).unwrap();
127        let tensor = planar_image_view(&planar).unwrap();
128        assert_eq!(tensor.descriptor().shape(), &[3, 2, 2]);
129        assert_eq!(tensor.allocation_bytes().as_ptr(), planar.as_slice().as_ptr().cast());
130    }
131
132    #[test]
133    fn strided_interleaved_roi_is_explicitly_packed() {
134        let storage = (0_u8..30).collect::<Vec<_>>();
135        let parent = ImageView::<u8, 3>::new(3, 3, 10, &storage).unwrap();
136        let roi = parent.subview(ImageRegion::new(1, 1, 2, 2)).unwrap();
137        let tensor = pack_interleaved_image(roi).unwrap();
138        assert_eq!(tensor.descriptor().shape(), &[2, 2, 3]);
139        assert_eq!(tensor.allocation_bytes(), &[13, 14, 15, 16, 17, 18, 23, 24, 25, 26, 27, 28]);
140    }
141
142    #[test]
143    fn padded_planar_view_is_explicitly_packed() {
144        let storage = (0_u16..24).collect::<Vec<_>>();
145        let view = PlanarImageView::<u16, 2>::new(2, 2, 3, 12, &storage).unwrap();
146        let tensor = pack_planar_image(view).unwrap();
147        assert_eq!(tensor.descriptor().shape(), &[2, 2, 2]);
148        assert_eq!(
149            bytemuck::cast_slice::<u8, u16>(tensor.allocation_bytes()),
150            &[0, 1, 3, 4, 12, 13, 15, 16]
151        );
152    }
153}
spatialrust_tensor/image.rs

spatialrust_tensor/
image.rs
