xcc.es
[READ-ONLY] Mirror of https://github.com/mrgnw/spatial-maker.
2.7 kB
96 lines
1pub mod depth;
2pub mod error;
3pub mod image_loader;
4pub mod model;
5pub mod output;
6pub mod stereo;
7pub mod video;
8
9#[cfg(all(target_os = "macos", feature = "coreml"))]
10pub mod depth_coreml;
11
12pub use error::{SpatialError, SpatialResult};
13pub use image_loader::load_image;
14pub use model::{find_model, get_checkpoint_dir, model_exists};
15pub use output::{create_sbs_image, save_stereo_image, ImageEncoding, MVHEVCConfig, OutputFormat, OutputOptions};
16pub use stereo::generate_stereo_pair;
17pub use video::{get_video_metadata, process_video, ProgressCallback, VideoMetadata, VideoProgress};
18
19#[cfg(all(target_os = "macos", feature = "coreml"))]
20pub use depth_coreml::CoreMLDepthEstimator;
21
22#[cfg(feature = "onnx")]
23pub use depth::OnnxDepthEstimator;
24
25use std::path::Path;
26
27#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
28pub struct SpatialConfig {
29 pub encoder_size: String,
30 pub max_disparity: u32,
31 pub target_depth_size: u32,
32}
33
34pub type StereoOutputFormat = OutputFormat;
35
36impl Default for SpatialConfig {
37 fn default() -> Self {
38 Self {
39 encoder_size: "s".to_string(),
40 max_disparity: 30,
41 target_depth_size: 518,
42 }
43 }
44}
45
46pub async fn process_photo(
47 input_path: &Path,
48 output_path: &Path,
49 config: SpatialConfig,
50 output_options: OutputOptions,
51) -> SpatialResult<()> {
52 let input_image = load_image(input_path).await?;
53
54 model::ensure_model_exists::<fn(u64, u64)>(&config.encoder_size, None).await?;
55
56 #[cfg(all(target_os = "macos", feature = "coreml"))]
57 let depth_map = {
58 let model_path = model::find_model(&config.encoder_size)?;
59 let model_str = model_path.to_str().ok_or_else(|| {
60 SpatialError::ModelError("Invalid model path encoding".to_string())
61 })?;
62 let estimator = CoreMLDepthEstimator::new(model_str)?;
63 estimator.estimate(&input_image)?
64 };
65
66 #[cfg(not(all(target_os = "macos", feature = "coreml")))]
67 let depth_map = {
68 #[cfg(feature = "onnx")]
69 {
70 let model_path = model::find_model(&config.encoder_size)?;
71 let estimator = OnnxDepthEstimator::new(model_path.to_str().unwrap())?;
72 estimator.estimate(&input_image)?
73 }
74 #[cfg(not(feature = "onnx"))]
75 {
76 return Err(SpatialError::ConfigError(
77 "No depth backend enabled. Enable 'coreml' (macOS) or 'onnx' feature.".to_string(),
78 ));
79 }
80 };
81
82 let (left, right) = generate_stereo_pair(&input_image, &depth_map, config.max_disparity)?;
83
84 save_stereo_image(&left, &right, output_path, output_options)?;
85
86 Ok(())
87}
88
89pub async fn process_video_sbs(
90 input_path: &Path,
91 output_path: &Path,
92 config: SpatialConfig,
93 progress_cb: Option<ProgressCallback>,
94) -> SpatialResult<()> {
95 video::process_video(input_path, output_path, config, progress_cb).await
96}