burn_vision/backends/cpu/morphology/
mod.rs

1use std::fmt::Debug;
2
3use burn_tensor::{
4    BasicOps, Bool, BoolStore, DType, Element, ElementLimits, ElementOrdered, Shape, Tensor,
5    TensorData, backend::Backend, cast::ToElement, ops::BoolTensor,
6};
7use filter::{MaxOp, MinOp, MorphOperator, VecMorphOperator};
8use filter_engine::{ColFilter, Filter, Filter2D, FilterEngine, RowFilter};
9use macerator::{Simd, VOrd};
10
11use crate::{BorderType, MorphOptions, Point, Size};
12
13use super::MinMax;
14
15mod filter;
16mod filter_engine;
17
18/// A morphology operation.
19/// TODO: Implement composite ops
20#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
21pub enum MorphOp {
22    Erode,
23    Dilate,
24}
25
26pub enum MorphKernel<B: Element> {
27    Rect {
28        size: Size,
29        anchor: Point,
30    },
31    Other {
32        kernel: Vec<B>,
33        size: Size,
34        anchor: Point,
35    },
36}
37
38pub fn morph<B: Backend, K: BasicOps<B>>(
39    input: Tensor<B, 3, K>,
40    kernel: BoolTensor<B>,
41    op: MorphOp,
42    opts: MorphOptions<B, K>,
43) -> Tensor<B, 3, K> {
44    let device = input.device();
45
46    let kernel = Tensor::<B, 2, Bool>::new(kernel);
47    let kshape = kernel.shape().dims();
48    let [kh, kw] = kshape;
49
50    let kernel = kernel.into_data().into_vec::<B::BoolElem>().unwrap();
51    let is_rect = kernel.iter().all(|it| it.to_bool());
52    let anchor = opts.anchor.unwrap_or(Point::new(kw / 2, kh / 2));
53    let iter = opts.iterations;
54    let btype = opts.border_type;
55    let bvalue = opts.border_value.map(|it| it.into_data());
56
57    let size = Size::new(kw, kh);
58    let kernel = if is_rect {
59        MorphKernel::Rect { size, anchor }
60    } else {
61        MorphKernel::Other {
62            kernel,
63            size,
64            anchor,
65        }
66    };
67
68    let shape = input.shape();
69    let data = input.into_data();
70    match data.dtype {
71        DType::F64 => {
72            morph_typed::<B, K, f64>(data, shape, kernel, op, iter, btype, bvalue, &device)
73        }
74        DType::F32 | DType::Flex32 => {
75            morph_typed::<B, K, f32>(data, shape, kernel, op, iter, btype, bvalue, &device)
76        }
77        DType::F16 | DType::BF16 => morph_typed::<B, K, f32>(
78            data.convert::<f32>(),
79            shape,
80            kernel,
81            op,
82            iter,
83            btype,
84            bvalue,
85            &device,
86        ),
87        DType::I64 => {
88            morph_typed::<B, K, i64>(data, shape, kernel, op, iter, btype, bvalue, &device)
89        }
90        DType::I32 => {
91            morph_typed::<B, K, i32>(data, shape, kernel, op, iter, btype, bvalue, &device)
92        }
93        DType::I16 => {
94            morph_typed::<B, K, i16>(data, shape, kernel, op, iter, btype, bvalue, &device)
95        }
96        DType::I8 => morph_typed::<B, K, i8>(data, shape, kernel, op, iter, btype, bvalue, &device),
97        DType::U64 => {
98            morph_typed::<B, K, u64>(data, shape, kernel, op, iter, btype, bvalue, &device)
99        }
100        DType::U32 | DType::Bool(BoolStore::U32) => {
101            morph_typed::<B, K, u32>(data, shape, kernel, op, iter, btype, bvalue, &device)
102        }
103        DType::U16 => {
104            morph_typed::<B, K, u16>(data, shape, kernel, op, iter, btype, bvalue, &device)
105        }
106        DType::U8 | DType::Bool(BoolStore::U8) => {
107            morph_typed::<B, K, u8>(data, shape, kernel, op, iter, btype, bvalue, &device)
108        }
109        DType::Bool(BoolStore::Native) => {
110            morph_bool::<B, K>(data, shape, kernel, op, iter, btype, bvalue, &device)
111        }
112        DType::QFloat(_) => unimplemented!(),
113    }
114}
115
116#[allow(clippy::too_many_arguments)]
117fn morph_typed<B: Backend, K: BasicOps<B>, T: VOrd + MinMax + ElementOrdered>(
118    mut input: TensorData,
119    shape: Shape,
120    kernel: MorphKernel<B::BoolElem>,
121    op: MorphOp,
122    iter: usize,
123    btype: BorderType,
124    bvalue: Option<TensorData>,
125    device: &B::Device,
126) -> Tensor<B, 3, K> {
127    let data = input.as_mut_slice::<T>().unwrap();
128    let bvalue = border_value(btype, bvalue, op, &shape);
129    run_morph(data, shape, kernel, op, iter, btype, &bvalue);
130    Tensor::from_data(input, device)
131}
132
133#[allow(clippy::too_many_arguments)]
134fn morph_bool<B: Backend, K: BasicOps<B>>(
135    mut input: TensorData,
136    shape: Shape,
137    kernel: MorphKernel<B::BoolElem>,
138    op: MorphOp,
139    iter: usize,
140    btype: BorderType,
141    bvalue: Option<TensorData>,
142    device: &B::Device,
143) -> Tensor<B, 3, K> {
144    let data = input.as_mut_slice::<bool>().unwrap();
145    // SAFETY: Morph can't produce invalid boolean values
146    let data = unsafe { core::mem::transmute::<&mut [bool], &mut [u8]>(data) };
147    let bvalue = border_value(btype, bvalue, op, &shape);
148    run_morph(data, shape.clone(), kernel, op, iter, btype, &bvalue);
149    Tensor::from_data(input, device)
150}
151
152fn border_value<T: Element + ElementLimits>(
153    btype: BorderType,
154    bvalue: Option<TensorData>,
155    op: MorphOp,
156    shape: &Shape,
157) -> Vec<T> {
158    let [_, _, ch] = shape.dims();
159    match (btype, bvalue) {
160        (BorderType::Constant, Some(value)) => value.convert::<T>().into_vec().unwrap(),
161        (BorderType::Constant, None) => match op {
162            MorphOp::Erode => vec![T::MAX; ch],
163            MorphOp::Dilate => vec![T::MIN; ch],
164        },
165        _ => vec![],
166    }
167}
168
169fn run_morph<T: VOrd + MinMax + Element, B: Element>(
170    input: &mut [T],
171    shape: Shape,
172    kernel: MorphKernel<B>,
173    op: MorphOp,
174    iter: usize,
175    btype: BorderType,
176    bvalue: &[T],
177) {
178    match op {
179        MorphOp::Erode => {
180            let filter = filter::<T, MinOp, B>(kernel);
181            dispatch_morph(input, shape, filter, btype, bvalue, iter);
182        }
183        MorphOp::Dilate => {
184            let filter = filter::<T, MaxOp, B>(kernel);
185            dispatch_morph(input, shape, filter, btype, bvalue, iter);
186        }
187    };
188}
189
190fn filter<T: VOrd + MinMax, Op: MorphOperator<T> + VecMorphOperator<T>, B: Element>(
191    kernel: MorphKernel<B>,
192) -> Filter<T, Op> {
193    match kernel {
194        MorphKernel::Rect { size, anchor } => {
195            let row_filter = RowFilter::new(size.width, anchor.x);
196            let col_filter = ColFilter::new(size.height, anchor.y);
197            Filter::Separable {
198                row_filter,
199                col_filter,
200            }
201        }
202        MorphKernel::Other {
203            kernel,
204            size,
205            anchor,
206        } => {
207            let filter = Filter2D::new(&kernel, size, anchor);
208            Filter::Fallback(filter)
209        }
210    }
211}
212
213#[inline(always)]
214#[allow(clippy::too_many_arguments)]
215#[macerator::with_simd]
216fn dispatch_morph<
217    'a,
218    S: Simd,
219    T: VOrd + MinMax + Debug,
220    Op: MorphOperator<T> + VecMorphOperator<T>,
221>(
222    buffer: &'a mut [T],
223    buffer_shape: Shape,
224    filter: filter_engine::Filter<T, Op>,
225    border_type: BorderType,
226    border_value: &'a [T],
227    iterations: usize,
228) where
229    'a: 'a,
230{
231    let [_, _, ch] = buffer_shape.dims();
232    let mut engine = FilterEngine::<S, _, _>::new(filter, border_type, border_value, ch);
233    engine.apply(buffer, buffer_shape.clone());
234    for _ in 1..iterations {
235        engine.apply(buffer, buffer_shape.clone());
236    }
237}
238
239/// Shape of the structuring element
240#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
241pub enum KernelShape {
242    /// Rectangular kernel
243    Rect,
244    /// Cross shaped kernel
245    Cross,
246    /// Ellipse shaped kernel
247    Ellipse,
248}
249
250/// Create a structuring element tensor for use with morphology ops
251pub fn create_structuring_element<B: Backend>(
252    shape: KernelShape,
253    ksize: Size,
254    anchor: Option<Point>,
255    device: &B::Device,
256) -> Tensor<B, 2, Bool> {
257    fn create_kernel(shape: KernelShape, ksize: Size, anchor: Option<Point>) -> Vec<bool> {
258        let anchor = anchor.unwrap_or(Point::new(ksize.width / 2, ksize.height / 2));
259        let mut r = 0;
260        let mut c = 0;
261        let mut inv_r2 = 0.0;
262
263        if (ksize.width == 1 && ksize.height == 1) || shape == KernelShape::Rect {
264            return vec![true; ksize.height * ksize.width];
265        }
266
267        if shape == KernelShape::Ellipse {
268            r = ksize.height / 2;
269            c = ksize.width / 2;
270            inv_r2 = if r > 0 { 1.0 / (r * r) as f64 } else { 0.0 }
271        }
272
273        let mut elem = vec![false; ksize.height * ksize.width];
274
275        for i in 0..ksize.height {
276            let mut j1 = 0;
277            let mut j2 = 0;
278            if shape == KernelShape::Cross && i == anchor.y {
279                j2 = ksize.width;
280            } else if shape == KernelShape::Cross {
281                j1 = anchor.x;
282                j2 = j1 + 1;
283            } else {
284                let dy = i as isize - r as isize;
285                if dy.abs() <= r as isize {
286                    let dx = (c as f64 * ((r * r - (dy * dy) as usize) as f64 * inv_r2).sqrt())
287                        .round() as isize;
288                    j1 = (c as isize - dx).max(0) as usize;
289                    j2 = (c + dx as usize + 1).min(ksize.width);
290                }
291            }
292
293            for j in j1..j2 {
294                elem[i * ksize.width + j] = true;
295            }
296        }
297        elem
298    }
299
300    let elem = create_kernel(shape, ksize, anchor);
301
302    let data = TensorData::new(elem, [ksize.height, ksize.width]);
303    Tensor::from_data(data, device)
304}
burn_vision/backends/cpu/morphology/mod.rs

burn_vision/backends/cpu/morphology/
mod.rs
