hrx/image/
composite.rs

1use super::*;
2
3/// Quantization contract for affine RGB compositing.
4#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
5pub enum RgbComposite {
6    /// Interpolate the fractional crop, alpha-blend, then truncate to bytes.
7    Truncate,
8    /// Round sampled crop half-up, alpha-blend and truncate to bytes, then
9    /// blend that byte result with the original using ties-to-even rounding.
10    QuantizedBlend,
11}
12
13impl ImageOps {
14    /// Prepare affine compositing of one F32 NHWC RGB crop and its F32
15    /// `[height,width]` alpha mask into one U8 NHWC frame. RGB uses replicated
16    /// borders; mask sampling uses zero borders and clamps alpha to `[0,1]`.
17    ///
18    /// Inputs: frame, crop, mask, F32 `[1,2,3]` frame-to-crop matrix, F32 `[5]`
19    /// containing `[left,top,right,bottom,blend]`. Region bounds are exclusive
20    /// on right/bottom. Pixels outside it (or with invalid coordinates) are
21    /// unchanged. Blend is clamped to `[0,1]` and used only for QuantizedBlend.
22    pub fn prepare_composite_rgb(
23        &self,
24        image: &TensorDesc,
25        crop: &TensorDesc,
26        mode: RgbComposite,
27    ) -> Result<Arc<PreparedModel>> {
28        self.composite_plans
29            .get_or_prepare((image.clone(), crop.clone(), mode), || {
30                self.composite_rgb_fragment(image, crop, mode)?.prepare(3)
31            })
32    }
33
34    /// Record compositing with directly bound frame, crop, mask, matrix and
35    /// region/blend inputs. Descriptors and arithmetic match `prepare_composite_rgb`.
36    pub fn composite_rgb_fragment(
37        &self,
38        image: &TensorDesc,
39        crop: &TensorDesc,
40        mode: RgbComposite,
41    ) -> Result<ModelFragment> {
42        for (desc, dtype) in [(image, DType::U8), (crop, DType::F32)] {
43            if desc.dtype() != dtype
44                || desc.layout() != Layout::Nhwc
45                || !desc.is_contiguous()
46                || desc.is_empty()
47                || desc.shape().len() != 4
48                || desc.shape()[0] != 1
49                || desc.shape()[3] != 3
50                || desc.elements() > i32::MAX as usize
51                || desc.shape()[1..3].iter().any(|&n| n > 1 << 24)
52            {
53                return Err(Error::Message(
54                    "composite requires one contiguous RGB frame and crop".into(),
55                ));
56            }
57        }
58        let mask = TensorDesc::new(DType::F32, crop.shape()[1..3].to_vec())?;
59        let matrix = TensorDesc::new(DType::F32, vec![1, 2, 3])?;
60        let region = TensorDesc::new(DType::F32, vec![5])?;
61        {
62            let count = image.elements();
63            let mut source = include_str!("composite_rgb.loom").to_owned();
64            for (name, value) in [
65                ("COUNT", count),
66                ("GRID", count.div_ceil(256)),
67                ("WIDTH", image.shape()[2]),
68                ("LAST_OUTPUT", count - 1),
69                ("LAST_CROP", crop.elements() - 1),
70                ("LAST_MASK", mask.elements() - 1),
71                ("LAST_X", crop.shape()[2] - 1),
72                ("LAST_Y", crop.shape()[1] - 1),
73                ("CROP", crop.elements()),
74                ("MASK", mask.elements()),
75                ("CROP_WIDTH", crop.shape()[2]),
76                ("CROP_HEIGHT", crop.shape()[1]),
77            ] {
78                source = source.replace(&format!("@{name}@"), &value.to_string());
79            }
80            source = source.replace("@CROP_ROUND@", if mode == RgbComposite::QuantizedBlend {
81                "%half = scalar.constant 0.5 : f32\n%shift = scalar.addf %face0, %half : f32\n%face = scalar.floorf %shift : f32"
82            } else { "%face = scalar.addf %face0, %fzero : f32" });
83            source = source.replace("@FINAL_BLEND@", if mode == RgbComposite::QuantizedBlend {
84                "%quantized = scalar.floorf %clamped : f32\n%blend_at = index.constant 4 : index\n%blend0 = view.load %rv[%blend_at] : view<5xf32> -> f32\n%blendlo = scalar.maxnumf %blend0, %fzero : f32\n%blend = scalar.minnumf %blendlo, %fone : f32\n%inverse = scalar.subf %fone, %blend : f32\n%before2 = scalar.mulf %original_float, %inverse : f32\n%after2 = scalar.mulf %quantized, %blend : f32\n%mixed = scalar.addf %before2, %after2 : f32\n%final = scalar.roundevenf %mixed : f32"
85            } else { "%final = scalar.addf %clamped, %fzero : f32" });
86            let mut model = ModelSession::in_context(&self.context)?;
87            let descs = [image.clone(), crop.clone(), mask, matrix, region];
88            let inputs = descs
89                .into_iter()
90                .map(|d| Ok((model.allocate(d.bytes())?, d)))
91                .collect::<Result<Vec<_>>>()?;
92            let dst = model.allocate(image.bytes())?;
93            let mut bindings = inputs.iter().map(|(r, _)| r.read()).collect::<Vec<_>>();
94            bindings.push(dst.write());
95            // The source guards float-to-integer conversion and bounds every access.
96            let kernel =
97                unsafe { model.compile(&[(&source, Specialization::new("composite_rgb"))])? }[0];
98            unsafe {
99                model.freeze(&self.context)?.fragment(
100                    &[Command::Dispatch(Dispatch::indices(
101                        kernel,
102                        [0],
103                        [count.div_ceil(256) as u32, 1, 1],
104                        bindings,
105                    ))],
106                    &inputs,
107                    &[(dst, image.clone())],
108                )
109            }
110        }
111    }
112
113    /// Composite resident pixels. Bounds and blend are small host metadata;
114    /// image, crop, mask and matrix never cross the host boundary.
115    #[allow(clippy::too_many_arguments)]
116    pub fn composite_rgb(
117        &self,
118        image: &DeviceTensor,
119        crop: &DeviceTensor,
120        mask: &DeviceTensor,
121        matrix: &DeviceTensor,
122        region: [usize; 4],
123        blend: f32,
124        mode: RgbComposite,
125    ) -> Result<Inference> {
126        for tensor in [image, crop, mask, matrix] {
127            self.context.validate(tensor)?;
128        }
129        let plan = self.prepare_composite_rgb(image.desc(), crop.desc(), mode)?;
130        let s = image.desc().shape();
131        if region[0] > region[2]
132            || region[1] > region[3]
133            || region[2] > s[2]
134            || region[3] > s[1]
135            || !blend.is_finite()
136            || !(0.0..=1.0).contains(&blend)
137        {
138            return Err(Error::Message("invalid composite region or blend".into()));
139        }
140        let bytes = region
141            .into_iter()
142            .map(|v| v as f32)
143            .chain([blend])
144            .flat_map(f32::to_le_bytes)
145            .collect::<Vec<_>>();
146        let metadata = self
147            .context
148            .upload(TensorDesc::new(DType::F32, vec![5])?, &bytes)?;
149        plan.submit(&[
150            image.clone(),
151            crop.clone(),
152            mask.clone(),
153            matrix.clone(),
154            metadata,
155        ])
156    }
157}
hrx/image/composite.rs

hrx/image/
composite.rs
