hrx/image/
affine.rs

1use super::*;
2
3/// Explicit byte-image interpolation contracts. The multiplication order is
4/// part of each contract; algebraic rearrangement can change a rounded byte.
5#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
6pub enum RgbSampling {
7    /// Black border, `(value * wx) * wy`, then ties-to-even rounding.
8    BlackTiesEven,
9    /// Replicated border, `value * (wx * wy)`, then `floor(value + 0.5)`.
10    ReplicateHalfUp,
11}
12
13/// Invert a finite FP32 2×3 affine matrix, rejecting singular or overflowing maps.
14pub fn invert_affine(t: [[f32; 3]; 2]) -> Result<[[f32; 3]; 2]> {
15    let det = t[0][0] * t[1][1] - t[0][1] * t[1][0];
16    if !t.iter().flatten().all(|v| v.is_finite()) || !det.is_finite() || det.abs() <= 1e-20 {
17        return Err(Error::Message("degenerate affine transform".into()));
18    }
19    let a = t[1][1] / det;
20    let b = -t[0][1] / det;
21    let c = -t[1][0] / det;
22    let d = t[0][0] / det;
23    let inverse = [
24        [a, b, -a * t[0][2] - b * t[1][2]],
25        [c, d, -c * t[0][2] - d * t[1][2]],
26    ];
27    if !inverse.iter().flatten().all(|v| v.is_finite()) {
28        return Err(Error::Message("non-finite inverse transform".into()));
29    }
30    Ok(inverse)
31}
32
33impl ImageOps {
34    /// Prepare several affine crops from one packed RGB image `[1,h,w,3]`.
35    /// The second input is contiguous F32 `[crops,2,3]` (General layout), mapping
36    /// integer output pixel centers to source coordinates. Output is U8 NHWC.
37    /// Sampling uses the explicit FP32 border/arithmetic/rounding contract.
38    /// Non-finite or excessively large coordinates produce black, never an unsafe
39    /// integer conversion. Matrices are runtime data, not part of the cache key.
40    pub fn prepare_affine_rgb(
41        &self,
42        input: &TensorDesc,
43        crops: usize,
44        height: usize,
45        width: usize,
46        sampling: RgbSampling,
47    ) -> Result<Arc<PreparedModel>> {
48        self.affine_plans
49            .get_or_prepare((input.clone(), crops, height, width, sampling), || {
50                self.affine_rgb_fragment(input, crops, height, width, sampling)?
51                    .prepare(3)
52            })
53    }
54
55    /// Recordable affine crop sampling for composition with model kernels.
56    pub fn affine_rgb_fragment(
57        &self,
58        input: &TensorDesc,
59        crops: usize,
60        height: usize,
61        width: usize,
62        sampling: RgbSampling,
63    ) -> Result<ModelFragment> {
64        self.affine_fragment(input, crops, height, width, sampling, false)
65    }
66
67    /// Sample crops from differently sized RGB images in one packed byte buffer.
68    /// Inputs are U8 `[bytes]`, F32 `[crops,2,3]` inverse matrices, and U32
69    /// `[crops,3]` records `(byte_offset, width, height)`. Repeated records sample
70    /// several faces from one image without duplicating its pixels. Geometry is
71    /// runtime data: changing image sizes does not compile another graph.
72    /// Axes must be 1..=32767. Invalid records produce black; source reads are
73    /// bounded to the byte buffer.
74    pub fn affine_packed_rgb_fragment(
75        &self,
76        input: &TensorDesc,
77        crops: usize,
78        height: usize,
79        width: usize,
80        sampling: RgbSampling,
81    ) -> Result<ModelFragment> {
82        self.affine_fragment(input, crops, height, width, sampling, true)
83    }
84
85    fn affine_fragment(
86        &self,
87        input: &TensorDesc,
88        crops: usize,
89        height: usize,
90        width: usize,
91        sampling: RgbSampling,
92        packed: bool,
93    ) -> Result<ModelFragment> {
94        let shape = input.shape();
95        if input.dtype() != DType::U8
96            || if packed {
97                input.layout() != Layout::General || shape.len() != 1
98            } else {
99                input.layout() != Layout::Nhwc || shape.len() != 4 || shape[0] != 1 || shape[3] != 3
100            }
101            || !input.is_contiguous()
102            || input.is_empty()
103            || input.elements() > i32::MAX as usize
104            || crops == 0
105            || height == 0
106            || width == 0
107        {
108            return Err(Error::Message(
109                "affine requires one nonempty NHWC RGB image and nonempty crops".into(),
110            ));
111        }
112        let output =
113            TensorDesc::new(DType::U8, vec![crops, height, width, 3])?.with_layout(Layout::Nhwc)?;
114        let matrices = TensorDesc::new(DType::F32, vec![crops, 2, 3])?;
115        if output.elements() > i32::MAX as usize || matrices.elements() > i32::MAX as usize {
116            return Err(Error::Message(
117                "affine dimensions exceed kernel indexing limits".into(),
118            ));
119        }
120        let count = output.elements();
121        let mut source = include_str!("affine_rgb.loom").to_owned();
122        source = source.replace(
123            "@GEOMETRY_ARG@",
124            if packed { ", %geometry: buffer" } else { "" },
125        );
126        source = source.replace("@GEOMETRY@", if packed {
127            include_str!("affine_packed_geometry.loom")
128        } else {
129            "%sw = index.constant @SOURCE_WIDTH@ : index\n%fw = scalar.constant @SOURCE_WIDTH@.0 : f32\n%fh = scalar.constant @SOURCE_HEIGHT@.0 : f32\n%source_offset = index.constant 0 : index\n%geometry_ok = scalar.constant true : i1"
130        });
131        for (name, value) in [
132            ("COUNT", count),
133            ("LAST_OUTPUT", count - 1),
134            ("INPUT", input.elements()),
135            ("LAST_INPUT", input.elements() - 1),
136            ("MATRICES", matrices.elements()),
137            ("LAST_MATRIX", matrices.elements() - 1),
138            ("GRID", count.div_ceil(256)),
139            ("WIDTH", width),
140            ("PIXELS", height * width),
141            (
142                "SOURCE_WIDTH",
143                if packed { input.elements() } else { shape[2] },
144            ),
145            (
146                "SOURCE_HEIGHT",
147                if packed { input.elements() } else { shape[1] },
148            ),
149            ("GEOMETRY_COUNT", crops * 3),
150            ("GEOMETRY_LAST", crops * 3 - 1),
151        ] {
152            source = source.replace(&format!("@{name}@"), &value.to_string());
153        }
154        source = source.replace(
155            "@REPLICATE@",
156            if sampling == RgbSampling::ReplicateHalfUp {
157                "true"
158            } else {
159                "false"
160            },
161        );
162        source = source.replace(
163            "@WEIGHTED_VALUE@",
164            match sampling {
165                RgbSampling::BlackTiesEven => {
166                    "%vx = scalar.mulf %v, %wx : f32\n%vxy = scalar.mulf %vx, %wy : f32"
167                }
168                RgbSampling::ReplicateHalfUp => {
169                    "%weight = scalar.mulf %wx, %wy : f32\n%vxy = scalar.mulf %v, %weight : f32"
170                }
171            },
172        );
173        source = source.replace("@ROUND@", match sampling {
174                    RgbSampling::BlackTiesEven => "%rounded = scalar.roundevenf %value : f32",
175                    RgbSampling::ReplicateHalfUp => "%half = scalar.constant 0.5 : f32\n%shifted = scalar.addf %value, %half : f32\n%rounded = scalar.floorf %shifted : f32",
176                });
177        let mut model = ModelSession::in_context(&self.context)?;
178        let src = model.allocate(input.bytes())?;
179        let maps = model.allocate(matrices.bytes())?;
180        let dst = model.allocate(output.bytes())?;
181        let mut inputs = vec![(src, input.clone()), (maps, matrices)];
182        let mut args = vec![src.read(), maps.read(), dst.write()];
183        if packed {
184            let desc = TensorDesc::new(DType::U32, vec![crops, 3])?;
185            let geometry = model.allocate(desc.bytes())?;
186            inputs.push((geometry, desc));
187            args.push(geometry.read());
188        }
189        // Kernel loads are guarded before conversion and bounded to these extents.
190        let kernel = unsafe { model.compile(&[(&source, Specialization::new("affine_rgb"))])? }[0];
191        unsafe {
192            model.freeze(&self.context)?.fragment(
193                &[Command::Dispatch(Dispatch::indices(
194                    kernel,
195                    [0],
196                    [count.div_ceil(256) as u32, 1, 1],
197                    args,
198                ))],
199                &inputs,
200                &[(dst, output)],
201            )
202        }
203    }
204
205    /// Sample resident RGB into resident crops without a host wait or readback.
206    pub fn affine_rgb(
207        &self,
208        image: &DeviceTensor,
209        inverse: &DeviceTensor,
210        height: usize,
211        width: usize,
212        sampling: RgbSampling,
213    ) -> Result<Inference> {
214        self.context.validate(image)?;
215        self.context.validate(inverse)?;
216        let shape = inverse.desc().shape();
217        if shape.len() != 3
218            || shape[1..] != [2, 3]
219            || inverse.desc().dtype() != DType::F32
220            || inverse.desc().layout() != Layout::General
221            || !inverse.desc().is_contiguous()
222        {
223            return Err(Error::Message(
224                "affine matrices must be contiguous F32 [crops,2,3] with General layout".into(),
225            ));
226        }
227        self.prepare_affine_rgb(image.desc(), shape[0], height, width, sampling)?
228            .submit(&[image.clone(), inverse.clone()])
229    }
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235    #[test]
236    fn inverse_validation() {
237        assert_eq!(
238            invert_affine([[2., 0., 4.], [0., 2., -6.]]).unwrap(),
239            [[0.5, -0., -2.], [-0., 0.5, 3.]]
240        );
241        for t in [[[0.; 3]; 2], [[f32::NAN; 3]; 2], [[f32::MAX; 3]; 2]] {
242            assert!(invert_affine(t).is_err());
243        }
244    }
245}
hrx/image/affine.rs

hrx/image/
affine.rs
