xcc.es / spatial-maker
[READ-ONLY] Mirror of https://github.com/mrgnw/spatial-maker.
0
Fork 0

Configure Feed

Select the types of activity you want to include in your feed.

v0.4.0: add --output-types flag with depth map export (avif/png/png16)

Replace --format and --mvhevc with unified --output-types flag.
Supports depth:avif,png,png16, sbs, tab, sep, spatial as
comma-separated output targets. Pipeline skips stereo generation
when only depth is requested. Video depth exports as grayscale
AV1 .mov. Also adds depth filtering (bilateral, gaussian blur,
temporal EMA), progress bars, and fixes CoreML CVPixelBuffer
row padding.

author
Morgan Williams date (Feb 19, 2026, 2:59 PM +0100) commit 5bddd51c parent 043def2d
+1186 -129
+66 -1
Cargo.lock
··· 313 313 checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" 314 314 315 315 [[package]] 316 + name = "console" 317 + version = "0.15.11" 318 + source = "registry+https://github.com/rust-lang/crates.io-index" 319 + checksum = "054ccb5b10f9f2cbf51eb355ca1d05c2d279ce1804688d0db74b4733a5aeafd8" 320 + dependencies = [ 321 + "encode_unicode", 322 + "libc", 323 + "once_cell", 324 + "unicode-width", 325 + "windows-sys 0.59.0", 326 + ] 327 + 328 + [[package]] 316 329 name = "core-foundation" 317 330 version = "0.9.4" 318 331 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 436 449 checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" 437 450 438 451 [[package]] 452 + name = "encode_unicode" 453 + version = "1.0.0" 454 + source = "registry+https://github.com/rust-lang/crates.io-index" 455 + checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" 456 + 457 + [[package]] 439 458 name = "encoding_rs" 440 459 version = "0.8.35" 441 460 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 1035 1054 ] 1036 1055 1037 1056 [[package]] 1057 + name = "indicatif" 1058 + version = "0.17.11" 1059 + source = "registry+https://github.com/rust-lang/crates.io-index" 1060 + checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235" 1061 + dependencies = [ 1062 + "console", 1063 + "number_prefix", 1064 + "portable-atomic", 1065 + "unicode-width", 1066 + "web-time", 1067 + ] 1068 + 1069 + [[package]] 1038 1070 name = "interpolate_name" 1039 1071 version = "0.2.4" 1040 1072 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 1534 1566 ] 1535 1567 1536 1568 [[package]] 1569 + name = "number_prefix" 1570 + version = "0.4.0" 1571 + source = "registry+https://github.com/rust-lang/crates.io-index" 1572 + checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" 1573 + 1574 + [[package]] 1537 1575 name = "once_cell" 1538 1576 version = "1.21.3" 1539 1577 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 2239 2277 2240 2278 [[package]] 2241 2279 name = "spatial-maker" 2242 - version = "0.3.3" 2280 + version = "0.4.0" 2243 2281 dependencies = [ 2244 2282 "clap", 2283 + "console", 2245 2284 "dirs", 2246 2285 "flate2", 2247 2286 "futures-util", 2248 2287 "image", 2288 + "indicatif", 2249 2289 "jxl-oxide", 2250 2290 "libheif-rs", 2251 2291 "ndarray 0.16.1", ··· 2626 2666 checksum = "537dd038a89878be9b64dd4bd1b260315c1bb94f4d784956b81e27a088d9a09e" 2627 2667 2628 2668 [[package]] 2669 + name = "unicode-width" 2670 + version = "0.2.2" 2671 + source = "registry+https://github.com/rust-lang/crates.io-index" 2672 + checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" 2673 + 2674 + [[package]] 2629 2675 name = "unicode-xid" 2630 2676 version = "0.2.6" 2631 2677 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 2880 2926 ] 2881 2927 2882 2928 [[package]] 2929 + name = "web-time" 2930 + version = "1.1.0" 2931 + source = "registry+https://github.com/rust-lang/crates.io-index" 2932 + checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" 2933 + dependencies = [ 2934 + "js-sys", 2935 + "wasm-bindgen", 2936 + ] 2937 + 2938 + [[package]] 2883 2939 name = "webpki-root-certs" 2884 2940 version = "1.0.6" 2885 2941 source = "registry+https://github.com/rust-lang/crates.io-index" ··· 2974 3030 version = "0.52.0" 2975 3031 source = "registry+https://github.com/rust-lang/crates.io-index" 2976 3032 checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" 3033 + dependencies = [ 3034 + "windows-targets 0.52.6", 3035 + ] 3036 + 3037 + [[package]] 3038 + name = "windows-sys" 3039 + version = "0.59.0" 3040 + source = "registry+https://github.com/rust-lang/crates.io-index" 3041 + checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" 2977 3042 dependencies = [ 2978 3043 "windows-targets 0.52.6", 2979 3044 ]
+10 -2
Cargo.toml
··· 1 1 [package] 2 2 name = "spatial-maker" 3 - version = "0.3.3" 3 + version = "0.4.0" 4 4 edition = "2021" 5 5 license = "MIT" 6 6 description = "Convert 2D images and videos to stereoscopic 3D spatial content for Apple Vision Pro using AI depth estimation" ··· 42 42 default = ["coreml", "cli"] 43 43 coreml = [] 44 44 onnx = ["dep:ort"] 45 - cli = ["dep:clap"] 45 + cli = ["dep:clap", "dep:indicatif", "dep:console"] 46 46 47 47 # Native format decoders (optional - falls back to ffmpeg if not enabled) 48 48 avif = ["image/avif"] ··· 74 74 [dependencies.clap] 75 75 version = "4" 76 76 features = ["derive"] 77 + optional = true 78 + 79 + [dependencies.indicatif] 80 + version = "0.17" 81 + optional = true 82 + 83 + [dependencies.console] 84 + version = "0.15" 77 85 optional = true 78 86 79 87 [dev-dependencies]
+38 -12
build.rs
··· 10 10 let out_dir = env::var("OUT_DIR").unwrap(); 11 11 let lib_path = format!("{}/libCoreMLDepth.a", out_dir); 12 12 13 - let status = Command::new("swiftc") 14 - .args(&[ 15 - "-emit-library", 16 - "-static", 17 - "-module-name", 18 - "CoreMLDepth", 19 - "-O", 20 - "swift-bridge/CoreMLDepth.swift", 21 - "-o", 22 - &lib_path, 23 - ]) 13 + let target = env::var("TARGET").unwrap(); 14 + let swift_target = match target.as_str() { 15 + "aarch64-apple-darwin" => "arm64-apple-macosx13.0", 16 + "x86_64-apple-darwin" => "x86_64-apple-macosx13.0", 17 + _ => "", 18 + }; 19 + 20 + let mut cmd = Command::new("swiftc"); 21 + cmd.args(&[ 22 + "-emit-library", 23 + "-static", 24 + "-module-name", 25 + "CoreMLDepth", 26 + "-O", 27 + "swift-bridge/CoreMLDepth.swift", 28 + "-o", 29 + &lib_path, 30 + ]); 31 + if !swift_target.is_empty() { 32 + cmd.args(&["-target", swift_target]); 33 + } 34 + let status = cmd 24 35 .status() 25 36 .expect("Failed to compile Swift code. Is Xcode Command Line Tools installed?"); 26 37 ··· 44 55 .stdout, 45 56 ) 46 57 .unwrap(); 47 - 48 58 let sdk_path = sdk_path.trim(); 49 59 println!("cargo:rustc-link-search=native={}/usr/lib/swift", sdk_path); 50 60 println!("cargo:rustc-link-search=native=/usr/lib/swift"); 61 + 62 + let toolchain_path = String::from_utf8( 63 + Command::new("xcrun") 64 + .args(&["--toolchain", "default", "--find", "swiftc"]) 65 + .output() 66 + .expect("Failed to find Swift toolchain") 67 + .stdout, 68 + ) 69 + .unwrap(); 70 + if let Some(lib_dir) = std::path::Path::new(toolchain_path.trim()) 71 + .ancestors() 72 + .nth(2) 73 + .map(|p| p.join("lib/swift/macosx")) 74 + { 75 + println!("cargo:rustc-link-search=native={}", lib_dir.display()); 76 + } 51 77 } 52 78 } 53 79 }
+30 -3
src/depth_coreml.rs
··· 40 40 Ok(Self { model }) 41 41 } 42 42 43 - pub fn estimate_raw(&self, image: &DynamicImage) -> SpatialResult<ImageBuffer<Luma<f32>, Vec<f32>>> { 44 - let (orig_width, orig_height) = (image.width(), image.height()); 45 - 43 + fn infer_raw(&self, image: &DynamicImage) -> SpatialResult<Vec<f32>> { 46 44 let resized = image.resize_exact( 47 45 INPUT_SIZE, 48 46 INPUT_SIZE, ··· 71 69 result 72 70 ))); 73 71 } 72 + 73 + Ok(output_data) 74 + } 75 + 76 + pub fn estimate_unnormalized(&self, image: &DynamicImage) -> SpatialResult<Array2<f32>> { 77 + let (orig_width, orig_height) = (image.width(), image.height()); 78 + let output_data = self.infer_raw(image)?; 79 + 80 + let depth_image = ImageBuffer::from_fn(INPUT_SIZE, INPUT_SIZE, |x, y| { 81 + let idx = (y * INPUT_SIZE + x) as usize; 82 + Luma([output_data[idx]]) 83 + }); 84 + 85 + let resized_depth = image::imageops::resize( 86 + &depth_image, 87 + orig_width, 88 + orig_height, 89 + image::imageops::FilterType::Lanczos3, 90 + ); 91 + 92 + let (w, h) = resized_depth.dimensions(); 93 + let data: Vec<f32> = resized_depth.pixels().map(|p| p[0]).collect(); 94 + Array2::from_shape_vec((h as usize, w as usize), data) 95 + .map_err(|e| SpatialError::TensorError(format!("Failed to reshape depth: {}", e))) 96 + } 97 + 98 + pub fn estimate_raw(&self, image: &DynamicImage) -> SpatialResult<ImageBuffer<Luma<f32>, Vec<f32>>> { 99 + let (orig_width, orig_height) = (image.width(), image.height()); 100 + let mut output_data = self.infer_raw(image)?; 74 101 75 102 let min_val = output_data.iter().copied().fold(f32::INFINITY, f32::min); 76 103 let max_val = output_data
+206
src/depth_filter.rs
··· 1 + use crate::NormalizeMode; 2 + use ndarray::Array2; 3 + 4 + pub struct DepthProcessor { 5 + prev_depth: Option<Array2<f32>>, 6 + ema_min: f32, 7 + ema_max: f32, 8 + global_min: f32, 9 + global_max: f32, 10 + temporal_alpha: f32, 11 + bilateral_sigma_space: f32, 12 + bilateral_sigma_color: f32, 13 + depth_blur_sigma: f32, 14 + normalize_mode: NormalizeMode, 15 + frame_index: u32, 16 + } 17 + 18 + impl DepthProcessor { 19 + pub fn new( 20 + temporal_alpha: f32, 21 + bilateral_sigma_space: f32, 22 + bilateral_sigma_color: f32, 23 + depth_blur_sigma: f32, 24 + normalize_mode: NormalizeMode, 25 + ) -> Self { 26 + Self { 27 + prev_depth: None, 28 + ema_min: 0.0, 29 + ema_max: 0.0, 30 + global_min: f32::INFINITY, 31 + global_max: f32::NEG_INFINITY, 32 + temporal_alpha, 33 + bilateral_sigma_space, 34 + bilateral_sigma_color, 35 + depth_blur_sigma, 36 + normalize_mode, 37 + frame_index: 0, 38 + } 39 + } 40 + 41 + pub fn set_global_range(&mut self, min: f32, max: f32) { 42 + self.global_min = min; 43 + self.global_max = max; 44 + } 45 + 46 + pub fn update_global_range(&mut self, raw_depth: &Array2<f32>) { 47 + let min = raw_depth.iter().copied().fold(f32::INFINITY, f32::min); 48 + let max = raw_depth.iter().copied().fold(f32::NEG_INFINITY, f32::max); 49 + self.global_min = self.global_min.min(min); 50 + self.global_max = self.global_max.max(max); 51 + } 52 + 53 + pub fn process(&mut self, raw_depth: Array2<f32>) -> Array2<f32> { 54 + let mut depth = self.normalize(raw_depth); 55 + 56 + if self.bilateral_sigma_space > 0.0 { 57 + depth = bilateral_filter(&depth, self.bilateral_sigma_space, self.bilateral_sigma_color); 58 + } 59 + 60 + if self.depth_blur_sigma > 0.0 { 61 + depth = gaussian_blur(&depth, self.depth_blur_sigma); 62 + } 63 + 64 + if self.temporal_alpha > 0.0 && self.temporal_alpha < 1.0 { 65 + if let Some(ref prev) = self.prev_depth { 66 + if prev.dim() == depth.dim() { 67 + let alpha = self.temporal_alpha; 68 + depth.zip_mut_with(prev, |curr, &prev_val| { 69 + *curr = alpha * *curr + (1.0 - alpha) * prev_val; 70 + }); 71 + } 72 + } 73 + self.prev_depth = Some(depth.clone()); 74 + } 75 + 76 + self.frame_index += 1; 77 + depth 78 + } 79 + 80 + fn normalize(&mut self, raw: Array2<f32>) -> Array2<f32> { 81 + match self.normalize_mode { 82 + NormalizeMode::PerFrame => normalize_minmax(raw), 83 + NormalizeMode::RunningEMA => { 84 + let min = raw.iter().copied().fold(f32::INFINITY, f32::min); 85 + let max = raw.iter().copied().fold(f32::NEG_INFINITY, f32::max); 86 + 87 + let adapt_rate = 0.05; 88 + if self.frame_index == 0 { 89 + self.ema_min = min; 90 + self.ema_max = max; 91 + } else { 92 + self.ema_min = self.ema_min + adapt_rate * (min - self.ema_min); 93 + self.ema_max = self.ema_max + adapt_rate * (max - self.ema_max); 94 + } 95 + 96 + let range = self.ema_max - self.ema_min; 97 + if range > 1e-6 { 98 + raw.mapv(|v| ((v - self.ema_min) / range).clamp(0.0, 1.0)) 99 + } else { 100 + raw.mapv(|_| 0.5) 101 + } 102 + } 103 + NormalizeMode::Global => { 104 + let range = self.global_max - self.global_min; 105 + if range > 1e-6 { 106 + raw.mapv(|v| ((v - self.global_min) / range).clamp(0.0, 1.0)) 107 + } else { 108 + raw.mapv(|_| 0.5) 109 + } 110 + } 111 + } 112 + } 113 + } 114 + 115 + fn normalize_minmax(mut depth: Array2<f32>) -> Array2<f32> { 116 + let min = depth.iter().copied().fold(f32::INFINITY, f32::min); 117 + let max = depth.iter().copied().fold(f32::NEG_INFINITY, f32::max); 118 + let range = max - min; 119 + if range > 1e-6 { 120 + depth.mapv_inplace(|v| (v - min) / range); 121 + } 122 + depth 123 + } 124 + 125 + pub fn bilateral_filter(depth: &Array2<f32>, sigma_space: f32, sigma_color: f32) -> Array2<f32> { 126 + let (h, w) = depth.dim(); 127 + let mut out = Array2::zeros((h, w)); 128 + let radius = (sigma_space * 2.0).ceil() as i32; 129 + let space_coeff = -0.5 / (sigma_space * sigma_space); 130 + let color_coeff = -0.5 / (sigma_color * sigma_color); 131 + 132 + for y in 0..h { 133 + for x in 0..w { 134 + let center = depth[[y, x]]; 135 + let mut sum = 0.0f32; 136 + let mut weight_sum = 0.0f32; 137 + 138 + let y0 = (y as i32 - radius).max(0) as usize; 139 + let y1 = (y as i32 + radius).min(h as i32 - 1) as usize; 140 + let x0 = (x as i32 - radius).max(0) as usize; 141 + let x1 = (x as i32 + radius).min(w as i32 - 1) as usize; 142 + 143 + for ny in y0..=y1 { 144 + for nx in x0..=x1 { 145 + let dy = ny as f32 - y as f32; 146 + let dx = nx as f32 - x as f32; 147 + let spatial_dist = dx * dx + dy * dy; 148 + let val = depth[[ny, nx]]; 149 + let color_dist = (val - center) * (val - center); 150 + 151 + let weight = (spatial_dist * space_coeff + color_dist * color_coeff).exp(); 152 + sum += val * weight; 153 + weight_sum += weight; 154 + } 155 + } 156 + 157 + out[[y, x]] = if weight_sum > 0.0 { sum / weight_sum } else { center }; 158 + } 159 + } 160 + 161 + out 162 + } 163 + 164 + pub fn gaussian_blur(depth: &Array2<f32>, sigma: f32) -> Array2<f32> { 165 + let radius = (sigma * 3.0).ceil() as i32; 166 + let kernel_size = (2 * radius + 1) as usize; 167 + let mut kernel = vec![0.0f32; kernel_size]; 168 + let coeff = -0.5 / (sigma * sigma); 169 + 170 + for i in 0..kernel_size { 171 + let d = i as f32 - radius as f32; 172 + kernel[i] = (d * d * coeff).exp(); 173 + } 174 + let ksum: f32 = kernel.iter().sum(); 175 + for v in &mut kernel { 176 + *v /= ksum; 177 + } 178 + 179 + let (h, w) = depth.dim(); 180 + 181 + let mut temp = Array2::zeros((h, w)); 182 + for y in 0..h { 183 + for x in 0..w { 184 + let mut sum = 0.0f32; 185 + for i in 0..kernel_size { 186 + let nx = (x as i32 + i as i32 - radius).clamp(0, w as i32 - 1) as usize; 187 + sum += depth[[y, nx]] * kernel[i]; 188 + } 189 + temp[[y, x]] = sum; 190 + } 191 + } 192 + 193 + let mut out = Array2::zeros((h, w)); 194 + for y in 0..h { 195 + for x in 0..w { 196 + let mut sum = 0.0f32; 197 + for i in 0..kernel_size { 198 + let ny = (y as i32 + i as i32 - radius).clamp(0, h as i32 - 1) as usize; 199 + sum += temp[[ny, x]] * kernel[i]; 200 + } 201 + out[[y, x]] = sum; 202 + } 203 + } 204 + 205 + out 206 + }
+92 -7
src/lib.rs
··· 1 1 pub mod depth; 2 + pub mod depth_filter; 2 3 pub mod error; 3 4 pub mod image_loader; 4 5 pub mod model; ··· 9 10 #[cfg(all(target_os = "macos", feature = "coreml"))] 10 11 pub mod depth_coreml; 11 12 13 + pub use depth_filter::DepthProcessor; 12 14 pub use error::{SpatialError, SpatialResult}; 13 15 pub use image_loader::load_image; 14 16 pub use model::{find_model, get_checkpoint_dir, model_exists}; 15 - pub use output::{create_sbs_image, save_stereo_image, ImageEncoding, MVHEVCConfig, OutputFormat, OutputOptions}; 17 + pub use output::{ 18 + create_sbs_image, save_stereo_image, 19 + DepthFormat, ImageEncoding, MVHEVCConfig, OutputFormat, OutputOptions, OutputType, 20 + depth_formats, needs_depth, needs_stereo, parse_output_types, save_depth_map, stereo_types, 21 + }; 16 22 pub use stereo::generate_stereo_pair; 17 23 pub use video::{get_video_metadata, process_video, ProgressCallback, VideoMetadata, VideoProgress}; 18 24 ··· 25 31 use std::path::Path; 26 32 27 33 #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] 34 + pub enum NormalizeMode { 35 + PerFrame, 36 + RunningEMA, 37 + Global, 38 + } 39 + 40 + impl Default for NormalizeMode { 41 + fn default() -> Self { 42 + Self::RunningEMA 43 + } 44 + } 45 + 46 + impl std::fmt::Display for NormalizeMode { 47 + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 48 + match self { 49 + Self::PerFrame => write!(f, "per-frame"), 50 + Self::RunningEMA => write!(f, "running"), 51 + Self::Global => write!(f, "global"), 52 + } 53 + } 54 + } 55 + 56 + impl std::str::FromStr for NormalizeMode { 57 + type Err = String; 58 + fn from_str(s: &str) -> Result<Self, Self::Err> { 59 + match s.to_lowercase().as_str() { 60 + "per-frame" | "perframe" | "frame" => Ok(Self::PerFrame), 61 + "running" | "ema" | "running-ema" => Ok(Self::RunningEMA), 62 + "global" | "two-pass" | "twopass" => Ok(Self::Global), 63 + _ => Err(format!("Unknown normalize mode: '{}'. Use: per-frame, running, global", s)), 64 + } 65 + } 66 + } 67 + 68 + #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] 28 69 pub struct SpatialConfig { 29 70 pub encoder_size: String, 30 71 pub max_disparity: u32, 31 72 pub target_depth_size: u32, 73 + pub temporal_alpha: f32, 74 + pub bilateral_sigma_space: f32, 75 + pub bilateral_sigma_color: f32, 76 + pub depth_blur_sigma: f32, 77 + pub normalize_mode: NormalizeMode, 32 78 } 33 79 34 80 pub type StereoOutputFormat = OutputFormat; ··· 39 85 encoder_size: "s".to_string(), 40 86 max_disparity: 30, 41 87 target_depth_size: 518, 88 + temporal_alpha: 0.7, 89 + bilateral_sigma_space: 5.0, 90 + bilateral_sigma_color: 0.1, 91 + depth_blur_sigma: 1.5, 92 + normalize_mode: NormalizeMode::RunningEMA, 42 93 } 43 94 } 44 95 } 45 96 97 + pub struct ProcessPhotoOutput { 98 + pub depth_paths: Vec<std::path::PathBuf>, 99 + pub stereo_paths: Vec<std::path::PathBuf>, 100 + } 101 + 46 102 pub async fn process_photo( 47 103 input_path: &Path, 48 - output_path: &Path, 104 + output_base_path: &Path, 49 105 config: SpatialConfig, 106 + output_types: &[OutputType], 50 107 output_options: OutputOptions, 51 - ) -> SpatialResult<()> { 108 + ) -> SpatialResult<ProcessPhotoOutput> { 52 109 let input_image = load_image(input_path).await?; 53 110 54 111 model::ensure_model_exists::<fn(u64, u64)>(&config.encoder_size, None).await?; ··· 79 136 } 80 137 }; 81 138 82 - let (left, right) = generate_stereo_pair(&input_image, &depth_map, config.max_disparity)?; 139 + let mut result = ProcessPhotoOutput { 140 + depth_paths: Vec::new(), 141 + stereo_paths: Vec::new(), 142 + }; 143 + 144 + if needs_depth(output_types) { 145 + let parent = output_base_path.parent().unwrap_or_else(|| Path::new(".")); 146 + let stem = output_base_path.file_stem().and_then(|s| s.to_str()).unwrap_or("output"); 147 + 148 + for fmt in depth_formats(output_types) { 149 + let filename = format!("{}-depth{}.{}", stem, fmt.suffix(), fmt.extension()); 150 + let depth_path = parent.join(&filename); 151 + save_depth_map(&depth_map, &depth_path, fmt)?; 152 + result.depth_paths.push(depth_path.clone()); 153 + } 154 + } 83 155 84 - save_stereo_image(&left, &right, output_path, output_options)?; 156 + if needs_stereo(output_types) { 157 + let (left, right) = generate_stereo_pair(&input_image, &depth_map, config.max_disparity)?; 158 + let src_ext = input_path.extension().and_then(|s| s.to_str()).unwrap_or("").to_lowercase(); 159 + let stereo_ext = match src_ext.as_str() { 160 + "heic" | "heif" | "avif" | "jxl" => "jpg", 161 + "" => "jpg", 162 + other => other, 163 + }; 164 + let parent = output_base_path.parent().unwrap_or_else(|| Path::new(".")); 165 + let stem = output_base_path.file_stem().and_then(|s| s.to_str()).unwrap_or("output"); 166 + let stereo_path = parent.join(format!("{}-spatial.{}", stem, stereo_ext)); 167 + save_stereo_image(&left, &right, &stereo_path, output_options)?; 168 + result.stereo_paths.push(stereo_path); 169 + } 85 170 86 - Ok(()) 171 + Ok(result) 87 172 } 88 173 89 174 pub async fn process_video_sbs( ··· 92 177 config: SpatialConfig, 93 178 progress_cb: Option<ProgressCallback>, 94 179 ) -> SpatialResult<()> { 95 - video::process_video(input_path, output_path, config, progress_cb).await 180 + video::process_video(input_path, output_path, config, &[OutputType::Spatial], progress_cb).await 96 181 }
+246 -65
src/main.rs
··· 1 1 use clap::{Parser, Subcommand}; 2 + use console::style; 3 + use indicatif::{ProgressBar, ProgressStyle}; 2 4 use spatial_maker::{ 3 - process_photo, process_video, ImageEncoding, MVHEVCConfig, OutputFormat, OutputOptions, 4 - SpatialConfig, VideoProgress, 5 + process_photo, process_video, ImageEncoding, MVHEVCConfig, NormalizeMode, OutputFormat, 6 + OutputOptions, OutputType, SpatialConfig, VideoProgress, 7 + needs_stereo, parse_output_types, 5 8 }; 6 9 use std::path::PathBuf; 7 10 ··· 33 36 #[arg(long, default_value = "30")] 34 37 max_disparity: u32, 35 38 36 - /// Output format for photos: sbs (side-by-side), tab (top-and-bottom), sep (separate L/R) 37 - #[arg(long, default_value = "sbs")] 38 - format: String, 39 + /// Output types (comma-separated): depth, depth:avif,png,png16, sbs, tab, sep, spatial 40 + #[arg(long, default_value = "spatial")] 41 + output_types: String, 39 42 40 43 /// JPEG quality for photos (1-100) 41 44 #[arg(long, default_value = "95")] 42 45 quality: u8, 43 46 44 - /// Enable MV-HEVC packaging for photos (requires 'spatial' CLI in PATH) 45 - #[arg(long)] 46 - mvhevc: bool, 47 + /// Temporal EMA blend factor for video depth (0=off, 1=no smoothing, default 0.7) 48 + #[arg(long, default_value = "0.7")] 49 + temporal_alpha: f32, 50 + 51 + /// Bilateral filter spatial sigma (0=off, default 5.0) 52 + #[arg(long, default_value = "5.0")] 53 + bilateral_sigma: f32, 54 + 55 + /// Bilateral filter range sigma (default 0.1) 56 + #[arg(long, default_value = "0.1")] 57 + bilateral_range: f32, 58 + 59 + /// Gaussian blur sigma for depth edge softening (0=off, default 1.5) 60 + #[arg(long, default_value = "1.5")] 61 + depth_blur: f32, 62 + 63 + /// Depth normalization mode for video: running (default), per-frame, global (two-pass) 64 + #[arg(long, default_value = "running")] 65 + normalize: String, 47 66 } 48 67 49 68 #[derive(Subcommand)] ··· 82 101 } 83 102 } 84 103 85 - fn generate_output_path(input: &PathBuf, media_type: &MediaType) -> PathBuf { 104 + fn generate_output_base(input: &PathBuf, model: &str) -> PathBuf { 86 105 let stem = input.file_stem().and_then(|s| s.to_str()).unwrap_or("output"); 87 - let src_ext = input 88 - .extension() 89 - .and_then(|s| s.to_str()) 90 - .unwrap_or("") 91 - .to_lowercase(); 106 + let parent = input.parent().unwrap_or_else(|| std::path::Path::new(".")); 107 + parent.join(format!("{}-{}", stem, model)) 108 + } 92 109 93 - let extension = match media_type { 94 - MediaType::Video => "mov", 95 - MediaType::Photo => match src_ext.as_str() { 96 - "heic" | "heif" | "avif" | "jxl" => "jpg", 97 - "" => "jpg", 98 - _ => &src_ext, 99 - }, 100 - }; 110 + 101 111 102 - let parent = input.parent().unwrap_or_else(|| std::path::Path::new(".")); 103 - parent.join(format!("{}-spatial.{}", stem, extension)) 112 + fn model_display_name(encoder_size: &str) -> (&str, u32) { 113 + match encoder_size { 114 + "s" | "small" => ("small", 48), 115 + "b" | "base" => ("base", 186), 116 + "l" | "large" => ("large", 638), 117 + _ => (encoder_size, 0), 118 + } 104 119 } 105 120 106 121 async fn process_single( 107 122 input: &PathBuf, 108 123 output: PathBuf, 109 124 config: SpatialConfig, 125 + output_types: &[OutputType], 110 126 cli: &Cli, 111 127 ) -> Result<(), Box<dyn std::error::Error>> { 112 128 let media_type = detect_media_type(input); 113 129 114 130 match media_type { 115 131 MediaType::Photo => { 116 - let layout = match cli.format.as_str() { 117 - "sbs" => OutputFormat::SideBySide, 118 - "tab" => OutputFormat::TopAndBottom, 119 - "sep" => OutputFormat::Separate, 120 - _ => { 121 - eprintln!("Invalid format '{}'. Use: sbs, tab, or sep", cli.format); 122 - std::process::exit(1); 132 + let has_stereo = needs_stereo(output_types); 133 + 134 + let layout = if has_stereo { 135 + let stereo = spatial_maker::stereo_types(output_types); 136 + match stereo.first() { 137 + Some(OutputType::TopAndBottom) => OutputFormat::TopAndBottom, 138 + Some(OutputType::Separate) => OutputFormat::Separate, 139 + _ => OutputFormat::SideBySide, 123 140 } 141 + } else { 142 + OutputFormat::SideBySide 124 143 }; 144 + 145 + let has_spatial = output_types.iter().any(|t| matches!(t, OutputType::Spatial)); 125 146 126 147 let output_options = OutputOptions { 127 148 layout, 128 149 image_format: ImageEncoding::Jpeg { 129 150 quality: cli.quality, 130 151 }, 131 - mvhevc: if cli.mvhevc { 152 + mvhevc: if has_spatial { 132 153 Some(MVHEVCConfig { 133 154 spatial_cli_path: None, 134 155 enabled: true, 135 156 quality: cli.quality, 136 - keep_intermediate: false, 157 + keep_intermediate: has_stereo && output_types.iter().any(|t| matches!(t, OutputType::SideBySide | OutputType::TopAndBottom | OutputType::Separate)), 137 158 }) 138 159 } else { 139 160 None 140 161 }, 141 162 }; 142 163 143 - eprintln!("Processing photo: {:?}", input); 144 - process_photo(input, &output, config, output_options).await?; 145 - eprintln!("Saved to: {:?}", output); 164 + let filename = input.file_name().and_then(|s| s.to_str()).unwrap_or("?"); 165 + eprintln!("{} {}", style("🖼").cyan(), style(filename).bold()); 166 + 167 + let spinner = ProgressBar::new_spinner(); 168 + spinner.set_style( 169 + ProgressStyle::default_spinner() 170 + .template("{spinner:.cyan} {msg}") 171 + .unwrap() 172 + ); 173 + spinner.set_message("Processing..."); 174 + spinner.enable_steady_tick(std::time::Duration::from_millis(80)); 175 + 176 + let result = process_photo(input, &output, config.clone(), output_types, output_options).await?; 177 + 178 + spinner.finish_and_clear(); 179 + let (model_name, model_mb) = model_display_name(&cli.model); 180 + eprintln!( 181 + "{} {} / {} MB / depth-anything-v2-{}", 182 + style("✔").green().bold(), 183 + style("done").green(), 184 + model_mb, 185 + model_name, 186 + ); 187 + 188 + for path in &result.depth_paths { 189 + let name = path.file_name().and_then(|s| s.to_str()).unwrap_or("?"); 190 + eprintln!("{} {}", style("→").dim(), style(name).dim()); 191 + } 192 + for path in &result.stereo_paths { 193 + let name = path.file_name().and_then(|s| s.to_str()).unwrap_or("?"); 194 + eprintln!("{} {}", style("→").dim(), style(name).dim()); 195 + } 146 196 } 147 197 MediaType::Video => { 148 - eprintln!("Processing video: {:?}", input); 198 + let filename = input.file_name().and_then(|s| s.to_str()).unwrap_or("?"); 199 + eprintln!("{} {}", style("🎥").cyan(), style(filename).bold()); 200 + 201 + let (model_name, model_mb) = model_display_name(&cli.model); 202 + let model_info = format!("model loaded / {} MB / depth-anything-v2-{}", model_mb, model_name); 203 + 204 + let spinner = ProgressBar::new_spinner(); 205 + spinner.set_style( 206 + ProgressStyle::default_spinner() 207 + .template("{spinner:.cyan} {msg}") 208 + .unwrap() 209 + ); 210 + spinner.set_message("Loading model..."); 211 + spinner.enable_steady_tick(std::time::Duration::from_millis(80)); 212 + 149 213 let start = std::time::Instant::now(); 150 214 215 + let pb = ProgressBar::new(100); 216 + pb.set_style( 217 + ProgressStyle::default_bar() 218 + .template("{msg} {bar:20.cyan/blue} {pos:>3}% {prefix}") 219 + .unwrap() 220 + .progress_chars("━╸─") 221 + ); 222 + 223 + let pb_inner = pb.clone(); 224 + let start_clone = start.clone(); 225 + let model_loaded = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); 226 + let model_loaded_clone = model_loaded.clone(); 227 + let spinner_clone = spinner.clone(); 228 + 151 229 process_video( 152 230 input, 153 231 &output, 154 232 config, 155 - Some(Box::new(|progress: VideoProgress| { 156 - eprint!( 157 - "\r[{}] Frame {}/{} ({:.1}%)", 158 - progress.stage, progress.current_frame, progress.total_frames, progress.percent 159 - ); 233 + output_types, 234 + Some(Box::new(move |progress: VideoProgress| { 235 + if !model_loaded_clone.load(std::sync::atomic::Ordering::Relaxed) { 236 + model_loaded_clone.store(true, std::sync::atomic::Ordering::Relaxed); 237 + spinner_clone.finish_and_clear(); 238 + eprintln!( 239 + "{} {}", 240 + style("✔").green().bold(), 241 + model_info, 242 + ); 243 + } 244 + 245 + let elapsed = start_clone.elapsed().as_secs_f64(); 246 + let fps = if elapsed > 0.1 && progress.current_frame > 0 { 247 + progress.current_frame as f64 / elapsed 248 + } else { 249 + 0.0 250 + }; 251 + 252 + match progress.stage.as_str() { 253 + "scanning" => { 254 + pb_inner.set_message(format!("{}", style("scanning depths").yellow())); 255 + let pct = if progress.total_frames > 0 { 256 + (progress.current_frame as f64 / progress.total_frames as f64 * 100.0) as u64 257 + } else { 258 + 0 259 + }; 260 + pb_inner.set_position(pct); 261 + pb_inner.set_prefix(format!("{}/{}", progress.current_frame, progress.total_frames)); 262 + } 263 + "extracting" => { 264 + pb_inner.set_message(format!("{}", style("loading").yellow())); 265 + pb_inner.set_position(0); 266 + } 267 + "processing" => { 268 + let eta_secs = if fps > 0.1 { 269 + let remaining = progress.total_frames.saturating_sub(progress.current_frame); 270 + remaining as f64 / fps 271 + } else { 272 + 0.0 273 + }; 274 + 275 + let eta_str = if eta_secs > 60.0 { 276 + format!("{}m {:02}s", eta_secs as u64 / 60, eta_secs as u64 % 60) 277 + } else { 278 + format!("{:.0}s", eta_secs) 279 + }; 280 + 281 + pb_inner.set_message(format!( 282 + "{:>5.1} fps", 283 + fps, 284 + )); 285 + pb_inner.set_prefix(format!( 286 + "{}/{} eta {}", 287 + progress.current_frame, 288 + progress.total_frames, 289 + eta_str, 290 + )); 291 + pb_inner.set_position(progress.percent as u64); 292 + } 293 + "encoding" => { 294 + pb_inner.set_message(format!("{}", style("encoding").yellow())); 295 + pb_inner.set_position(100); 296 + pb_inner.set_prefix(String::new()); 297 + } 298 + "packaging" => { 299 + pb_inner.set_message(format!("{}", style("packaging MV-HEVC").yellow())); 300 + } 301 + "complete" => { 302 + pb_inner.finish_and_clear(); 303 + } 304 + _ => {} 305 + } 160 306 })), 161 307 ) 162 308 .await?; 163 309 164 - eprintln!(); 165 - eprintln!("Saved to: {:?}", output); 166 - eprintln!("Total time: {:.1}s", start.elapsed().as_secs_f64()); 310 + pb.finish_and_clear(); 311 + 312 + let elapsed = start.elapsed().as_secs_f64(); 313 + let out_name = output.file_name().and_then(|s| s.to_str()).unwrap_or("?"); 314 + eprintln!( 315 + "{} {} ({:.1}s)", 316 + style("→").green().bold(), 317 + style(out_name).white(), 318 + elapsed, 319 + ); 167 320 } 168 321 } 169 322 ··· 189 342 std::process::exit(1); 190 343 } 191 344 345 + let output_types = parse_output_types(&cli.output_types).unwrap_or_else(|e| { 346 + eprintln!("Invalid --output-types: {}", e); 347 + std::process::exit(1); 348 + }); 349 + 350 + let normalize_mode: NormalizeMode = cli.normalize.parse().unwrap_or_else(|e| { 351 + eprintln!("{}", e); 352 + std::process::exit(1); 353 + }); 354 + 192 355 let config = SpatialConfig { 193 356 encoder_size: cli.model.clone(), 194 357 max_disparity: cli.max_disparity, 195 358 target_depth_size: 518, 359 + temporal_alpha: cli.temporal_alpha, 360 + bilateral_sigma_space: cli.bilateral_sigma, 361 + bilateral_sigma_color: cli.bilateral_range, 362 + depth_blur_sigma: cli.depth_blur, 363 + normalize_mode, 196 364 }; 197 365 198 366 let total = cli.inputs.len(); ··· 200 368 201 369 for (i, input) in cli.inputs.iter().enumerate() { 202 370 if total > 1 { 203 - eprintln!("[{}/{}]", i + 1, total); 371 + eprintln!( 372 + "{} {}/{}", 373 + style("─").dim(), 374 + style(i + 1).bold(), 375 + style(total).dim(), 376 + ); 204 377 } 205 378 206 - let media_type = detect_media_type(input); 207 379 let output = cli 208 380 .output 209 381 .clone() 210 - .unwrap_or_else(|| generate_output_path(input, &media_type)); 382 + .unwrap_or_else(|| generate_output_base(input, &cli.model)); 211 383 212 - if let Err(e) = process_single(input, output, config.clone(), &cli).await { 213 - eprintln!("Error processing {:?}: {}", input, e); 384 + if let Err(e) = process_single(input, output, config.clone(), &output_types, &cli).await { 385 + eprintln!("{} {:?}: {}", style("✗").red().bold(), input, e); 214 386 errors.push((input.clone(), e)); 215 387 } 216 388 } 217 389 218 390 if !errors.is_empty() { 219 - eprintln!("\n{}/{} files failed", errors.len(), total); 391 + eprintln!( 392 + "\n{} {}/{} files failed", 393 + style("✗").red().bold(), 394 + errors.len(), 395 + total, 396 + ); 220 397 std::process::exit(1); 221 398 } 222 399 ··· 232 409 233 410 let client = reqwest::Client::new(); 234 411 235 - // Get latest release from GitHub API 236 412 let release: serde_json::Value = client 237 413 .get(format!("https://api.github.com/repos/{}/releases/latest", repo)) 238 414 .header("User-Agent", "spatial-maker") ··· 253 429 254 430 eprintln!("New version available: v{} -> v{}", current_version, latest_version); 255 431 256 - // Determine current arch 257 432 let target = if cfg!(target_arch = "aarch64") { 258 433 "aarch64-apple-darwin" 259 434 } else if cfg!(target_arch = "x86_64") { ··· 264 439 265 440 let asset_name = format!("spatial-maker-{}-{}.tar.gz", latest_tag, target); 266 441 267 - // Find the download URL from release assets 268 442 let assets = release["assets"] 269 443 .as_array() 270 444 .ok_or("No assets in release")?; ··· 275 449 .and_then(|a| a["browser_download_url"].as_str()) 276 450 .ok_or_else(|| format!("No release asset found for {}", target))?; 277 451 278 - eprintln!("Downloading {}...", asset_name); 452 + let pb = ProgressBar::new_spinner(); 453 + pb.set_style( 454 + ProgressStyle::default_spinner() 455 + .template("{spinner:.cyan} {msg}") 456 + .unwrap() 457 + ); 458 + pb.set_message(format!("Downloading {}...", asset_name)); 459 + pb.enable_steady_tick(std::time::Duration::from_millis(80)); 279 460 280 - // Download to a temp file 281 461 let response = client 282 462 .get(download_url) 283 463 .header("User-Agent", "spatial-maker") ··· 286 466 287 467 let bytes = response.bytes().await?; 288 468 289 - // Extract the binary from the tarball 469 + pb.set_message("Extracting..."); 470 + 290 471 let decoder = flate2::read::GzDecoder::new(&bytes[..]); 291 472 let mut archive = tar::Archive::new(decoder); 292 473 ··· 298 479 299 480 let new_binary = temp_dir.join("spatial-maker"); 300 481 if !new_binary.exists() { 482 + pb.finish_and_clear(); 301 483 return Err("Binary not found in release archive".into()); 302 484 } 303 485 304 - // Determine where to install 305 486 let current_exe = std::env::current_exe()?; 306 487 let install_path = if is_writable(&current_exe) { 307 488 current_exe.clone() 308 489 } else { 309 - // Try ~/.local/bin as a user-writable alternative 310 490 let local_bin = dirs::home_dir() 311 491 .ok_or("Could not determine home directory")? 312 492 .join(".local/bin"); ··· 320 500 alt_path 321 501 }; 322 502 323 - // Replace the binary atomically: copy new -> rename over old 324 503 let staging = install_path.with_extension("new"); 325 504 std::fs::copy(&new_binary, &staging)?; 326 505 ··· 332 511 333 512 std::fs::rename(&staging, &install_path)?; 334 513 335 - // Cleanup 336 514 let _ = std::fs::remove_dir_all(&temp_dir); 337 515 338 - eprintln!("Updated to v{} at {}", latest_version, install_path.display()); 516 + pb.finish_and_clear(); 517 + eprintln!( 518 + "{} Updated to {} at {}", 519 + style("✔").green().bold(), 520 + style(format!("v{}", latest_version)).green(), 521 + install_path.display(), 522 + ); 339 523 340 - // Remind about PATH if installed to ~/.local/bin 341 524 if install_path != current_exe { 342 525 eprintln!("Make sure {} is in your PATH", install_path.parent().unwrap().display()); 343 526 } ··· 356 539 357 540 fn is_writable(path: &PathBuf) -> bool { 358 541 if path.exists() { 359 - // Try opening for write to test permissions 360 542 std::fs::OpenOptions::new() 361 543 .write(true) 362 544 .open(path) 363 545 .is_ok() 364 546 } else { 365 - // Check if parent directory is writable 366 547 path.parent() 367 548 .map(|p| { 368 549 let test = p.join(".spatial-maker-write-test");
+268 -1
src/output.rs
··· 1 1 use crate::error::{SpatialError, SpatialResult}; 2 2 use image::DynamicImage; 3 + use ndarray::Array2; 3 4 use std::path::{Path, PathBuf}; 4 5 use std::process::Command; 5 6 6 7 #[derive(Clone, Copy, Debug, PartialEq, Eq)] 8 + pub enum DepthFormat { 9 + Avif, 10 + Png, 11 + Png16, 12 + } 13 + 14 + impl DepthFormat { 15 + pub fn extension(&self) -> &'static str { 16 + match self { 17 + DepthFormat::Avif => "avif", 18 + DepthFormat::Png => "png", 19 + DepthFormat::Png16 => "png", 20 + } 21 + } 22 + 23 + pub fn suffix(&self) -> &'static str { 24 + match self { 25 + DepthFormat::Avif => "", 26 + DepthFormat::Png => "", 27 + DepthFormat::Png16 => "-16bit", 28 + } 29 + } 30 + } 31 + 32 + pub const DEFAULT_DEPTH_FORMAT: DepthFormat = DepthFormat::Avif; 33 + 34 + #[derive(Clone, Debug, PartialEq, Eq)] 35 + pub enum OutputType { 36 + Depth(Vec<DepthFormat>), 37 + SideBySide, 38 + TopAndBottom, 39 + Separate, 40 + Spatial, 41 + } 42 + 43 + pub fn needs_depth(types: &[OutputType]) -> bool { 44 + types.iter().any(|t| matches!(t, OutputType::Depth(_))) 45 + } 46 + 47 + pub fn needs_stereo(types: &[OutputType]) -> bool { 48 + types.iter().any(|t| matches!(t, OutputType::SideBySide | OutputType::TopAndBottom | OutputType::Separate | OutputType::Spatial)) 49 + } 50 + 51 + pub fn depth_formats(types: &[OutputType]) -> Vec<DepthFormat> { 52 + types.iter().filter_map(|t| { 53 + if let OutputType::Depth(fmts) = t { Some(fmts.clone()) } else { None } 54 + }).flatten().collect() 55 + } 56 + 57 + pub fn stereo_types(types: &[OutputType]) -> Vec<&OutputType> { 58 + types.iter().filter(|t| matches!(t, OutputType::SideBySide | OutputType::TopAndBottom | OutputType::Separate | OutputType::Spatial)).collect() 59 + } 60 + 61 + fn is_depth_format(s: &str) -> bool { 62 + matches!(s, "avif" | "png" | "png16") 63 + } 64 + 65 + fn is_stereo_type(s: &str) -> bool { 66 + matches!(s, "sbs" | "tab" | "sep" | "spatial") 67 + } 68 + 69 + fn parse_depth_format(s: &str) -> Result<DepthFormat, String> { 70 + match s { 71 + "avif" => Ok(DepthFormat::Avif), 72 + "png" => Ok(DepthFormat::Png), 73 + "png16" => Ok(DepthFormat::Png16), 74 + _ => Err(format!("Unknown depth format: '{}'. Use: avif, png, png16", s)), 75 + } 76 + } 77 + 78 + fn parse_stereo_type(s: &str) -> Result<OutputType, String> { 79 + match s { 80 + "sbs" => Ok(OutputType::SideBySide), 81 + "tab" => Ok(OutputType::TopAndBottom), 82 + "sep" => Ok(OutputType::Separate), 83 + "spatial" => Ok(OutputType::Spatial), 84 + _ => Err(format!("Unknown output type: '{}'", s)), 85 + } 86 + } 87 + 88 + pub fn parse_output_types(s: &str) -> Result<Vec<OutputType>, String> { 89 + let parts: Vec<&str> = s.split(',').map(|p| p.trim()).filter(|p| !p.is_empty()).collect(); 90 + let mut types = Vec::new(); 91 + let mut depth_fmts = Vec::new(); 92 + let mut has_depth = false; 93 + 94 + for part in &parts { 95 + if *part == "depth" { 96 + has_depth = true; 97 + continue; 98 + } 99 + 100 + if let Some(after_colon) = part.strip_prefix("depth:") { 101 + has_depth = true; 102 + depth_fmts.push(parse_depth_format(after_colon)?); 103 + continue; 104 + } 105 + 106 + if has_depth && is_depth_format(part) { 107 + depth_fmts.push(parse_depth_format(part)?); 108 + continue; 109 + } 110 + 111 + if is_stereo_type(part) { 112 + types.push(parse_stereo_type(part)?); 113 + } else if is_depth_format(part) { 114 + return Err(format!( 115 + "'{}' must be specified as a depth sub-format: depth:{}", part, part 116 + )); 117 + } else { 118 + return Err(format!("Unknown output type: '{}'", part)); 119 + } 120 + } 121 + 122 + if has_depth { 123 + if depth_fmts.is_empty() { 124 + depth_fmts.push(DEFAULT_DEPTH_FORMAT); 125 + } 126 + types.insert(0, OutputType::Depth(depth_fmts)); 127 + } 128 + 129 + if types.is_empty() { 130 + return Err("No output types specified".to_string()); 131 + } 132 + 133 + Ok(types) 134 + } 135 + 136 + // --- Depth map saving --- 137 + 138 + fn normalize_depth(depth: &Array2<f32>) -> (f32, f32) { 139 + let mut min_val = f32::INFINITY; 140 + let mut max_val = f32::NEG_INFINITY; 141 + for &v in depth.iter() { 142 + if v < min_val { min_val = v; } 143 + if v > max_val { max_val = v; } 144 + } 145 + (min_val, max_val) 146 + } 147 + 148 + pub fn save_depth_png8(depth: &Array2<f32>, path: &Path) -> SpatialResult<()> { 149 + let (h, w) = depth.dim(); 150 + let (min_val, max_val) = normalize_depth(depth); 151 + let range = max_val - min_val; 152 + 153 + let pixels: Vec<u8> = depth.iter().map(|&v| { 154 + if range > 1e-6 { 155 + ((v - min_val) / range * 255.0).round() as u8 156 + } else { 157 + 128u8 158 + } 159 + }).collect(); 160 + 161 + let img = image::GrayImage::from_raw(w as u32, h as u32, pixels) 162 + .ok_or_else(|| SpatialError::ImageError("Failed to create grayscale image".to_string()))?; 163 + 164 + img.save(path) 165 + .map_err(|e| SpatialError::ImageError(format!("Failed to save depth PNG: {}", e)))?; 166 + 167 + Ok(()) 168 + } 169 + 170 + pub fn save_depth_png16(depth: &Array2<f32>, path: &Path) -> SpatialResult<()> { 171 + let (h, w) = depth.dim(); 172 + let (min_val, max_val) = normalize_depth(depth); 173 + let range = max_val - min_val; 174 + 175 + let pixels: Vec<u16> = depth.iter().map(|&v| { 176 + if range > 1e-6 { 177 + ((v - min_val) / range * 65535.0).round() as u16 178 + } else { 179 + 32768u16 180 + } 181 + }).collect(); 182 + 183 + let file = std::fs::File::create(path) 184 + .map_err(|e| SpatialError::ImageError(format!("Failed to create output file: {}", e)))?; 185 + let writer = std::io::BufWriter::new(file); 186 + 187 + let encoder = image::codecs::png::PngEncoder::new(writer); 188 + use image::ImageEncoder; 189 + 190 + let byte_data: Vec<u8> = pixels.iter().flat_map(|&v| v.to_be_bytes()).collect(); 191 + 192 + encoder.write_image( 193 + &byte_data, 194 + w as u32, 195 + h as u32, 196 + image::ExtendedColorType::L16, 197 + ).map_err(|e| SpatialError::ImageError(format!("Failed to encode 16-bit PNG: {}", e)))?; 198 + 199 + Ok(()) 200 + } 201 + 202 + pub fn save_depth_avif(depth: &Array2<f32>, path: &Path) -> SpatialResult<()> { 203 + let (h, w) = depth.dim(); 204 + let (min_val, max_val) = normalize_depth(depth); 205 + let range = max_val - min_val; 206 + 207 + let pixels: Vec<u8> = depth.iter().map(|&v| { 208 + if range > 1e-6 { 209 + ((v - min_val) / range * 255.0).round() as u8 210 + } else { 211 + 128u8 212 + } 213 + }).collect(); 214 + 215 + let rgb_pixels: Vec<u8> = pixels.iter().flat_map(|&v| [v, v, v]).collect(); 216 + 217 + let path_str = path.to_str() 218 + .ok_or_else(|| SpatialError::ImageError("Invalid output path".to_string()))?; 219 + 220 + let mut child = Command::new("ffmpeg") 221 + .args([ 222 + "-f", "rawvideo", 223 + "-pix_fmt", "rgb24", 224 + "-s", &format!("{}x{}", w, h), 225 + "-i", "-", 226 + "-frames:v", "1", 227 + "-c:v", "libsvtav1", 228 + "-crf", "23", 229 + "-y", 230 + path_str, 231 + ]) 232 + .stdin(std::process::Stdio::piped()) 233 + .stdout(std::process::Stdio::null()) 234 + .stderr(std::process::Stdio::piped()) 235 + .spawn() 236 + .map_err(|e| SpatialError::Other(format!("Failed to spawn ffmpeg for AVIF encoding: {}", e)))?; 237 + 238 + if let Some(mut stdin) = child.stdin.take() { 239 + use std::io::Write; 240 + stdin.write_all(&rgb_pixels) 241 + .map_err(|e| SpatialError::IoError(format!("Failed to write depth data to ffmpeg: {}", e)))?; 242 + } 243 + 244 + let output = child.wait_with_output() 245 + .map_err(|e| SpatialError::Other(format!("ffmpeg AVIF encoding failed: {}", e)))?; 246 + 247 + if !output.status.success() { 248 + let stderr = String::from_utf8_lossy(&output.stderr); 249 + return Err(SpatialError::ImageError(format!("ffmpeg AVIF encoding failed: {}", stderr))); 250 + } 251 + 252 + Ok(()) 253 + } 254 + 255 + pub fn save_depth_map(depth: &Array2<f32>, path: &Path, format: DepthFormat) -> SpatialResult<()> { 256 + if let Some(parent) = path.parent() { 257 + std::fs::create_dir_all(parent).map_err(|e| { 258 + SpatialError::ImageError(format!("Failed to create output directory: {}", e)) 259 + })?; 260 + } 261 + 262 + match format { 263 + DepthFormat::Avif => save_depth_avif(depth, path)?, 264 + DepthFormat::Png => save_depth_png8(depth, path)?, 265 + DepthFormat::Png16 => save_depth_png16(depth, path)?, 266 + } 267 + 268 + Ok(()) 269 + } 270 + 271 + // --- Existing stereo output --- 272 + 273 + #[derive(Clone, Copy, Debug, PartialEq, Eq)] 7 274 pub enum OutputFormat { 8 275 SideBySide, 9 276 TopAndBottom, ··· 229 496 Ok(()) 230 497 } 231 498 232 - fn encode_mvhevc(stereo_path: &Path, config: &MVHEVCConfig) -> SpatialResult<()> { 499 + pub fn encode_mvhevc(stereo_path: &Path, config: &MVHEVCConfig) -> SpatialResult<()> { 233 500 let spatial_path = config 234 501 .spatial_cli_path 235 502 .as_ref()
+206 -22
src/video.rs
··· 1 + use crate::depth_filter::DepthProcessor; 1 2 use crate::error::{SpatialError, SpatialResult}; 3 + use crate::output::{needs_depth, needs_stereo, OutputType}; 2 4 use crate::stereo::generate_stereo_pair; 3 - use crate::SpatialConfig; 5 + use crate::{NormalizeMode, SpatialConfig}; 4 6 use image::{DynamicImage, ImageBuffer, RgbImage}; 7 + use ndarray::Array2; 5 8 use std::path::Path; 6 9 use std::process::Stdio; 7 10 use tokio::io::{AsyncReadExt, AsyncWriteExt}; ··· 154 157 155 158 let input_path = input_path.to_path_buf(); 156 159 160 + let vf_scale = format!("scale={}:{}", width, height); 161 + 157 162 tokio::spawn(async move { 158 163 let mut child = Command::new("ffmpeg") 159 164 .args([ 160 165 "-i", 161 166 input_path.to_str().unwrap(), 167 + "-vf", 168 + &vf_scale, 162 169 "-f", 163 170 "rawvideo", 164 171 "-pix_fmt", ··· 288 295 Ok(()) 289 296 } 290 297 298 + async fn encode_depth_video( 299 + output_path: std::path::PathBuf, 300 + metadata: VideoMetadata, 301 + mut rx: mpsc::Receiver<Array2<f32>>, 302 + ) -> SpatialResult<()> { 303 + let width = metadata.width; 304 + let height = metadata.height; 305 + let fps = metadata.fps; 306 + 307 + let mut child = Command::new("ffmpeg") 308 + .args([ 309 + "-f", "rawvideo", 310 + "-pix_fmt", "gray", 311 + "-s", &format!("{}x{}", width, height), 312 + "-r", &format!("{}", fps), 313 + "-i", "-", 314 + "-c:v", "libsvtav1", 315 + "-crf", "23", 316 + "-pix_fmt", "yuv420p", 317 + "-y", 318 + output_path.to_str().unwrap(), 319 + ]) 320 + .stdin(Stdio::piped()) 321 + .stdout(Stdio::null()) 322 + .stderr(Stdio::piped()) 323 + .spawn() 324 + .map_err(|e| SpatialError::Other(format!("Failed to spawn ffmpeg depth encoder: {}", e)))?; 325 + 326 + let mut stdin = child.stdin.take().expect("Failed to capture stdin"); 327 + 328 + while let Some(depth) = rx.recv().await { 329 + let mut min_val = f32::INFINITY; 330 + let mut max_val = f32::NEG_INFINITY; 331 + for &v in depth.iter() { 332 + if v < min_val { min_val = v; } 333 + if v > max_val { max_val = v; } 334 + } 335 + let range = max_val - min_val; 336 + 337 + let pixels: Vec<u8> = depth.iter().map(|&v| { 338 + if range > 1e-6 { 339 + ((v - min_val) / range * 255.0).round() as u8 340 + } else { 341 + 128u8 342 + } 343 + }).collect(); 344 + 345 + stdin 346 + .write_all(&pixels) 347 + .await 348 + .map_err(|e| SpatialError::IoError(format!("Failed to write depth frame: {}", e)))?; 349 + } 350 + 351 + drop(stdin); 352 + 353 + let output = child 354 + .wait_with_output() 355 + .await 356 + .map_err(|e| SpatialError::Other(format!("ffmpeg depth encoding failed: {}", e)))?; 357 + 358 + if !output.status.success() { 359 + let stderr = String::from_utf8_lossy(&output.stderr); 360 + return Err(SpatialError::Other(format!( 361 + "ffmpeg depth encoding exited with error: {}", stderr 362 + ))); 363 + } 364 + 365 + Ok(()) 366 + } 367 + 291 368 fn is_spatial_cli_available() -> bool { 292 369 std::process::Command::new("spatial") 293 370 .arg("--version") ··· 370 447 input_path: &Path, 371 448 output_path: &Path, 372 449 config: SpatialConfig, 450 + output_types: &[OutputType], 373 451 progress_cb: Option<ProgressCallback>, 374 452 ) -> SpatialResult<()> { 375 453 if !input_path.exists() { ··· 379 457 ))); 380 458 } 381 459 382 - let metadata = get_video_metadata(input_path).await?; 383 - let use_spatial = is_spatial_cli_available(); 460 + let do_depth = needs_depth(output_types); 461 + let do_stereo = needs_stereo(output_types); 462 + 463 + let mut metadata = get_video_metadata(input_path).await?; 464 + metadata.width = metadata.width & !1; 465 + metadata.height = metadata.height & !1; 466 + let use_spatial = do_stereo && is_spatial_cli_available(); 467 + 468 + let stereo_output = { 469 + let stem = output_path.file_stem().and_then(|s| s.to_str()).unwrap_or("output"); 470 + let parent = output_path.parent().unwrap_or_else(|| Path::new(".")); 471 + parent.join(format!("{}-spatial.mov", stem)) 472 + }; 384 473 385 474 let sbs_path = if use_spatial { 386 475 let temp_dir = std::env::temp_dir(); ··· 391 480 .unwrap_or_default() 392 481 .as_millis() 393 482 )) 483 + } else if do_stereo { 484 + stereo_output.clone() 394 485 } else { 395 486 output_path.to_path_buf() 396 487 }; ··· 406 497 std::sync::Arc::new(crate::depth_coreml::CoreMLDepthEstimator::new(model_str)?) 407 498 }; 408 499 500 + let mut depth_processor = DepthProcessor::new( 501 + config.temporal_alpha, 502 + config.bilateral_sigma_space, 503 + config.bilateral_sigma_color, 504 + config.depth_blur_sigma, 505 + config.normalize_mode.clone(), 506 + ); 507 + 508 + let total_frames = metadata.total_frames; 509 + 510 + if matches!(config.normalize_mode, NormalizeMode::Global) { 511 + if let Some(ref cb) = progress_cb { 512 + cb(VideoProgress::new(0, total_frames, "scanning".to_string())); 513 + } 514 + 515 + let mut scan_rx = extract_frames(input_path, &metadata).await?; 516 + let mut scan_count = 0u32; 517 + while let Some(frame_data) = scan_rx.recv().await { 518 + let frame = frame_to_image(&frame_data, metadata.width, metadata.height)?; 519 + scan_count += 1; 520 + 521 + #[cfg(all(target_os = "macos", feature = "coreml"))] 522 + { 523 + let raw = estimator.estimate_unnormalized(&frame)?; 524 + depth_processor.update_global_range(&raw); 525 + } 526 + 527 + #[cfg(not(all(target_os = "macos", feature = "coreml")))] 528 + { 529 + let _ = frame; 530 + } 531 + 532 + if let Some(ref cb) = progress_cb { 533 + if scan_count % 10 == 0 || scan_count == total_frames { 534 + cb(VideoProgress::new( 535 + scan_count, 536 + total_frames, 537 + "scanning".to_string(), 538 + )); 539 + } 540 + } 541 + } 542 + } 543 + 409 544 let mut frame_rx = extract_frames(input_path, &metadata).await?; 410 545 411 - let (processed_tx, processed_rx) = mpsc::channel::<(DynamicImage, DynamicImage)>(10); 546 + let stereo_tx_opt; 547 + let stereo_handle; 412 548 413 - let encode_handle = tokio::spawn(encode_stereo_video( 414 - sbs_path.clone(), 415 - metadata.clone(), 416 - processed_rx, 417 - )); 549 + if do_stereo { 550 + let (tx, rx) = mpsc::channel::<(DynamicImage, DynamicImage)>(10); 551 + stereo_tx_opt = Some(tx); 552 + stereo_handle = Some(tokio::spawn(encode_stereo_video( 553 + sbs_path.clone(), 554 + metadata.clone(), 555 + rx, 556 + ))); 557 + } else { 558 + stereo_tx_opt = None; 559 + stereo_handle = None; 560 + } 561 + 562 + let depth_tx_opt; 563 + let depth_handle; 564 + 565 + if do_depth { 566 + let depth_path = { 567 + let stem = output_path.file_stem().and_then(|s| s.to_str()).unwrap_or("output"); 568 + let parent = output_path.parent().unwrap_or_else(|| Path::new(".")); 569 + parent.join(format!("{}-depth.mov", stem)) 570 + }; 571 + 572 + let (tx, rx) = mpsc::channel::<Array2<f32>>(10); 573 + depth_tx_opt = Some(tx); 574 + depth_handle = Some(tokio::spawn(encode_depth_video( 575 + depth_path, 576 + metadata.clone(), 577 + rx, 578 + ))); 579 + } else { 580 + depth_tx_opt = None; 581 + depth_handle = None; 582 + } 418 583 419 584 let mut frame_count = 0u32; 420 - let total_frames = metadata.total_frames; 421 585 422 586 if let Some(ref cb) = progress_cb { 423 587 cb(VideoProgress::new(0, total_frames, "extracting".to_string())); ··· 438 602 } 439 603 440 604 #[cfg(all(target_os = "macos", feature = "coreml"))] 441 - let depth_map = estimator.estimate(&frame)?; 605 + let depth_map = { 606 + let raw = estimator.estimate_unnormalized(&frame)?; 607 + depth_processor.process(raw) 608 + }; 442 609 443 610 #[cfg(not(all(target_os = "macos", feature = "coreml")))] 444 611 let depth_map = { ··· 456 623 } 457 624 }; 458 625 459 - let (left, right) = generate_stereo_pair(&frame, &depth_map, config.max_disparity)?; 626 + if let Some(ref depth_tx) = depth_tx_opt { 627 + if depth_tx.send(depth_map.clone()).await.is_err() { 628 + return Err(SpatialError::Other( 629 + "Depth encoder stopped unexpectedly".to_string(), 630 + )); 631 + } 632 + } 460 633 461 - if processed_tx.send((left, right)).await.is_err() { 462 - return Err(SpatialError::Other( 463 - "Encoder stopped unexpectedly".to_string(), 464 - )); 634 + if let Some(ref stereo_tx) = stereo_tx_opt { 635 + let (left, right) = generate_stereo_pair(&frame, &depth_map, config.max_disparity)?; 636 + if stereo_tx.send((left, right)).await.is_err() { 637 + return Err(SpatialError::Other( 638 + "Encoder stopped unexpectedly".to_string(), 639 + )); 640 + } 465 641 } 466 642 } 467 643 468 - drop(processed_tx); 644 + drop(stereo_tx_opt); 645 + drop(depth_tx_opt); 469 646 470 647 if let Some(ref cb) = progress_cb { 471 648 cb(VideoProgress::new( ··· 475 652 )); 476 653 } 477 654 478 - encode_handle 479 - .await 480 - .map_err(|e| SpatialError::Other(format!("Encoding task failed: {}", e)))??; 655 + if let Some(handle) = stereo_handle { 656 + handle 657 + .await 658 + .map_err(|e| SpatialError::Other(format!("Stereo encoding task failed: {}", e)))??; 659 + } 660 + 661 + if let Some(handle) = depth_handle { 662 + handle 663 + .await 664 + .map_err(|e| SpatialError::Other(format!("Depth encoding task failed: {}", e)))??; 665 + } 481 666 482 667 if use_spatial { 483 668 if let Some(ref cb) = progress_cb { ··· 488 673 )); 489 674 } 490 675 491 - let result = encode_mvhevc_video(&sbs_path, output_path, input_path, &metadata).await; 492 - let _ = tokio::fs::remove_file(&sbs_path).await; 676 + let result = encode_mvhevc_video(&sbs_path, &stereo_output, input_path, &metadata).await; 493 677 result?; 494 678 } 495 679
+24 -16
swift-bridge/CoreMLDepth.swift
··· 18 18 // Compile the model first 19 19 let compiledURL = try MLModel.compileModel(at: url) 20 20 modelURL = compiledURL 21 - print("✓ Model compiled to: \(compiledURL.path)") 22 21 } else { 23 22 modelURL = url 24 23 } 25 24 26 - // Configure to use all compute units (ANE + GPU + CPU) 27 25 let config = MLModelConfiguration() 28 26 config.computeUnits = .all 29 27 30 28 let model = try MLModel(contentsOf: modelURL, configuration: config) 31 - print("✓ CoreML model loaded successfully") 32 29 33 30 // Return retained pointer to model 34 31 return Unmanaged.passRetained(model as AnyObject).toOpaque() ··· 60 57 // Create CVPixelBuffer from RGB8 data 61 58 // CoreML expects BGRA format, so we need to convert RGB -> BGRA 62 59 var pixelBuffer: CVPixelBuffer? 63 - let bytesPerRow = w * 4 64 60 let status = CVPixelBufferCreate( 65 61 kCFAllocatorDefault, 66 62 w, ··· 82 78 let baseAddress = CVPixelBufferGetBaseAddress(pixelBuffer)! 83 79 let bgraData = baseAddress.assumingMemoryBound(to: UInt8.self) 84 80 85 - // Convert RGB (HWC) to BGRA 86 - for i in 0..<(w * h) { 87 - let rgbIdx = i * 3 88 - let bgraIdx = i * 4 89 - bgraData[bgraIdx + 0] = rgbData[rgbIdx + 2] // B 90 - bgraData[bgraIdx + 1] = rgbData[rgbIdx + 1] // G 91 - bgraData[bgraIdx + 2] = rgbData[rgbIdx + 0] // R 92 - bgraData[bgraIdx + 3] = 255 // A 81 + // Convert RGB (HWC) to BGRA, respecting CVPixelBuffer row padding 82 + let actualBytesPerRow = CVPixelBufferGetBytesPerRow(pixelBuffer) 83 + for y in 0..<h { 84 + for x in 0..<w { 85 + let rgbIdx = (y * w + x) * 3 86 + let bgraIdx = y * actualBytesPerRow + x * 4 87 + bgraData[bgraIdx + 0] = rgbData[rgbIdx + 2] // B 88 + bgraData[bgraIdx + 1] = rgbData[rgbIdx + 1] // G 89 + bgraData[bgraIdx + 2] = rgbData[rgbIdx + 0] // R 90 + bgraData[bgraIdx + 3] = 255 // A 91 + } 93 92 } 94 93 95 94 // Run inference with CVPixelBuffer ··· 107 106 return -3 108 107 } 109 108 110 - // Convert Float16 output to Float32 for Rust 111 - let outputCount = w * h 109 + // Convert Float16 output to Float32 for Rust, respecting MLMultiArray strides 112 110 let srcPtr16 = depthArray.dataPointer.assumingMemoryBound(to: Float16.self) 113 - for i in 0..<outputCount { 114 - outputPtr[i] = Float(srcPtr16[i]) 111 + let strides = depthArray.strides.map { $0.intValue } 112 + if strides.count == 3 { 113 + for y in 0..<h { 114 + for x in 0..<w { 115 + let srcIdx = y * strides[1] + x * strides[2] 116 + outputPtr[y * w + x] = Float(srcPtr16[srcIdx]) 117 + } 118 + } 119 + } else { 120 + for i in 0..<(w * h) { 121 + outputPtr[i] = Float(srcPtr16[i]) 122 + } 115 123 } 116 124 117 125 return 0