Skip to main content

spatialrust_gpu/image/
vision_chain.rs

1//! Explicit upload-once GPU-resident Vision 2 chain.
2
3use spatialrust_core::SpatialResult;
4
5use super::{
6    box_blur_gpu, morphology_gpu, pack_ai_chw_gpu, resize_nearest_gpu, rgb_to_gray_gpu, sobel_gpu,
7    GpuAiTensor, GpuImage, GpuImageBorder, GpuMorphology,
8};
9use crate::WgpuRuntime;
10
11/// Configuration for the resident resize-to-AI chain.
12#[derive(Clone, Copy, Debug, PartialEq)]
13pub struct GpuVisionChainOptions {
14    /// Resize output width.
15    pub width: u32,
16    /// Resize output height.
17    pub height: u32,
18    /// Odd box-blur kernel size.
19    pub blur_kernel: u32,
20    /// Odd morphology kernel size.
21    pub morphology_kernel: u32,
22    /// Morphology operation after Sobel.
23    pub morphology: GpuMorphology,
24    /// Input-value scale for AI packing.
25    pub scale: f32,
26    /// Per-channel means.
27    pub mean: [f32; 4],
28    /// Per-channel standard deviations.
29    pub std: [f32; 4],
30}
31
32impl Default for GpuVisionChainOptions {
33    fn default() -> Self {
34        Self {
35            width: 640,
36            height: 480,
37            blur_kernel: 3,
38            morphology_kernel: 3,
39            morphology: GpuMorphology::Dilate,
40            scale: 1.0 / 255.0,
41            mean: [0.0; 4],
42            std: [1.0; 4],
43        }
44    }
45}
46
47/// Runs resize → gray → blur → Sobel → morphology → planar AI packing.
48///
49/// The caller owns the uploaded source. Every intermediate is recycled after
50/// its consumer has been submitted, and no host readback occurs.
51pub fn run_gpu_vision_chain(
52    runtime: &WgpuRuntime,
53    source: &GpuImage,
54    options: GpuVisionChainOptions,
55) -> SpatialResult<GpuAiTensor> {
56    let resized = resize_nearest_gpu(runtime, source, options.width, options.height)?;
57    let gray = rgb_to_gray_gpu(runtime, &resized)?;
58    resized.recycle(runtime);
59    let blurred = box_blur_gpu(
60        runtime,
61        &gray,
62        options.blur_kernel,
63        options.blur_kernel,
64        GpuImageBorder::Replicate,
65    )?;
66    gray.recycle(runtime);
67    let edges = sobel_gpu(runtime, &blurred)?;
68    blurred.recycle(runtime);
69    let morphology = morphology_gpu(
70        runtime,
71        &edges,
72        options.morphology_kernel,
73        options.morphology_kernel,
74        options.morphology,
75    )?;
76    edges.recycle(runtime);
77    let tensor = pack_ai_chw_gpu(runtime, &morphology, options.scale, options.mean, options.std)?;
78    morphology.recycle(runtime);
79    Ok(tensor)
80}