spatialrust_vision/
adapters.rs

1//! Explicit image ↔ tensor adapters for AI pipelines.
2//!
3//! These helpers materialize contiguous host tensors and vision maps. They never
4//! call an inference backend and never imply device transfers.
5
6use spatialrust_image::{ImageView, PlanarImage};
7use spatialrust_tensor::{DataType, Device, TensorBuffer, TensorDescriptor};
8
9use crate::{
10    letterbox, pack_chw, BoundingBox2, DepthMap, Detection, Interpolation, LetterboxTransform,
11    VisionError, VisionResult,
12};
13
14/// Letterboxes an RGB `u8` image and packs it as contiguous NCHW `f32` `[1,3,H,W]`.
15pub fn rgb_u8_to_nchw_f32(
16    image: ImageView<'_, u8, 3>,
17    width: usize,
18    height: usize,
19    interpolation: Interpolation,
20    pad_value: [u8; 3],
21    scale: f32,
22    mean: [f32; 3],
23    std: [f32; 3],
24) -> VisionResult<(TensorBuffer, LetterboxTransform)> {
25    let (letterboxed, mapping) = letterbox(image, width, height, interpolation, pad_value)?;
26    let planar = pack_chw(letterboxed.view(), scale, mean, std)?;
27    Ok((planar_f32_to_nchw(&planar)?, mapping))
28}
29
30/// Packs a planar CHW `f32` image into contiguous NCHW `[1,C,H,W]`.
31pub fn planar_f32_to_nchw<const CHANNELS: usize>(
32    planar: &PlanarImage<f32, CHANNELS>,
33) -> VisionResult<TensorBuffer> {
34    let values = planar.as_slice().to_vec();
35    TensorBuffer::try_from_f32(
36        values,
37        TensorDescriptor::contiguous(
38            DataType::F32,
39            vec![1, CHANNELS, planar.height(), planar.width()],
40            Device::CPU,
41        ),
42    )
43    .map_err(|error| VisionError::InvalidParameter(error.to_string()))
44}
45
46/// Converts a host depth tensor into a [`DepthMap`].
47///
48/// Accepted shapes: `[H,W]`, `[1,H,W]`, or `[1,1,H,W]`.
49pub fn depth_tensor_to_depth_map(tensor: &TensorBuffer) -> VisionResult<DepthMap> {
50    let (height, width, values) = flatten_single_channel_f32(tensor, "depth")?;
51    DepthMap::try_new(width, height, values)
52}
53
54/// Converts a host score tensor into a binary mask via threshold.
55///
56/// Accepted shapes: `[H,W]`, `[1,H,W]`, or `[1,1,H,W]`.
57pub fn score_tensor_to_binary_mask(
58    tensor: &TensorBuffer,
59    threshold: f32,
60) -> VisionResult<crate::BinaryMask> {
61    if !threshold.is_finite() {
62        return Err(VisionError::InvalidParameter("mask threshold must be finite".into()));
63    }
64    let (height, width, values) = flatten_single_channel_f32(tensor, "scores")?;
65    crate::BinaryMask::try_new(
66        width,
67        height,
68        values.into_iter().map(|value| u8::from(value.is_finite() && value >= threshold)).collect(),
69    )
70}
71
72/// Decodes axis-aligned detections from an `[N,6]` tensor of
73/// `x0,y0,x1,y1,score,class`.
74pub fn detection_tensor_to_detections(tensor: &TensorBuffer) -> VisionResult<Vec<Detection>> {
75    let descriptor = tensor.descriptor();
76    if descriptor.dtype() != DataType::F32 {
77        return Err(VisionError::InvalidParameter(format!(
78            "detections require f32, found {:?}",
79            descriptor.dtype()
80        )));
81    }
82    let shape = descriptor.shape();
83    if shape.len() != 2 || shape[1] != 6 {
84        return Err(VisionError::ShapeMismatch(format!(
85            "detections expect [N,6], found {shape:?}"
86        )));
87    }
88    if !descriptor.is_c_contiguous() || descriptor.byte_offset() != 0 {
89        return Err(VisionError::InvalidParameter(
90            "detections tensor must be contiguous with byte_offset=0".into(),
91        ));
92    }
93    let values = f32_slice(tensor, "detections")?;
94    let mut detections = Vec::with_capacity(shape[0]);
95    for row in values.chunks_exact(6) {
96        let bbox = BoundingBox2::try_new(row[0], row[1], row[2], row[3])
97            .map_err(|error| VisionError::InvalidParameter(error.to_string()))?;
98        detections.push(Detection { bbox, score: row[4], class_id: row[5] as i64 });
99    }
100    Ok(detections)
101}
102
103fn flatten_single_channel_f32(
104    tensor: &TensorBuffer,
105    name: &str,
106) -> VisionResult<(usize, usize, Vec<f32>)> {
107    let descriptor = tensor.descriptor();
108    if descriptor.dtype() != DataType::F32 {
109        return Err(VisionError::InvalidParameter(format!(
110            "{name} requires f32, found {:?}",
111            descriptor.dtype()
112        )));
113    }
114    if !descriptor.is_c_contiguous() || descriptor.byte_offset() != 0 {
115        return Err(VisionError::InvalidParameter(format!(
116            "{name} tensor must be contiguous with byte_offset=0"
117        )));
118    }
119    let shape = descriptor.shape();
120    let (height, width) = match shape {
121        [h, w] => (*h, *w),
122        [1, h, w] => (*h, *w),
123        [1, 1, h, w] => (*h, *w),
124        other => {
125            return Err(VisionError::ShapeMismatch(format!(
126                "{name} expects [H,W], [1,H,W], or [1,1,H,W]; found {other:?}"
127            )));
128        }
129    };
130    let values = f32_slice(tensor, name)?.to_vec();
131    if values.len() != height.saturating_mul(width) {
132        return Err(VisionError::ShapeMismatch(format!(
133            "{name} length {} does not match {width}x{height}",
134            values.len()
135        )));
136    }
137    Ok((height, width, values))
138}
139
140fn f32_slice<'a>(tensor: &'a TensorBuffer, name: &str) -> VisionResult<&'a [f32]> {
141    let bytes = tensor.allocation_bytes();
142    if bytes.len() % 4 != 0 {
143        return Err(VisionError::InvalidParameter(format!(
144            "{name} f32 allocation is not a multiple of 4 bytes"
145        )));
146    }
147    Ok(bytemuck::cast_slice(bytes))
148}
149
150#[cfg(test)]
151mod tests {
152    use super::{
153        depth_tensor_to_depth_map, detection_tensor_to_detections, planar_f32_to_nchw,
154        rgb_u8_to_nchw_f32, score_tensor_to_binary_mask,
155    };
156    use crate::Interpolation;
157    use spatialrust_image::{Image, PlanarImage};
158    use spatialrust_tensor::{DataType, Device, TensorBuffer, TensorDescriptor};
159
160    #[test]
161    fn letterbox_pack_produces_batch_nchw() {
162        let image = Image::<u8, 3>::try_new(4, 2, vec![255; 4 * 2 * 3]).unwrap();
163        let (tensor, mapping) = rgb_u8_to_nchw_f32(
164            image.view(),
165            4,
166            4,
167            Interpolation::Nearest,
168            [0, 0, 0],
169            1.0 / 255.0,
170            [0.0; 3],
171            [1.0; 3],
172        )
173        .unwrap();
174        assert_eq!(tensor.descriptor().shape(), &[1, 3, 4, 4]);
175        assert!(mapping.scale > 0.0);
176    }
177
178    #[test]
179    fn planar_to_nchw_keeps_values() {
180        let planar =
181            PlanarImage::<f32, 3>::try_new(2, 2, (0..12).map(|v| v as f32).collect()).unwrap();
182        let tensor = planar_f32_to_nchw(&planar).unwrap();
183        assert_eq!(tensor.descriptor().shape(), &[1, 3, 2, 2]);
184        assert_eq!(bytemuck::cast_slice::<u8, f32>(tensor.allocation_bytes()), planar.as_slice());
185    }
186
187    #[test]
188    fn depth_and_mask_decode_single_channel_shapes() {
189        let depth = TensorBuffer::try_from_f32(
190            vec![1.0, 2.0, 3.0, 4.0],
191            TensorDescriptor::contiguous(DataType::F32, vec![1, 1, 2, 2], Device::CPU),
192        )
193        .unwrap();
194        let map = depth_tensor_to_depth_map(&depth).unwrap();
195        assert_eq!(map.image().width(), 2);
196        assert_eq!(map.image().as_slice(), &[1.0, 2.0, 3.0, 4.0]);
197
198        let scores = TensorBuffer::try_from_f32(
199            vec![0.1, 0.9, 0.4, 0.8],
200            TensorDescriptor::contiguous(DataType::F32, vec![2, 2], Device::CPU),
201        )
202        .unwrap();
203        let mask = score_tensor_to_binary_mask(&scores, 0.5).unwrap();
204        assert_eq!(mask.image().as_slice(), &[0, 1, 0, 1]);
205    }
206
207    #[test]
208    fn detections_decode_rows() {
209        let tensor = TensorBuffer::try_from_f32(
210            vec![0.0, 0.0, 2.0, 2.0, 0.9, 1.0],
211            TensorDescriptor::contiguous(DataType::F32, vec![1, 6], Device::CPU),
212        )
213        .unwrap();
214        let detections = detection_tensor_to_detections(&tensor).unwrap();
215        assert_eq!(detections.len(), 1);
216        assert_eq!(detections[0].class_id, 1);
217        assert!((detections[0].score - 0.9).abs() < 1e-6);
218    }
219}
spatialrust_vision/adapters.rs

spatialrust_vision/
adapters.rs
