xcc.es
[READ-ONLY] Mirror of https://github.com/mrgnw/spatial-maker.
3.4 kB
131 lines
1use crate::error::{SpatialError, SpatialResult};
2use image::{DynamicImage, ImageBuffer, Luma};
3use ndarray::Array2;
4use std::ffi::CString;
5
6const INPUT_SIZE: u32 = 518;
7const IMAGENET_MEAN: [f32; 3] = [0.485, 0.456, 0.406];
8const IMAGENET_STD: [f32; 3] = [0.229, 0.224, 0.225];
9
10extern "C" {
11 fn coreml_load_model(path: *const std::os::raw::c_char) -> *mut std::os::raw::c_void;
12 fn coreml_unload_model(model: *mut std::os::raw::c_void);
13 fn coreml_infer_depth(
14 model: *mut std::os::raw::c_void,
15 rgb_data: *const f32,
16 width: i32,
17 height: i32,
18 output: *mut f32,
19 ) -> i32;
20}
21
22pub struct CoreMLDepthEstimator {
23 model: *mut std::os::raw::c_void,
24}
25
26impl CoreMLDepthEstimator {
27 pub fn new(model_path: &str) -> SpatialResult<Self> {
28 let c_path = CString::new(model_path)
29 .map_err(|e| SpatialError::ModelError(format!("Invalid model path: {}", e)))?;
30
31 let model = unsafe { coreml_load_model(c_path.as_ptr()) };
32
33 if model.is_null() {
34 return Err(SpatialError::ModelError(format!(
35 "Failed to load CoreML model: {}",
36 model_path
37 )));
38 }
39
40 tracing::info!("CoreML model loaded: {}", model_path);
41
42 Ok(Self { model })
43 }
44
45 pub fn estimate_raw(&self, image: &DynamicImage) -> SpatialResult<ImageBuffer<Luma<f32>, Vec<f32>>> {
46 let (orig_width, orig_height) = (image.width(), image.height());
47
48 let resized = image.resize_exact(
49 INPUT_SIZE,
50 INPUT_SIZE,
51 image::imageops::FilterType::Lanczos3,
52 );
53
54 let rgb = resized.to_rgb8();
55 let mut input_data = Vec::with_capacity(3 * INPUT_SIZE as usize * INPUT_SIZE as usize);
56
57 for c in 0..3 {
58 for pixel in rgb.pixels() {
59 let normalized = (pixel[c] as f32 / 255.0 - IMAGENET_MEAN[c]) / IMAGENET_STD[c];
60 input_data.push(normalized);
61 }
62 }
63
64 let output_size = (INPUT_SIZE * INPUT_SIZE) as usize;
65 let mut output_data = vec![0.0f32; output_size];
66
67 let result = unsafe {
68 coreml_infer_depth(
69 self.model,
70 input_data.as_ptr(),
71 INPUT_SIZE as i32,
72 INPUT_SIZE as i32,
73 output_data.as_mut_ptr(),
74 )
75 };
76
77 if result != 0 {
78 return Err(SpatialError::ModelError(format!(
79 "CoreML inference failed with error code: {}",
80 result
81 )));
82 }
83
84 let min_val = output_data.iter().copied().fold(f32::INFINITY, f32::min);
85 let max_val = output_data
86 .iter()
87 .copied()
88 .fold(f32::NEG_INFINITY, f32::max);
89 let range = max_val - min_val;
90
91 if range > 1e-6 {
92 for v in &mut output_data {
93 *v = (*v - min_val) / range;
94 }
95 }
96
97 let depth_image = ImageBuffer::from_fn(INPUT_SIZE, INPUT_SIZE, |x, y| {
98 let idx = (y * INPUT_SIZE + x) as usize;
99 Luma([output_data[idx]])
100 });
101
102 let resized_depth = image::imageops::resize(
103 &depth_image,
104 orig_width,
105 orig_height,
106 image::imageops::FilterType::Lanczos3,
107 );
108
109 Ok(resized_depth)
110 }
111
112 pub fn estimate(&self, image: &DynamicImage) -> SpatialResult<Array2<f32>> {
113 let depth_image = self.estimate_raw(image)?;
114 let (width, height) = depth_image.dimensions();
115 let data: Vec<f32> = depth_image.pixels().map(|p| p[0]).collect();
116 let depth_2d = Array2::from_shape_vec((height as usize, width as usize), data)
117 .map_err(|e| SpatialError::TensorError(format!("Failed to reshape depth: {}", e)))?;
118 Ok(depth_2d)
119 }
120}
121
122impl Drop for CoreMLDepthEstimator {
123 fn drop(&mut self) {
124 unsafe {
125 coreml_unload_model(self.model);
126 }
127 }
128}
129
130unsafe impl Send for CoreMLDepthEstimator {}
131unsafe impl Sync for CoreMLDepthEstimator {}