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            ImportMem, Renderer, Texture,
18            element::{
19                Kind,
20                memory::{MemoryRenderBuffer, MemoryRenderBufferRenderElement},
21            },
22        },
23    },
24    utils::{Physical, Point, Rectangle, Transform},
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 color 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    /// The loaded image and its size, for a backend that hands it to the GPU once instead of
226    /// blending it into every frame's host pixels. Straight alpha, R,G,B,A per pixel.
227    pub fn sprite(&self) -> Option<(&[u8], u32, u32)> {
228        self.wm_loaded
229            .then_some((self.wm_pixels.as_slice(), self.wm_width, self.wm_height))
230    }
231
232    /// Where [`Self::update_position`] last put the image's top-left corner in the frame.
233    pub fn position(&self) -> (i32, i32) {
234        (self.wm_pos_x, self.wm_pos_y)
235    }
236
237    /// True when the watermark moves and must be re-rendered every frame.
238    pub fn is_animated(&self) -> bool {
239        self.is_animated
240    }
241
242    /// Place the watermark for the current frame size. Fixed anchors (corners / middle) are
243    /// pure geometry, but the `AN` anchor makes the watermark bounce, so each call also advances
244    /// that animation one step and reflects it off the frame edges — which is precisely why an `AN`
245    /// watermark has to be re-rendered every frame (see `is_animated`). `loc_enum` is the i32 form
246    /// of `WatermarkLocation`.
247    pub fn update_position(&mut self, frame_width: i32, frame_height: i32, loc_enum: i32) {
248        if !self.wm_loaded {
249            return;
250        }
251
252        let loc = WatermarkLocation::from(loc_enum);
253        let w = self.wm_width as i32;
254        let h = self.wm_height as i32;
255
256        self.wm_prev_pos = Some((self.wm_pos_x, self.wm_pos_y));
257        self.is_animated = matches!(loc, WatermarkLocation::AN);
258
259        match loc {
260            WatermarkLocation::TL => {
261                self.wm_pos_x = 0;
262                self.wm_pos_y = 0;
263            }
264            WatermarkLocation::TR => {
265                self.wm_pos_x = frame_width - w;
266                self.wm_pos_y = 0;
267            }
268            WatermarkLocation::BL => {
269                self.wm_pos_x = 0;
270                self.wm_pos_y = frame_height - h;
271            }
272            WatermarkLocation::BR => {
273                self.wm_pos_x = frame_width - w;
274                self.wm_pos_y = frame_height - h;
275            }
276            WatermarkLocation::MI => {
277                self.wm_pos_x = (frame_width - w) / 2;
278                self.wm_pos_y = (frame_height - h) / 2;
279            }
280            WatermarkLocation::AN => {
281                self.wm_subpixel_x += self.wm_velocity_x;
282                self.wm_subpixel_y += self.wm_velocity_y;
283
284                if self.wm_subpixel_x <= 0.0 {
285                    self.wm_subpixel_x = 0.0;
286                    self.wm_velocity_x = self.wm_velocity_x.abs();
287                } else if self.wm_subpixel_x + (w as f64) >= frame_width as f64 {
288                    self.wm_subpixel_x = (frame_width - w) as f64;
289                    self.wm_velocity_x = -self.wm_velocity_x.abs();
290                }
291
292                if self.wm_subpixel_y <= 0.0 {
293                    self.wm_subpixel_y = 0.0;
294                    self.wm_velocity_y = self.wm_velocity_y.abs();
295                } else if self.wm_subpixel_y + (h as f64) >= frame_height as f64 {
296                    self.wm_subpixel_y = (frame_height - h) as f64;
297                    self.wm_velocity_y = -self.wm_velocity_y.abs();
298                }
299
300                self.wm_pos_x = self.wm_subpixel_x as i32;
301                self.wm_pos_y = self.wm_subpixel_y as i32;
302            }
303            WatermarkLocation::None => {}
304        }
305    }
306
307    /// Render element for the watermark; `None` when no watermark is loaded.
308    pub fn get_watermark_element<R>(
309        &self,
310        renderer: &mut R,
311    ) -> Option<MemoryRenderBufferRenderElement<R>>
312    where
313        R: Renderer + ImportMem,
314        R::TextureId: Texture + Clone + Send + 'static,
315    {
316        if let Some(buffer) = &self.render_buffer {
317            let location =
318                Point::<f64, Physical>::from((self.wm_pos_x as f64, self.wm_pos_y as f64));
319            MemoryRenderBufferRenderElement::from_buffer(
320                renderer,
321                location,
322                buffer,
323                Some(1.0),
324                None,
325                None,
326                Kind::Unspecified,
327            )
328            .ok()
329        } else {
330            None
331        }
332    }
333
334    /// Render element for a software cursor `image` at logical `pos` on an output with
335    /// fractional `scale`; the position converts to physical here so the composited cursor
336    /// lands where the damage tracker and clients (which work in physical pixels) expect it.
337    pub fn get_cursor_element<R>(
338        &self,
339        renderer: &mut R,
340        image: xcursor::parser::Image,
341        pos: Point<f64, smithay::utils::Logical>,
342        scale: f64,
343    ) -> Option<MemoryRenderBufferRenderElement<R>>
344    where
345        R: Renderer + ImportMem,
346        R::TextureId: Texture + Clone + Send + 'static,
347    {
348        let buffer = MemoryRenderBuffer::from_slice(
349            &image.pixels_rgba,
350            Fourcc::Abgr8888,
351            (image.width as i32, image.height as i32),
352            1,
353            Transform::Normal,
354            None,
355        );
356
357        let hot: Point<i32, smithay::utils::Physical> =
358            (image.xhot as i32, image.yhot as i32).into();
359        let phys_pos = pos
360            .to_physical(smithay::utils::Scale::from(scale))
361            .to_i32_round();
362
363        MemoryRenderBufferRenderElement::from_buffer(
364            renderer,
365            (phys_pos - hot).to_f64(),
366            &buffer,
367            Some(1.0),
368            None,
369            None,
370            Kind::Cursor,
371        )
372        .ok()
373    }
374}