spatialrust_ai/
mock.rs

1//! Deterministic in-process mock inference backends.
2
3use spatialrust_tensor::{DataType, Device, TensorBuffer, TensorDescriptor};
4
5use crate::{
6    AiError, AiResult, CopyPolicy, Dimension, InferenceBackend, ModelInfo, ModelSession,
7    ModelSource, NamedTensors, RunOptions, SessionOptions, TensorSpec,
8};
9
10/// Built-in mock model profiles selected through [`ModelSource::Mock`].
11#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
12pub enum MockProfile {
13    /// Emits metric depth `[1,1,H,W]` from RGB NCHW `[1,3,H,W]` using BT.601 luma.
14    ///
15    /// Depth meters are `1.0 + 0.5 * (1.0 - luminance)` so brighter pixels map
16    /// nearer, giving a deterministic textured plane for image → AI → XYZ demos.
17    SyntheticDepth,
18}
19
20/// Backend that never loads ONNX bytes and always uses [`MockProfile`].
21#[derive(Clone, Debug, Default)]
22pub struct MockInferenceBackend;
23
24impl InferenceBackend for MockInferenceBackend {
25    fn name(&self) -> &str {
26        "mock"
27    }
28
29    fn create_session(
30        &self,
31        source: &ModelSource,
32        options: &SessionOptions,
33    ) -> AiResult<Box<dyn ModelSession>> {
34        options.validate()?;
35        let profile = match source {
36            ModelSource::Mock(profile) => *profile,
37            ModelSource::Path(_) | ModelSource::Bytes(_) => {
38                return Err(AiError::Unsupported {
39                    backend: "mock".into(),
40                    operation: "filesystem or byte-backed model loading".into(),
41                });
42            }
43        };
44        Ok(Box::new(MockSession { profile, info: profile.model_info() }))
45    }
46}
47
48struct MockSession {
49    profile: MockProfile,
50    info: ModelInfo,
51}
52
53impl ModelSession for MockSession {
54    fn backend_name(&self) -> &str {
55        "mock"
56    }
57
58    fn model_info(&self) -> &ModelInfo {
59        &self.info
60    }
61
62    fn run_with_options(
63        &mut self,
64        inputs: NamedTensors,
65        options: RunOptions,
66    ) -> AiResult<NamedTensors> {
67        self.info.validate_inputs(&inputs)?;
68        // Output is always a newly allocated host `TensorBuffer`.
69        if options.output_copy == CopyPolicy::Forbid {
70            return Err(AiError::CopyRequired { direction: "output host", name: "depth".into() });
71        }
72        match self.profile {
73            MockProfile::SyntheticDepth => run_synthetic_depth(inputs, options.input_copy),
74        }
75    }
76}
77
78impl MockProfile {
79    fn model_info(self) -> ModelInfo {
80        match self {
81            Self::SyntheticDepth => ModelInfo {
82                name: Some("mock-synthetic-depth".into()),
83                inputs: vec![TensorSpec::new(
84                    "images",
85                    DataType::F32,
86                    vec![
87                        Dimension::Fixed(1),
88                        Dimension::Fixed(3),
89                        Dimension::Dynamic,
90                        Dimension::Dynamic,
91                    ],
92                )],
93                outputs: vec![TensorSpec::new(
94                    "depth",
95                    DataType::F32,
96                    vec![
97                        Dimension::Fixed(1),
98                        Dimension::Fixed(1),
99                        Dimension::Dynamic,
100                        Dimension::Dynamic,
101                    ],
102                )],
103            },
104        }
105    }
106}
107
108fn run_synthetic_depth(inputs: NamedTensors, input_copy: CopyPolicy) -> AiResult<NamedTensors> {
109    let input = inputs.get("images").ok_or_else(|| AiError::MissingInput("images".into()))?;
110    let descriptor = input.descriptor();
111    let shape = descriptor.shape();
112    if shape.len() != 4 || shape[0] != 1 || shape[1] != 3 {
113        return Err(AiError::ShapeMismatch {
114            name: "images".into(),
115            expected: vec![
116                Dimension::Fixed(1),
117                Dimension::Fixed(3),
118                Dimension::Dynamic,
119                Dimension::Dynamic,
120            ],
121            actual: shape.to_vec(),
122        });
123    }
124    if !descriptor.is_c_contiguous() || descriptor.byte_offset() != 0 {
125        if input_copy == CopyPolicy::Forbid {
126            return Err(AiError::CopyRequired { direction: "input host", name: "images".into() });
127        }
128        return Err(AiError::Unsupported {
129            backend: "mock".into(),
130            operation: "non-contiguous input packing for SyntheticDepth".into(),
131        });
132    }
133    let height = shape[2];
134    let width = shape[3];
135    let values = f32_values(input)?;
136    let plane = height.saturating_mul(width);
137    if values.len() != plane.saturating_mul(3) {
138        return Err(AiError::InvalidConfiguration(
139            "images storage length does not match NCHW shape".into(),
140        ));
141    }
142    let mut depth = Vec::with_capacity(plane);
143    for index in 0..plane {
144        let r = values[index];
145        let g = values[plane + index];
146        let b = values[2 * plane + index];
147        let luma = (0.299 * r + 0.587 * g + 0.114 * b).clamp(0.0, 1.0);
148        depth.push(1.0 + 0.5 * (1.0 - luma));
149    }
150    let output = TensorBuffer::try_from_f32(
151        depth,
152        TensorDescriptor::contiguous(DataType::F32, vec![1, 1, height, width], Device::CPU),
153    )
154    .map_err(|error| AiError::InvalidConfiguration(error.to_string()))?;
155    let mut outputs = NamedTensors::new();
156    outputs.insert("depth", output)?;
157    Ok(outputs)
158}
159
160fn f32_values(tensor: &TensorBuffer) -> AiResult<Vec<f32>> {
161    if tensor.descriptor().dtype() != DataType::F32 {
162        return Err(AiError::DataTypeMismatch {
163            name: "images".into(),
164            expected: DataType::F32,
165            actual: tensor.descriptor().dtype(),
166        });
167    }
168    if let Some(values) = tensor.shared_f32() {
169        return Ok(values.to_vec());
170    }
171    let bytes = tensor.allocation_bytes();
172    if bytes.len() % 4 != 0 {
173        return Err(AiError::InvalidConfiguration(
174            "f32 tensor allocation is not a multiple of 4 bytes".into(),
175        ));
176    }
177    Ok(bytemuck::cast_slice(bytes).to_vec())
178}
179
180#[cfg(test)]
181mod tests {
182    use super::{MockInferenceBackend, MockProfile};
183    use crate::{
184        AiError, CopyPolicy, InferenceBackend, ModelSource, NamedTensors, RunOptions,
185        SessionOptions,
186    };
187    use spatialrust_tensor::{DataType, Device, TensorBuffer, TensorDescriptor};
188
189    #[test]
190    fn synthetic_depth_emits_metric_plane() {
191        let backend = MockInferenceBackend;
192        let mut session = backend
193            .create_session(
194                &ModelSource::Mock(MockProfile::SyntheticDepth),
195                &SessionOptions::default(),
196            )
197            .unwrap();
198        // planar RRRRGGGG BBBB for 2x2
199        let values = vec![
200            0.0, 1.0, 0.5, 0.25, // R
201            0.0, 1.0, 0.5, 0.25, // G
202            0.0, 1.0, 0.5, 0.25, // B
203        ];
204        let input = TensorBuffer::try_from_f32(
205            values,
206            TensorDescriptor::contiguous(DataType::F32, vec![1, 3, 2, 2], Device::CPU),
207        )
208        .unwrap();
209        let mut inputs = NamedTensors::new();
210        inputs.insert("images", input).unwrap();
211        let outputs = session
212            .run_with_options(
213                inputs,
214                RunOptions { input_copy: CopyPolicy::Forbid, output_copy: CopyPolicy::Allow },
215            )
216            .unwrap();
217        let depth = outputs.get("depth").unwrap();
218        assert_eq!(depth.descriptor().shape(), &[1, 1, 2, 2]);
219        assert_eq!(session.model_info().name.as_deref(), Some("mock-synthetic-depth"));
220        let depth_values = depth.shared_f32().unwrap();
221        let near = depth_values[1]; // bright pixel -> nearer
222        let far = depth_values[0];
223        assert!(near < far);
224    }
225
226    #[test]
227    fn synthetic_depth_requires_output_copy_permission() {
228        let backend = MockInferenceBackend;
229        let mut session = backend
230            .create_session(
231                &ModelSource::Mock(MockProfile::SyntheticDepth),
232                &SessionOptions::default(),
233            )
234            .unwrap();
235        let input = TensorBuffer::try_from_f32(
236            vec![0.0; 12],
237            TensorDescriptor::contiguous(DataType::F32, vec![1, 3, 2, 2], Device::CPU),
238        )
239        .unwrap();
240        let mut inputs = NamedTensors::new();
241        inputs.insert("images", input).unwrap();
242        assert!(matches!(
243            session.run(inputs),
244            Err(AiError::CopyRequired { direction: "output host", .. })
245        ));
246    }
247}
spatialrust_ai/mock.rs

spatialrust_ai/
mock.rs
