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    /// Emits deterministic class IDs `[1,N]` and confidences `[1,N]` from
19    /// normalized point features `[1,4,N]`.
20    ///
21    /// The four channels are `x`, `y`, `z`, and horizontal radial distance.
22    /// This profile is a visual/demo fixture, not a learned ontology or a
23    /// substitute for a production model receipt.
24    SemanticClasses,
25}
26
27/// Backend that never loads ONNX bytes and always uses [`MockProfile`].
28#[derive(Clone, Debug, Default)]
29pub struct MockInferenceBackend;
30
31impl InferenceBackend for MockInferenceBackend {
32    fn name(&self) -> &str {
33        "mock"
34    }
35
36    fn create_session(
37        &self,
38        source: &ModelSource,
39        options: &SessionOptions,
40    ) -> AiResult<Box<dyn ModelSession>> {
41        options.validate()?;
42        let profile = match source {
43            ModelSource::Mock(profile) => *profile,
44            ModelSource::Path(_) | ModelSource::Bytes(_) => {
45                return Err(AiError::Unsupported {
46                    backend: "mock".into(),
47                    operation: "filesystem or byte-backed model loading".into(),
48                });
49            }
50        };
51        Ok(Box::new(MockSession { profile, info: profile.model_info() }))
52    }
53}
54
55struct MockSession {
56    profile: MockProfile,
57    info: ModelInfo,
58}
59
60impl ModelSession for MockSession {
61    fn backend_name(&self) -> &str {
62        "mock"
63    }
64
65    fn model_info(&self) -> &ModelInfo {
66        &self.info
67    }
68
69    fn run_with_options(
70        &mut self,
71        inputs: NamedTensors,
72        options: RunOptions,
73    ) -> AiResult<NamedTensors> {
74        self.info.validate_inputs(&inputs)?;
75        // Output is always a newly allocated host `TensorBuffer`.
76        if options.output_copy == CopyPolicy::Forbid {
77            let name = match self.profile {
78                MockProfile::SyntheticDepth => "depth",
79                MockProfile::SemanticClasses => "class_ids",
80            };
81            return Err(AiError::CopyRequired { direction: "output host", name: name.into() });
82        }
83        match self.profile {
84            MockProfile::SyntheticDepth => run_synthetic_depth(inputs, options.input_copy),
85            MockProfile::SemanticClasses => run_semantic_classes(inputs, options.input_copy),
86        }
87    }
88}
89
90impl MockProfile {
91    fn model_info(self) -> ModelInfo {
92        match self {
93            Self::SyntheticDepth => ModelInfo {
94                name: Some("mock-synthetic-depth".into()),
95                inputs: vec![TensorSpec::new(
96                    "images",
97                    DataType::F32,
98                    vec![
99                        Dimension::Fixed(1),
100                        Dimension::Fixed(3),
101                        Dimension::Dynamic,
102                        Dimension::Dynamic,
103                    ],
104                )],
105                outputs: vec![TensorSpec::new(
106                    "depth",
107                    DataType::F32,
108                    vec![
109                        Dimension::Fixed(1),
110                        Dimension::Fixed(1),
111                        Dimension::Dynamic,
112                        Dimension::Dynamic,
113                    ],
114                )],
115            },
116            Self::SemanticClasses => ModelInfo {
117                name: Some("mock-semantic-classes".into()),
118                inputs: vec![TensorSpec::new(
119                    "features",
120                    DataType::F32,
121                    vec![Dimension::Fixed(1), Dimension::Fixed(4), Dimension::Dynamic],
122                )],
123                outputs: vec![
124                    TensorSpec::new(
125                        "class_ids",
126                        DataType::U32,
127                        vec![Dimension::Fixed(1), Dimension::Dynamic],
128                    ),
129                    TensorSpec::new(
130                        "confidence",
131                        DataType::F32,
132                        vec![Dimension::Fixed(1), Dimension::Dynamic],
133                    ),
134                ],
135            },
136        }
137    }
138}
139
140fn run_synthetic_depth(inputs: NamedTensors, input_copy: CopyPolicy) -> AiResult<NamedTensors> {
141    let input = inputs.get("images").ok_or_else(|| AiError::MissingInput("images".into()))?;
142    let descriptor = input.descriptor();
143    let shape = descriptor.shape();
144    if shape.len() != 4 || shape[0] != 1 || shape[1] != 3 {
145        return Err(AiError::ShapeMismatch {
146            name: "images".into(),
147            expected: vec![
148                Dimension::Fixed(1),
149                Dimension::Fixed(3),
150                Dimension::Dynamic,
151                Dimension::Dynamic,
152            ],
153            actual: shape.to_vec(),
154        });
155    }
156    if !descriptor.is_c_contiguous() || descriptor.byte_offset() != 0 {
157        if input_copy == CopyPolicy::Forbid {
158            return Err(AiError::CopyRequired { direction: "input host", name: "images".into() });
159        }
160        return Err(AiError::Unsupported {
161            backend: "mock".into(),
162            operation: "non-contiguous input packing for SyntheticDepth".into(),
163        });
164    }
165    let height = shape[2];
166    let width = shape[3];
167    let values = f32_values(input)?;
168    let plane = height.saturating_mul(width);
169    if values.len() != plane.saturating_mul(3) {
170        return Err(AiError::InvalidConfiguration(
171            "images storage length does not match NCHW shape".into(),
172        ));
173    }
174    let mut depth = Vec::with_capacity(plane);
175    for index in 0..plane {
176        let r = values[index];
177        let g = values[plane + index];
178        let b = values[2 * plane + index];
179        let luma = (0.299 * r + 0.587 * g + 0.114 * b).clamp(0.0, 1.0);
180        depth.push(1.0 + 0.5 * (1.0 - luma));
181    }
182    let output = TensorBuffer::try_from_f32(
183        depth,
184        TensorDescriptor::contiguous(DataType::F32, vec![1, 1, height, width], Device::CPU),
185    )
186    .map_err(|error| AiError::InvalidConfiguration(error.to_string()))?;
187    let mut outputs = NamedTensors::new();
188    outputs.insert("depth", output)?;
189    Ok(outputs)
190}
191
192fn run_semantic_classes(inputs: NamedTensors, input_copy: CopyPolicy) -> AiResult<NamedTensors> {
193    let input = inputs.get("features").ok_or_else(|| AiError::MissingInput("features".into()))?;
194    let descriptor = input.descriptor();
195    let shape = descriptor.shape();
196    if shape.len() != 3 || shape[0] != 1 || shape[1] != 4 {
197        return Err(AiError::ShapeMismatch {
198            name: "features".into(),
199            expected: vec![Dimension::Fixed(1), Dimension::Fixed(4), Dimension::Dynamic],
200            actual: shape.to_vec(),
201        });
202    }
203    if !descriptor.is_c_contiguous() || descriptor.byte_offset() != 0 {
204        if input_copy == CopyPolicy::Forbid {
205            return Err(AiError::CopyRequired { direction: "input host", name: "features".into() });
206        }
207        return Err(AiError::Unsupported {
208            backend: "mock".into(),
209            operation: "non-contiguous input packing for SemanticClasses".into(),
210        });
211    }
212    let count = shape[2];
213    let values = f32_values(input)?;
214    let plane = count;
215    if values.len() != plane.saturating_mul(4) {
216        return Err(AiError::InvalidConfiguration(
217            "features storage length does not match [1,4,N] shape".into(),
218        ));
219    }
220    let mut class_ids = Vec::with_capacity(count);
221    let mut confidence = Vec::with_capacity(count);
222    for index in 0..count {
223        let x = values[index];
224        let y = values[plane + index];
225        let z = values[2 * plane + index];
226        let radial = values[3 * plane + index];
227        if [x, y, z, radial].iter().any(|value| !value.is_finite()) {
228            return Err(AiError::InvalidConfiguration(
229                "SemanticClasses features must be finite".into(),
230            ));
231        }
232        let class_id = if z < 0.28 {
233            0
234        } else if radial > 0.55 {
235            1
236        } else {
237            2
238        };
239        let margin = match class_id {
240            0 => (0.28 - z).abs(),
241            1 => (radial - 0.55).abs(),
242            _ => (z - 0.28).abs().min((radial - 0.55).abs()),
243        };
244        class_ids.push(class_id);
245        confidence.push((0.60 + 0.40 * (margin * 2.0).clamp(0.0, 1.0)).clamp(0.0, 1.0));
246    }
247    let class_output = TensorBuffer::try_from_u32(
248        class_ids,
249        TensorDescriptor::contiguous(DataType::U32, vec![1, count], Device::CPU),
250    )
251    .map_err(|error| AiError::InvalidConfiguration(error.to_string()))?;
252    let confidence_output = TensorBuffer::try_from_f32(
253        confidence,
254        TensorDescriptor::contiguous(DataType::F32, vec![1, count], Device::CPU),
255    )
256    .map_err(|error| AiError::InvalidConfiguration(error.to_string()))?;
257    let mut outputs = NamedTensors::new();
258    outputs.insert("class_ids", class_output)?;
259    outputs.insert("confidence", confidence_output)?;
260    Ok(outputs)
261}
262
263fn f32_values(tensor: &TensorBuffer) -> AiResult<Vec<f32>> {
264    if tensor.descriptor().dtype() != DataType::F32 {
265        return Err(AiError::DataTypeMismatch {
266            name: "images".into(),
267            expected: DataType::F32,
268            actual: tensor.descriptor().dtype(),
269        });
270    }
271    if let Some(values) = tensor.shared_f32() {
272        return Ok(values.to_vec());
273    }
274    let bytes = tensor.allocation_bytes();
275    if bytes.len() % 4 != 0 {
276        return Err(AiError::InvalidConfiguration(
277            "f32 tensor allocation is not a multiple of 4 bytes".into(),
278        ));
279    }
280    Ok(bytemuck::cast_slice(bytes).to_vec())
281}
282
283#[cfg(test)]
284mod tests {
285    use super::{MockInferenceBackend, MockProfile};
286    use crate::{
287        AiError, CopyPolicy, InferenceBackend, ModelSource, NamedTensors, RunOptions,
288        SessionOptions,
289    };
290    use spatialrust_tensor::{DataType, Device, TensorBuffer, TensorDescriptor};
291
292    #[test]
293    fn synthetic_depth_emits_metric_plane() {
294        let backend = MockInferenceBackend;
295        let mut session = backend
296            .create_session(
297                &ModelSource::Mock(MockProfile::SyntheticDepth),
298                &SessionOptions::default(),
299            )
300            .unwrap();
301        // planar RRRRGGGG BBBB for 2x2
302        let values = vec![
303            0.0, 1.0, 0.5, 0.25, // R
304            0.0, 1.0, 0.5, 0.25, // G
305            0.0, 1.0, 0.5, 0.25, // B
306        ];
307        let input = TensorBuffer::try_from_f32(
308            values,
309            TensorDescriptor::contiguous(DataType::F32, vec![1, 3, 2, 2], Device::CPU),
310        )
311        .unwrap();
312        let mut inputs = NamedTensors::new();
313        inputs.insert("images", input).unwrap();
314        let outputs = session
315            .run_with_options(
316                inputs,
317                RunOptions { input_copy: CopyPolicy::Forbid, output_copy: CopyPolicy::Allow },
318            )
319            .unwrap();
320        let depth = outputs.get("depth").unwrap();
321        assert_eq!(depth.descriptor().shape(), &[1, 1, 2, 2]);
322        assert_eq!(session.model_info().name.as_deref(), Some("mock-synthetic-depth"));
323        let depth_values = depth.shared_f32().unwrap();
324        let near = depth_values[1]; // bright pixel -> nearer
325        let far = depth_values[0];
326        assert!(near < far);
327    }
328
329    #[test]
330    fn synthetic_depth_requires_output_copy_permission() {
331        let backend = MockInferenceBackend;
332        let mut session = backend
333            .create_session(
334                &ModelSource::Mock(MockProfile::SyntheticDepth),
335                &SessionOptions::default(),
336            )
337            .unwrap();
338        let input = TensorBuffer::try_from_f32(
339            vec![0.0; 12],
340            TensorDescriptor::contiguous(DataType::F32, vec![1, 3, 2, 2], Device::CPU),
341        )
342        .unwrap();
343        let mut inputs = NamedTensors::new();
344        inputs.insert("images", input).unwrap();
345        assert!(matches!(
346            session.run(inputs),
347            Err(AiError::CopyRequired { direction: "output host", .. })
348        ));
349    }
350
351    #[test]
352    fn semantic_classes_are_deterministic_and_explicitly_host_bound() {
353        let backend = MockInferenceBackend;
354        let mut session = backend
355            .create_session(
356                &ModelSource::Mock(MockProfile::SemanticClasses),
357                &SessionOptions { deterministic: true, ..SessionOptions::default() },
358            )
359            .unwrap();
360        let values = vec![
361            0.1, 0.2, 0.3, // x
362            0.1, 0.2, 0.3, // y
363            0.1, 0.5, 0.8, // z
364            0.1, 0.7, 0.2, // radial
365        ];
366        let input = TensorBuffer::try_from_f32(
367            values,
368            TensorDescriptor::contiguous(DataType::F32, vec![1, 4, 3], Device::CPU),
369        )
370        .unwrap();
371        let mut inputs = NamedTensors::new();
372        inputs.insert("features", input).unwrap();
373        let outputs = session
374            .run_with_options(
375                inputs,
376                RunOptions { input_copy: CopyPolicy::Forbid, output_copy: CopyPolicy::Allow },
377            )
378            .unwrap();
379        assert_eq!(outputs.get("class_ids").unwrap().shared_u32().unwrap().as_ref(), &[0, 1, 2]);
380        assert_eq!(outputs.get("confidence").unwrap().descriptor().shape(), &[1, 3]);
381        assert_eq!(session.model_info().name.as_deref(), Some("mock-semantic-classes"));
382    }
383}
spatialrust_ai/mock.rs

spatialrust_ai/
mock.rs
