ad_plugins_rs/
transform.rs

1use std::sync::Arc;
2
3use ad_core_rs::color::NDColorMode;
4use ad_core_rs::ndarray::{NDArray, NDDataBuffer, NDDimension};
5use ad_core_rs::ndarray_pool::NDArrayPool;
6use ad_core_rs::plugin::runtime::{NDPluginProcess, ProcessResult};
7use parking_lot::Mutex;
8
9/// Transform types matching C++ `NDPluginTransformType_t`.
10///
11/// The numeric ordering is the C++ enum order:
12/// `None=0, Rotate90=1, Rotate180=2, Rotate270=3, Mirror=4,
13/// Rotate90Mirror=5, Rotate180Mirror=6, Rotate270Mirror=7`.
14///
15/// - `Mirror` is a horizontal flip.
16/// - `Rotate90Mirror` is the transpose (main-diagonal flip).
17/// - `Rotate180Mirror` is a vertical flip.
18/// - `Rotate270Mirror` is the anti-diagonal flip.
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20#[repr(u8)]
21pub enum TransformType {
22    None = 0,
23    Rot90CW = 1,
24    Rot180 = 2,
25    Rot90CCW = 3,
26    FlipHoriz = 4,
27    /// C++ `Rotate90Mirror`: transpose / main-diagonal flip.
28    FlipDiag = 5,
29    /// C++ `Rotate180Mirror`: vertical flip.
30    FlipVert = 6,
31    /// C++ `Rotate270Mirror`: anti-diagonal flip.
32    FlipAntiDiag = 7,
33}
34
35impl TransformType {
36    pub fn from_u8(v: u8) -> Self {
37        match v {
38            1 => Self::Rot90CW,
39            2 => Self::Rot180,
40            3 => Self::Rot90CCW,
41            4 => Self::FlipHoriz,
42            // C++ TransformRotate90Mirror == transpose.
43            5 => Self::FlipDiag,
44            // C++ TransformRotate180Mirror == vertical flip.
45            6 => Self::FlipVert,
46            7 => Self::FlipAntiDiag,
47            _ => Self::None,
48        }
49    }
50
51    /// Whether this transform swaps x and y dimensions.
52    pub fn swaps_dims(&self) -> bool {
53        matches!(
54            self,
55            Self::Rot90CW | Self::Rot90CCW | Self::FlipDiag | Self::FlipAntiDiag
56        )
57    }
58}
59
60/// Map source (x, y) to destination (x, y) for the given transform.
61fn map_coords(
62    sx: usize,
63    sy: usize,
64    src_w: usize,
65    src_h: usize,
66    transform: TransformType,
67) -> (usize, usize) {
68    match transform {
69        TransformType::None => (sx, sy),
70        TransformType::Rot90CW => (src_h - 1 - sy, sx),
71        TransformType::Rot180 => (src_w - 1 - sx, src_h - 1 - sy),
72        TransformType::Rot90CCW => (sy, src_w - 1 - sx),
73        TransformType::FlipHoriz => (src_w - 1 - sx, sy),
74        TransformType::FlipVert => (sx, src_h - 1 - sy),
75        TransformType::FlipDiag => (sy, sx),
76        TransformType::FlipAntiDiag => (src_h - 1 - sy, src_w - 1 - sx),
77    }
78}
79
80/// Per-color-mode element strides for a 2-D or 3-D image of the given
81/// X/Y/color sizes. Mirrors C++ `NDArray::getInfo` stride layout: returns
82/// `(x_stride, y_stride, color_stride)` and the destination dimension order.
83fn strides_for(color_mode: NDColorMode, xs: usize, ys: usize, cs: usize) -> (usize, usize, usize) {
84    match color_mode {
85        NDColorMode::RGB1 => (cs, xs * cs, 1),
86        NDColorMode::RGB2 => (1, xs * cs, xs),
87        // RGB3 / Mono / others: planar X-fastest layout.
88        _ => (1, xs, xs * ys),
89    }
90}
91
92/// Build the destination dimension vector for `color_mode` with the given
93/// X/Y/color sizes, matching the C++ dimension order per color mode.
94fn dims_for(
95    color_mode: NDColorMode,
96    xs: usize,
97    ys: usize,
98    cs: usize,
99    ndims: usize,
100) -> Vec<NDDimension> {
101    if ndims < 3 {
102        return vec![NDDimension::new(xs), NDDimension::new(ys)];
103    }
104    match color_mode {
105        NDColorMode::RGB1 => vec![
106            NDDimension::new(cs),
107            NDDimension::new(xs),
108            NDDimension::new(ys),
109        ],
110        NDColorMode::RGB2 => vec![
111            NDDimension::new(xs),
112            NDDimension::new(cs),
113            NDDimension::new(ys),
114        ],
115        _ => vec![
116            NDDimension::new(xs),
117            NDDimension::new(ys),
118            NDDimension::new(cs),
119        ],
120    }
121}
122
123/// Apply a transform to an NDArray.
124///
125/// Handles 2-D mono images and 3-D RGB1/RGB2/RGB3 color images. The per-color
126/// reindexing mirrors C++ `transformNDArray`: source `(x, y)` is geometrically
127/// mapped to destination `(x, y)` and every color component is copied with the
128/// destination strides recomputed for the (possibly swapped) X/Y sizes.
129pub fn apply_transform(src: &NDArray, transform: TransformType) -> NDArray {
130    if transform == TransformType::None || src.dims.len() < 2 {
131        return src.clone();
132    }
133
134    let info = src.info();
135    let src_w = info.x_size;
136    let src_h = info.y_size;
137    let color = info.color_size.max(1);
138    if src_w == 0 || src_h == 0 {
139        return src.clone();
140    }
141
142    let (dst_w, dst_h) = if transform.swaps_dims() {
143        (src_h, src_w)
144    } else {
145        (src_w, src_h)
146    };
147
148    let (sxs, sys, scs) = (
149        info.x_stride,
150        info.y_stride.max(1),
151        info.color_stride.max(1),
152    );
153    let (dxs, dys, dcs) = strides_for(info.color_mode, dst_w, dst_h, color);
154    let total = dst_w * dst_h * color;
155
156    macro_rules! transform_buf {
157        ($vec:expr, $zero:expr) => {{
158            let mut out = vec![$zero; total];
159            for sy in 0..src_h {
160                for sx in 0..src_w {
161                    let (dx, dy) = map_coords(sx, sy, src_w, src_h, transform);
162                    let s_base = sy * sys + sx * sxs;
163                    let d_base = dy * dys + dx * dxs;
164                    for c in 0..color {
165                        out[d_base + c * dcs] = $vec[s_base + c * scs];
166                    }
167                }
168            }
169            out
170        }};
171    }
172
173    let out_data = match &src.data {
174        NDDataBuffer::U8(v) => NDDataBuffer::U8(transform_buf!(v, 0)),
175        NDDataBuffer::U16(v) => NDDataBuffer::U16(transform_buf!(v, 0)),
176        NDDataBuffer::I8(v) => NDDataBuffer::I8(transform_buf!(v, 0)),
177        NDDataBuffer::I16(v) => NDDataBuffer::I16(transform_buf!(v, 0)),
178        NDDataBuffer::I32(v) => NDDataBuffer::I32(transform_buf!(v, 0)),
179        NDDataBuffer::U32(v) => NDDataBuffer::U32(transform_buf!(v, 0)),
180        NDDataBuffer::I64(v) => NDDataBuffer::I64(transform_buf!(v, 0)),
181        NDDataBuffer::U64(v) => NDDataBuffer::U64(transform_buf!(v, 0)),
182        NDDataBuffer::F32(v) => NDDataBuffer::F32(transform_buf!(v, 0.0)),
183        NDDataBuffer::F64(v) => NDDataBuffer::F64(transform_buf!(v, 0.0)),
184    };
185
186    let dims = dims_for(info.color_mode, dst_w, dst_h, color, src.dims.len());
187    let mut arr = NDArray::new(dims, src.data.data_type());
188    arr.data = out_data;
189    arr.unique_id = src.unique_id;
190    arr.timestamp = src.timestamp;
191    arr.time_stamp = src.time_stamp;
192    arr.attributes = src.attributes.clone();
193    arr
194}
195
196// --- New TransformProcessor (NDPluginProcess-based) ---
197
198/// Pure transform processing logic.
199pub struct TransformProcessor {
200    transform: Mutex<TransformType>,
201    transform_type_idx: Option<usize>,
202}
203
204impl TransformProcessor {
205    pub fn new(transform: TransformType) -> Self {
206        Self {
207            transform: Mutex::new(transform),
208            transform_type_idx: None,
209        }
210    }
211}
212
213impl NDPluginProcess for TransformProcessor {
214    fn process_array(&self, array: &NDArray, _pool: &NDArrayPool) -> ProcessResult {
215        // C reads the transform type under the port lock and releases it
216        // before `transformImage` (NDPluginTransform.cpp:500). A guard passed
217        // straight into the call would live to the end of the statement and
218        // hold across the whole rotation.
219        let transform = *self.transform.lock();
220        let out = apply_transform(array, transform);
221        ProcessResult::arrays(vec![Arc::new(out)])
222    }
223
224    fn plugin_type(&self) -> &str {
225        "NDPluginTransform"
226    }
227
228    fn register_params(
229        &mut self,
230        base: &mut asyn_rs::port::PortDriverBase,
231    ) -> asyn_rs::error::AsynResult<()> {
232        use asyn_rs::param::ParamType;
233        base.create_param("TRANSFORM_TYPE", ParamType::Int32)?;
234        self.transform_type_idx = base.find_param("TRANSFORM_TYPE");
235        Ok(())
236    }
237
238    fn on_param_change(
239        &self,
240        reason: usize,
241        params: &ad_core_rs::plugin::runtime::PluginParamSnapshot,
242    ) -> ad_core_rs::plugin::runtime::ParamChangeResult {
243        if Some(reason) == self.transform_type_idx {
244            *self.transform.lock() = TransformType::from_u8(params.value.as_i32() as u8);
245        }
246        ad_core_rs::plugin::runtime::ParamChangeResult::updates(vec![])
247    }
248}
249
250#[cfg(test)]
251mod tests {
252    use super::*;
253    use ad_core_rs::ndarray::NDDataType;
254
255    /// Create a 3x2 array:
256    /// [1, 2, 3]
257    /// [4, 5, 6]
258    fn make_3x2() -> NDArray {
259        let mut arr = NDArray::new(
260            vec![NDDimension::new(3), NDDimension::new(2)],
261            NDDataType::UInt8,
262        );
263        if let NDDataBuffer::U8(ref mut v) = arr.data {
264            *v = vec![1, 2, 3, 4, 5, 6];
265        }
266        arr
267    }
268
269    fn get_u8(arr: &NDArray) -> &[u8] {
270        match &arr.data {
271            NDDataBuffer::U8(v) => v,
272            _ => panic!("not u8"),
273        }
274    }
275
276    #[test]
277    fn test_none() {
278        let arr = make_3x2();
279        let out = apply_transform(&arr, TransformType::None);
280        assert_eq!(get_u8(&out), &[1, 2, 3, 4, 5, 6]);
281    }
282
283    #[test]
284    fn test_rot90cw() {
285        let arr = make_3x2();
286        let out = apply_transform(&arr, TransformType::Rot90CW);
287        assert_eq!(out.dims[0].size, 2);
288        assert_eq!(out.dims[1].size, 3);
289        // Expected:
290        // [4, 1]
291        // [5, 2]
292        // [6, 3]
293        assert_eq!(get_u8(&out), &[4, 1, 5, 2, 6, 3]);
294    }
295
296    #[test]
297    fn test_rot180() {
298        let arr = make_3x2();
299        let out = apply_transform(&arr, TransformType::Rot180);
300        assert_eq!(out.dims[0].size, 3);
301        assert_eq!(out.dims[1].size, 2);
302        assert_eq!(get_u8(&out), &[6, 5, 4, 3, 2, 1]);
303    }
304
305    #[test]
306    fn test_rot90ccw() {
307        let arr = make_3x2();
308        let out = apply_transform(&arr, TransformType::Rot90CCW);
309        assert_eq!(out.dims[0].size, 2);
310        assert_eq!(out.dims[1].size, 3);
311        // Expected:
312        // [3, 6]
313        // [2, 5]
314        // [1, 4]
315        assert_eq!(get_u8(&out), &[3, 6, 2, 5, 1, 4]);
316    }
317
318    #[test]
319    fn test_flip_horiz() {
320        let arr = make_3x2();
321        let out = apply_transform(&arr, TransformType::FlipHoriz);
322        assert_eq!(get_u8(&out), &[3, 2, 1, 6, 5, 4]);
323    }
324
325    #[test]
326    fn test_flip_vert() {
327        let arr = make_3x2();
328        let out = apply_transform(&arr, TransformType::FlipVert);
329        assert_eq!(get_u8(&out), &[4, 5, 6, 1, 2, 3]);
330    }
331
332    #[test]
333    fn test_flip_diag() {
334        let arr = make_3x2();
335        let out = apply_transform(&arr, TransformType::FlipDiag);
336        assert_eq!(out.dims[0].size, 2);
337        assert_eq!(out.dims[1].size, 3);
338        // Transpose:
339        // [1, 4]
340        // [2, 5]
341        // [3, 6]
342        assert_eq!(get_u8(&out), &[1, 4, 2, 5, 3, 6]);
343    }
344
345    #[test]
346    fn test_flip_anti_diag() {
347        let arr = make_3x2();
348        let out = apply_transform(&arr, TransformType::FlipAntiDiag);
349        assert_eq!(out.dims[0].size, 2);
350        assert_eq!(out.dims[1].size, 3);
351        // Anti-transpose:
352        // [6, 3]
353        // [5, 2]
354        // [4, 1]
355        assert_eq!(get_u8(&out), &[6, 3, 5, 2, 4, 1]);
356    }
357
358    #[test]
359    fn test_rot90_roundtrip() {
360        let arr = make_3x2();
361        let r1 = apply_transform(&arr, TransformType::Rot90CW);
362        let r2 = apply_transform(&r1, TransformType::Rot90CW);
363        let r3 = apply_transform(&r2, TransformType::Rot90CW);
364        let r4 = apply_transform(&r3, TransformType::Rot90CW);
365        assert_eq!(get_u8(&r4), get_u8(&arr));
366        assert_eq!(r4.dims[0].size, arr.dims[0].size);
367        assert_eq!(r4.dims[1].size, arr.dims[1].size);
368    }
369
370    #[test]
371    fn test_from_u8_cpp_enum_order() {
372        // C++ NDPluginTransformType_t order: value 5 is Rotate90Mirror
373        // (transpose), value 6 is Rotate180Mirror (vertical flip).
374        assert_eq!(TransformType::from_u8(0), TransformType::None);
375        assert_eq!(TransformType::from_u8(1), TransformType::Rot90CW);
376        assert_eq!(TransformType::from_u8(2), TransformType::Rot180);
377        assert_eq!(TransformType::from_u8(3), TransformType::Rot90CCW);
378        assert_eq!(TransformType::from_u8(4), TransformType::FlipHoriz);
379        assert_eq!(TransformType::from_u8(5), TransformType::FlipDiag);
380        assert_eq!(TransformType::from_u8(6), TransformType::FlipVert);
381        assert_eq!(TransformType::from_u8(7), TransformType::FlipAntiDiag);
382    }
383
384    #[test]
385    fn test_transform_5_is_transpose() {
386        // Selecting transform 5 from EPICS must produce a transpose.
387        let arr = make_3x2();
388        let out = apply_transform(&arr, TransformType::from_u8(5));
389        assert_eq!(out.dims[0].size, 2);
390        assert_eq!(out.dims[1].size, 3);
391        assert_eq!(get_u8(&out), &[1, 4, 2, 5, 3, 6]); // transpose
392    }
393
394    #[test]
395    fn test_transform_6_is_vertical_flip() {
396        // Selecting transform 6 from EPICS must produce a vertical flip.
397        let arr = make_3x2();
398        let out = apply_transform(&arr, TransformType::from_u8(6));
399        assert_eq!(out.dims[0].size, 3);
400        assert_eq!(out.dims[1].size, 2);
401        assert_eq!(get_u8(&out), &[4, 5, 6, 1, 2, 3]); // vertical flip
402    }
403
404    /// Build a 2x2 RGB1 image (color-interleaved): pixel (x,y) channel c.
405    /// dims = [color=3, x=2, y=2]. Pixel value encodes 100*y + 10*x + c.
406    fn make_rgb1_2x2() -> NDArray {
407        use ad_core_rs::attributes::{NDAttrSource, NDAttrValue, NDAttribute};
408        let mut arr = NDArray::new(
409            vec![
410                NDDimension::new(3),
411                NDDimension::new(2),
412                NDDimension::new(2),
413            ],
414            NDDataType::UInt8,
415        );
416        arr.attributes.add(NDAttribute::new_static(
417            "ColorMode",
418            "",
419            NDAttrSource::Driver,
420            NDAttrValue::Int32(NDColorMode::RGB1 as i32),
421        ));
422        if let NDDataBuffer::U8(ref mut v) = arr.data {
423            // layout: index = y*(x*c) + x*c + c, with x_stride=3, y_stride=6
424            for y in 0..2 {
425                for x in 0..2 {
426                    for c in 0..3 {
427                        v[y * 6 + x * 3 + c] = (100 * y + 10 * x + c) as u8;
428                    }
429                }
430            }
431        }
432        arr
433    }
434
435    #[test]
436    fn test_rgb1_flip_horiz_keeps_color_grouping() {
437        // Horizontal flip of an RGB1 image: each pixel's 3 channels stay
438        // together; only the x coordinate is mirrored.
439        let arr = make_rgb1_2x2();
440        let out = apply_transform(&arr, TransformType::FlipHoriz);
441        // dims unchanged for a non-swapping transform
442        assert_eq!(out.dims[0].size, 3);
443        assert_eq!(out.dims[1].size, 2);
444        assert_eq!(out.dims[2].size, 2);
445        if let NDDataBuffer::U8(v) = &out.data {
446            // pixel (x=0,y=0) should now hold source (x=1,y=0): 10,11,12
447            assert_eq!(&v[0..3], &[10, 11, 12]);
448            // pixel (x=1,y=0) holds source (x=0,y=0): 0,1,2
449            assert_eq!(&v[3..6], &[0, 1, 2]);
450            // pixel (x=0,y=1) holds source (x=1,y=1): 110,111,112
451            assert_eq!(&v[6..9], &[110, 111, 112]);
452        } else {
453            panic!("not u8");
454        }
455    }
456
457    #[test]
458    fn test_rgb1_rot90cw_swaps_dims_and_keeps_color() {
459        let arr = make_rgb1_2x2();
460        let out = apply_transform(&arr, TransformType::Rot90CW);
461        // x/y swapped (both 2 here), color dim preserved
462        assert_eq!(out.dims[0].size, 3);
463        assert_eq!(out.dims[1].size, 2);
464        assert_eq!(out.dims[2].size, 2);
465        if let NDDataBuffer::U8(v) = &out.data {
466            // Rot90CW maps src (sx,sy) -> (src_h-1-sy, sx).
467            // dest (0,0) <- src (sx,sy) with src_h-1-sy=0, sx=0 => sy=1,sx=0
468            // src (0,1) = 100,101,102
469            assert_eq!(&v[0..3], &[100, 101, 102]);
470        } else {
471            panic!("not u8");
472        }
473    }
474
475    // --- New TransformProcessor tests ---
476
477    #[test]
478    fn test_transform_processor() {
479        let proc = TransformProcessor::new(TransformType::Rot90CW);
480        let pool = NDArrayPool::new(1_000_000);
481
482        let arr = make_3x2();
483        let result = proc.process_array(&arr, &pool);
484        assert_eq!(result.output_arrays.len(), 1);
485        assert_eq!(result.output_arrays[0].dims[0].size, 2); // swapped
486        assert_eq!(result.output_arrays[0].dims[1].size, 3);
487        assert_eq!(get_u8(&result.output_arrays[0]), &[4, 1, 5, 2, 6, 3]);
488    }
489}
ad_plugins_rs/transform.rs

ad_plugins_rs/
transform.rs
