Skip to main content

pixelflux/encoders/
overlay.rs

1/*
2 * This Source Code Form is subject to the terms of the Mozilla Public
3 * License, v. 2.0. If a copy of the MPL was not distributed with this
4 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
5 */
6
7//! PNG watermark overlay composited onto captured frames before encoding.
8//!
9//! Supports static positioning (top-left, top-right, bottom-left, bottom-right, center) and an
10//! animated "DVD-screensaver" bounce mode. The watermark is rendered into the compositor frame on
11//! the GPU path (no readback) or blitted onto the host-ARGB buffer on the CPU path.
12
13use smithay::{
14    backend::{
15        allocator::Fourcc,
16        renderer::{
17            element::{
18                memory::{MemoryRenderBuffer, MemoryRenderBufferRenderElement},
19                Kind,
20            },
21            ImportMem, Renderer, Texture,
22        },
23    },
24    utils::{Point, Rectangle, Transform, Physical},
25};
26use std::path::Path;
27
28#[derive(Clone, Copy, PartialEq)]
29/// Watermark anchor: corners (TL/TR/BL/BR), middle (MI), or bouncing (AN).
30pub enum WatermarkLocation {
31    None = 0,
32    TL = 1,
33    TR = 2,
34    BL = 3,
35    BR = 4,
36    MI = 5,
37    AN = 6,
38}
39
40impl From<i32> for WatermarkLocation {
41    fn from(v: i32) -> Self {
42        match v {
43            1 => Self::TL,
44            2 => Self::TR,
45            3 => Self::BL,
46            4 => Self::BR,
47            5 => Self::MI,
48            6 => Self::AN,
49            _ => Self::None,
50        }
51    }
52}
53
54/// The watermark's uploaded pixels plus its current placement and bounce state, kept across
55/// frames so a moving (bouncing) watermark can be advanced and re-placed each tick without
56/// re-reading or re-uploading the image.
57pub struct OverlayState {
58    wm_width: u32,
59    wm_height: u32,
60    wm_pos_x: i32,
61    wm_pos_y: i32,
62    wm_prev_pos: Option<(i32, i32)>,
63    wm_velocity_x: f64,
64    wm_velocity_y: f64,
65    wm_subpixel_x: f64,
66    wm_subpixel_y: f64,
67    wm_loaded: bool,
68    is_animated: bool,
69    wm_pixels: Vec<u8>,
70    render_buffer: Option<MemoryRenderBuffer>,
71}
72
73impl Default for OverlayState {
74    fn default() -> Self {
75        Self {
76            wm_width: 0,
77            wm_height: 0,
78            wm_pos_x: 0,
79            wm_pos_y: 0,
80            wm_prev_pos: None,
81            wm_velocity_x: 2.0,
82            wm_velocity_y: 2.0,
83            wm_subpixel_x: 0.0,
84            wm_subpixel_y: 0.0,
85            wm_loaded: false,
86            is_animated: false,
87            wm_pixels: Vec::new(),
88            render_buffer: None,
89        }
90    }
91}
92
93/// Alpha-blend a source pixel (pre-split into r,g,b,a) over a BGRA destination pixel.
94///
95/// Overlay pixels are overwhelmingly either fully opaque or fully transparent, and this runs per
96/// pixel per frame on the CPU, so the two extremes are special-cased to skip the blend arithmetic
97/// entirely: an opaque source (`a == 255`) simply overwrites, a fully transparent source (`a == 0`)
98/// is left as-is, and only genuine edge pixels pay for the integer source-over. In every case only
99/// the B / G / R bytes are written; the destination's alpha byte is left as the capture delivered it.
100#[inline]
101pub(crate) fn blend_pixel(dst: &mut [u8], r: u8, g: u8, b: u8, a: u8) {
102    if a == 255 {
103        dst[0] = b;
104        dst[1] = g;
105        dst[2] = r;
106    } else if a > 0 {
107        let ia = 255 - a as u32;
108        dst[0] = ((b as u32 * a as u32 + dst[0] as u32 * ia) / 255) as u8;
109        dst[1] = ((g as u32 * a as u32 + dst[1] as u32 * ia) / 255) as u8;
110        dst[2] = ((r as u32 * a as u32 + dst[2] as u32 * ia) / 255) as u8;
111    }
112}
113
114/// Source-over compositing for PREMULTIPLIED sources (XFixes/Xcursor pixels are
115/// premultiplied by format definition): dst = src + dst*(1-a). The straight-alpha
116/// `blend_pixel` on premultiplied input multiplies alpha in twice and darkens
117/// every translucent pixel (visible dark fringes on the composited cursor).
118///
119/// The sum saturates: some toolkits ship cursors whose colour exceeds its alpha, and on those
120/// a wrapping cast would turn an over-bright pixel into a dark one.
121pub(crate) fn blend_pixel_premultiplied(dst: &mut [u8], r: u8, g: u8, b: u8, a: u8) {
122    if a == 255 {
123        dst[0] = b;
124        dst[1] = g;
125        dst[2] = r;
126    } else if a > 0 {
127        let ia = 255 - a as u32;
128        dst[0] = (b as u32 + dst[0] as u32 * ia / 255).min(255) as u8;
129        dst[1] = (g as u32 + dst[1] as u32 * ia / 255).min(255) as u8;
130        dst[2] = (r as u32 + dst[2] as u32 * ia / 255).min(255) as u8;
131    }
132}
133
134impl OverlayState {
135    /// Load the watermark image from disk; `output_scale` is the output's fractional
136    /// scale, ceiled to the integer buffer scale of the upload. A failed load clears the overlay.
137    pub fn load_watermark(&mut self, path: &str, output_scale: f64) {
138        if let Ok(img) = image::open(Path::new(path)) {
139            let rgba = img.to_rgba8();
140            self.wm_width = rgba.width();
141            self.wm_height = rgba.height();
142            self.wm_loaded = true;
143            let buffer_scale = output_scale.ceil().max(1.0) as i32;
144
145            let pixels = rgba.into_vec();
146            self.render_buffer = Some(MemoryRenderBuffer::from_slice(
147                &pixels,
148                Fourcc::Abgr8888,
149                (self.wm_width as i32, self.wm_height as i32),
150                buffer_scale,
151                Transform::Normal,
152                None,
153            ));
154            self.wm_pixels = pixels;
155        } else {
156            self.wm_loaded = false;
157            self.wm_pixels = Vec::new();
158            self.render_buffer = None;
159        }
160    }
161
162    /// Alpha-blend the watermark into a BGRA frame (row `stride` in bytes) at its
163    /// current position. Clips per pixel at the frame bounds because the animated
164    /// position — or a watermark larger than the capture — can leave part of the
165    /// image off-frame, and only the in-frame portion may be written.
166    pub fn blend_bgra(&self, frame: &mut [u8], stride: usize, frame_w: i32, frame_h: i32) {
167        if !self.wm_loaded {
168            return;
169        }
170        let (w, h) = (self.wm_width as i32, self.wm_height as i32);
171        for y in 0..h {
172            let ty = self.wm_pos_y + y;
173            if ty < 0 || ty >= frame_h {
174                continue;
175            }
176            for x in 0..w {
177                let tx = self.wm_pos_x + x;
178                if tx < 0 || tx >= frame_w {
179                    continue;
180                }
181                let src = ((y * w + x) * 4) as usize;
182                let (r, g, b, a) = (
183                    self.wm_pixels[src],
184                    self.wm_pixels[src + 1],
185                    self.wm_pixels[src + 2],
186                    self.wm_pixels[src + 3],
187                );
188                let off = ty as usize * stride + tx as usize * 4;
189                blend_pixel(&mut frame[off..off + 4], r, g, b, a);
190            }
191        }
192    }
193
194    /// Frame-clipped union of the watermark's current and previous rectangles —
195    /// the region a damage-gated encoder must repaint after a bounce step moved
196    /// the image (the vacated area needs repainting as much as the new one).
197    pub fn damage_rect(&self, frame_w: i32, frame_h: i32) -> Option<Rectangle<i32, Physical>> {
198        if !self.wm_loaded {
199            return None;
200        }
201        let (w, h) = (self.wm_width as i32, self.wm_height as i32);
202        let (mut x0, mut y0) = (self.wm_pos_x, self.wm_pos_y);
203        let (mut x1, mut y1) = (x0 + w, y0 + h);
204        if let Some((px, py)) = self.wm_prev_pos {
205            x0 = x0.min(px);
206            y0 = y0.min(py);
207            x1 = x1.max(px + w);
208            y1 = y1.max(py + h);
209        }
210        x0 = x0.max(0);
211        y0 = y0.max(0);
212        x1 = x1.min(frame_w);
213        y1 = y1.min(frame_h);
214        if x1 <= x0 || y1 <= y0 {
215            return None;
216        }
217        Some(Rectangle::new((x0, y0).into(), (x1 - x0, y1 - y0).into()))
218    }
219
220    /// True once a watermark image has been loaded.
221    pub fn is_active(&self) -> bool {
222        self.wm_loaded
223    }
224
225    /// True when the watermark moves and must be re-rendered every frame.
226    pub fn is_animated(&self) -> bool {
227        self.is_animated
228    }
229
230    /// Place the watermark for the current frame size. Fixed anchors (corners / middle) are
231    /// pure geometry, but the `AN` anchor makes the watermark bounce, so each call also advances
232    /// that animation one step and reflects it off the frame edges — which is precisely why an `AN`
233    /// watermark has to be re-rendered every frame (see `is_animated`). `loc_enum` is the i32 form
234    /// of `WatermarkLocation`.
235    pub fn update_position(&mut self, frame_width: i32, frame_height: i32, loc_enum: i32) {
236        if !self.wm_loaded {
237            return;
238        }
239
240        let loc = WatermarkLocation::from(loc_enum);
241        let w = self.wm_width as i32;
242        let h = self.wm_height as i32;
243
244        self.wm_prev_pos = Some((self.wm_pos_x, self.wm_pos_y));
245        self.is_animated = matches!(loc, WatermarkLocation::AN);
246
247        match loc {
248            WatermarkLocation::TL => {
249                self.wm_pos_x = 0;
250                self.wm_pos_y = 0;
251            }
252            WatermarkLocation::TR => {
253                self.wm_pos_x = frame_width - w;
254                self.wm_pos_y = 0;
255            }
256            WatermarkLocation::BL => {
257                self.wm_pos_x = 0;
258                self.wm_pos_y = frame_height - h;
259            }
260            WatermarkLocation::BR => {
261                self.wm_pos_x = frame_width - w;
262                self.wm_pos_y = frame_height - h;
263            }
264            WatermarkLocation::MI => {
265                self.wm_pos_x = (frame_width - w) / 2;
266                self.wm_pos_y = (frame_height - h) / 2;
267            }
268            WatermarkLocation::AN => {
269                self.wm_subpixel_x += self.wm_velocity_x;
270                self.wm_subpixel_y += self.wm_velocity_y;
271
272                if self.wm_subpixel_x <= 0.0 {
273                    self.wm_subpixel_x = 0.0;
274                    self.wm_velocity_x = self.wm_velocity_x.abs();
275                } else if self.wm_subpixel_x + (w as f64) >= frame_width as f64 {
276                    self.wm_subpixel_x = (frame_width - w) as f64;
277                    self.wm_velocity_x = -self.wm_velocity_x.abs();
278                }
279
280                if self.wm_subpixel_y <= 0.0 {
281                    self.wm_subpixel_y = 0.0;
282                    self.wm_velocity_y = self.wm_velocity_y.abs();
283                } else if self.wm_subpixel_y + (h as f64) >= frame_height as f64 {
284                    self.wm_subpixel_y = (frame_height - h) as f64;
285                    self.wm_velocity_y = -self.wm_velocity_y.abs();
286                }
287
288                self.wm_pos_x = self.wm_subpixel_x as i32;
289                self.wm_pos_y = self.wm_subpixel_y as i32;
290            }
291            WatermarkLocation::None => {}
292        }
293    }
294
295    /// Render element for the watermark; `None` when no watermark is loaded.
296    pub fn get_watermark_element<R>(
297        &self,
298        renderer: &mut R,
299    ) -> Option<MemoryRenderBufferRenderElement<R>>
300    where
301        R: Renderer + ImportMem,
302        R::TextureId: Texture + Clone + Send + 'static,
303    {
304        if let Some(buffer) = &self.render_buffer {
305            let location = Point::<f64, Physical>::from((self.wm_pos_x as f64, self.wm_pos_y as f64));
306            MemoryRenderBufferRenderElement::from_buffer(
307                renderer,
308                location,
309                buffer,
310                Some(1.0),
311                None,
312                None,
313                Kind::Unspecified,
314            )
315            .ok()
316        } else {
317            None
318        }
319    }
320
321    /// Render element for a software cursor `image` at logical `pos` on an output with
322    /// fractional `scale`; the position converts to physical here so the composited cursor
323    /// lands where the damage tracker and clients (which work in physical pixels) expect it.
324    pub fn get_cursor_element<R>(
325        &self,
326        renderer: &mut R,
327        image: xcursor::parser::Image,
328        pos: Point<f64, smithay::utils::Logical>,
329        scale: f64,
330    ) -> Option<MemoryRenderBufferRenderElement<R>>
331    where
332        R: Renderer + ImportMem,
333        R::TextureId: Texture + Clone + Send + 'static,
334    {
335        let buffer = MemoryRenderBuffer::from_slice(
336            &image.pixels_rgba,
337            Fourcc::Abgr8888,
338            (image.width as i32, image.height as i32),
339            1,
340            Transform::Normal,
341            None,
342        );
343
344        let hot: Point<i32, smithay::utils::Physical> =
345            (image.xhot as i32, image.yhot as i32).into();
346        let phys_pos = pos.to_physical(smithay::utils::Scale::from(scale)).to_i32_round();
347
348        MemoryRenderBufferRenderElement::from_buffer(
349            renderer,
350            (phys_pos - hot).to_f64(),
351            &buffer,
352            Some(1.0),
353            None,
354            None,
355            Kind::Cursor,
356        )
357        .ok()
358    }
359}