| """Brand watermark overlay for generated images and videos.
|
|
|
| A watermark is requested by passing a JSON spec (a string or an already-parsed
|
| dict) to a generation endpoint via the ``watermark`` parameter. When the spec is
|
| valid, a small badge — the brand logo with the site domain rendered beneath it —
|
| is composited onto the generated image (or every frame of the generated video,
|
| via a single ffmpeg overlay pass). An empty or invalid spec is a no-op, so any
|
| caller that omits the parameter is completely unaffected.
|
|
|
| Spec fields (JSON object; at least one of ``logo`` / ``domain`` must be present
|
| for the spec to count as "valid"):
|
|
|
| logo str logo image as an http(s) URL, a ``data:`` URI, or raw
|
| base64. Optional — a domain-only watermark is allowed.
|
| domain str text drawn under the logo, e.g. "nsfwailab.com". Optional.
|
| position str bottom-right (default), bottom-left, top-right, top-left,
|
| bottom-center or top-center.
|
| scale float badge width as a fraction of the media width (default 0.18,
|
| clamped to 0.05..0.6).
|
| opacity float overall badge opacity 0..1 (default 0.9).
|
| margin float badge margin from the edge as a fraction of media width
|
| (default 0.03).
|
|
|
| Every public entry point is wrapped so a watermarking failure degrades to
|
| "asset returned un-watermarked" rather than breaking generation.
|
| """
|
| from __future__ import annotations
|
|
|
| import base64
|
| import json
|
| import os
|
| import subprocess
|
| import tempfile
|
| import urllib.request
|
|
|
| _VIDEO_REF_W = 1024
|
|
|
|
|
|
|
| _logo_cache: dict[str, object] = {}
|
|
|
| _POSITIONS = {
|
| "bottom-right", "bottom-left", "top-right", "top-left",
|
| "bottom-center", "top-center",
|
| }
|
|
|
|
|
| def parse_spec(spec) -> dict | None:
|
| """Validate/normalise a watermark spec; return a clean dict or None.
|
|
|
| Accepts a JSON string or a dict. Returns None (a no-op signal) when the spec
|
| is empty, unparseable, or carries neither a logo nor a domain.
|
| """
|
| if not spec:
|
| return None
|
| data = spec
|
| if isinstance(spec, str):
|
| s = spec.strip()
|
| if not s:
|
| return None
|
| try:
|
| data = json.loads(s)
|
| except (ValueError, TypeError):
|
| return None
|
| if not isinstance(data, dict):
|
| return None
|
|
|
| logo = data.get("logo") or data.get("logo_url") or ""
|
| domain = (data.get("domain") or data.get("text") or "").strip()
|
| if not (logo or domain):
|
| return None
|
|
|
| def _f(key, default, lo, hi):
|
| try:
|
| return max(lo, min(hi, float(data.get(key, default))))
|
| except (TypeError, ValueError):
|
| return default
|
|
|
| position = str(data.get("position", "bottom-right")).strip().lower()
|
| if position not in _POSITIONS:
|
| position = "bottom-right"
|
|
|
| return {
|
| "logo": logo if isinstance(logo, str) else "",
|
| "domain": domain,
|
| "position": position,
|
| "scale": _f("scale", 0.18, 0.05, 0.6),
|
| "opacity": _f("opacity", 0.9, 0.1, 1.0),
|
| "margin": _f("margin", 0.03, 0.0, 0.2),
|
| }
|
|
|
|
|
| def is_valid(spec) -> bool:
|
| """True when ``spec`` would produce a watermark."""
|
| return parse_spec(spec) is not None
|
|
|
|
|
| def _load_logo(logo: str):
|
| """Decode a logo string (URL / data-URI / base64) into an RGBA PIL image."""
|
| if not logo:
|
| return None
|
| if logo in _logo_cache:
|
| return _logo_cache[logo]
|
| from PIL import Image
|
| import io
|
|
|
| img = None
|
| try:
|
| raw = None
|
| if logo.startswith("data:"):
|
| _, _, b64 = logo.partition(",")
|
| raw = base64.b64decode(b64)
|
| elif logo.startswith("http://") or logo.startswith("https://"):
|
| req = urllib.request.Request(logo, headers={"User-Agent": "watermark/1.0"})
|
| with urllib.request.urlopen(req, timeout=10) as resp:
|
| raw = resp.read()
|
| else:
|
| raw = base64.b64decode(logo)
|
| if raw:
|
| img = Image.open(io.BytesIO(raw)).convert("RGBA")
|
| except Exception as exc:
|
| print(f"[watermark] logo decode failed: {exc}")
|
| img = None
|
| _logo_cache[logo] = img
|
| return img
|
|
|
|
|
| _FONT_CANDIDATES = [
|
| "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
|
| "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
|
| "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf",
|
| "DejaVuSans-Bold.ttf",
|
| ]
|
|
|
|
|
| def _font(size: int):
|
| from PIL import ImageFont
|
|
|
| for path in _FONT_CANDIDATES:
|
| try:
|
| return ImageFont.truetype(path, size)
|
| except Exception:
|
| continue
|
| return ImageFont.load_default()
|
|
|
|
|
| def _render_badge(cfg: dict, media_w: int):
|
| """Build the watermark badge (logo + domain on a translucent plate)."""
|
| from PIL import Image, ImageDraw
|
|
|
| badge_w = max(72, int(media_w * cfg["scale"]))
|
| pad = max(6, badge_w // 16)
|
|
|
| logo = _load_logo(cfg["logo"])
|
| logo_img = None
|
| logo_h = 0
|
| if logo is not None and logo.width > 0:
|
| lw = max(1, badge_w - 2 * pad)
|
| lh = max(1, int(logo.height * lw / logo.width))
|
| logo_img = logo.resize((lw, lh), Image.LANCZOS)
|
| logo_h = lh
|
|
|
| domain = cfg["domain"]
|
| font = None
|
| text_w = text_h = 0
|
| bbox = (0, 0, 0, 0)
|
| if domain:
|
| measure = ImageDraw.Draw(Image.new("RGBA", (1, 1)))
|
| fs = max(11, int(badge_w * 0.14))
|
| font = _font(fs)
|
| bbox = measure.textbbox((0, 0), domain, font=font)
|
| text_w, text_h = bbox[2] - bbox[0], bbox[3] - bbox[1]
|
| max_tw = badge_w - 2 * pad
|
| if text_w > max_tw and text_w > 0:
|
| fs = max(8, int(fs * max_tw / text_w))
|
| font = _font(fs)
|
| bbox = measure.textbbox((0, 0), domain, font=font)
|
| text_w, text_h = bbox[2] - bbox[0], bbox[3] - bbox[1]
|
|
|
| gap = pad if (logo_img is not None and domain) else 0
|
| badge_h = logo_h + gap + text_h + 2 * pad
|
|
|
| badge = Image.new("RGBA", (badge_w, badge_h), (0, 0, 0, 0))
|
| draw = ImageDraw.Draw(badge)
|
| radius = max(8, badge_w // 12)
|
| draw.rounded_rectangle([0, 0, badge_w - 1, badge_h - 1], radius=radius,
|
| fill=(0, 0, 0, 110))
|
|
|
| y = pad
|
| if logo_img is not None:
|
| badge.alpha_composite(logo_img, ((badge_w - logo_img.width) // 2, y))
|
| y += logo_h + gap
|
| if domain:
|
| tx = (badge_w - text_w) // 2 - bbox[0]
|
| draw.text((tx, y - bbox[1]), domain, font=font, fill=(255, 255, 255, 235))
|
|
|
| opacity = cfg["opacity"]
|
| if opacity < 1.0:
|
| alpha = badge.split()[3].point(lambda v: int(v * opacity))
|
| badge.putalpha(alpha)
|
| return badge
|
|
|
|
|
| def _offset(position: str, media_w: int, media_h: int, bw: int, bh: int, margin: int):
|
| """Top-left pixel offset for the badge given a named position."""
|
| if "right" in position:
|
| x = media_w - bw - margin
|
| elif "left" in position:
|
| x = margin
|
| else:
|
| x = (media_w - bw) // 2
|
| y = margin if position.startswith("top") else media_h - bh - margin
|
| return max(0, x), max(0, y)
|
|
|
|
|
| def apply_to_image(img, spec):
|
| """Return ``img`` with the watermark composited on, or unchanged on no-op."""
|
| cfg = parse_spec(spec)
|
| if cfg is None:
|
| return img
|
| try:
|
| base = img.convert("RGBA")
|
| badge = _render_badge(cfg, base.width)
|
| margin = int(cfg["margin"] * base.width)
|
| x, y = _offset(cfg["position"], base.width, base.height,
|
| badge.width, badge.height, margin)
|
| base.alpha_composite(badge, (x, y))
|
| return base.convert("RGB")
|
| except Exception as exc:
|
| print(f"[watermark] image overlay failed: {exc}")
|
| return img
|
|
|
|
|
| def _video_dims(path: str):
|
| try:
|
| import cv2
|
|
|
| cap = cv2.VideoCapture(path)
|
| w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
| h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
| cap.release()
|
| if w > 0 and h > 0:
|
| return w, h
|
| except Exception:
|
| pass
|
|
|
|
|
| try:
|
| import re
|
|
|
| import imageio_ffmpeg
|
|
|
| ff = imageio_ffmpeg.get_ffmpeg_exe()
|
| err = subprocess.run([ff, "-i", path], capture_output=True, text=True).stderr
|
| m = re.search(r"Video:.*?(\d{2,5})x(\d{2,5})", err)
|
| if m:
|
| return int(m.group(1)), int(m.group(2))
|
| except Exception:
|
| pass
|
| return None
|
|
|
|
|
| def apply_to_video(in_path: str, spec, out_path: str | None = None) -> str:
|
| """Overlay the watermark on a video via one ffmpeg pass.
|
|
|
| Returns the path to the watermarked file, or the original ``in_path`` when
|
| the spec is a no-op or anything fails (so upload still proceeds).
|
| """
|
| cfg = parse_spec(spec)
|
| if cfg is None:
|
| return in_path
|
| badge_png = None
|
| try:
|
| import imageio_ffmpeg
|
|
|
| ffmpeg = imageio_ffmpeg.get_ffmpeg_exe()
|
| dims = _video_dims(in_path)
|
| media_w, media_h = dims if dims else (_VIDEO_REF_W, _VIDEO_REF_W)
|
|
|
| badge = _render_badge(cfg, media_w)
|
| margin = int(cfg["margin"] * media_w)
|
| x, y = _offset(cfg["position"], media_w, media_h,
|
| badge.width, badge.height, margin)
|
|
|
| fd, badge_png = tempfile.mkstemp(suffix=".png")
|
| os.close(fd)
|
| badge.save(badge_png)
|
|
|
| if out_path is None:
|
| fd, out_path = tempfile.mkstemp(suffix=".mp4")
|
| os.close(fd)
|
|
|
| cmd = [
|
| ffmpeg, "-y", "-i", in_path, "-i", badge_png,
|
| "-filter_complex", f"overlay={x}:{y}:format=auto",
|
| "-c:a", "copy", "-c:v", "libx264", "-pix_fmt", "yuv420p",
|
| "-crf", "18", "-preset", "veryfast", "-movflags", "+faststart",
|
| out_path,
|
| ]
|
| subprocess.run(cmd, check=True, capture_output=True)
|
| return out_path
|
| except Exception as exc:
|
| detail = exc.stderr.decode("utf-8", "ignore")[-400:] if isinstance(exc, subprocess.CalledProcessError) else exc
|
| print(f"[watermark] video overlay failed: {detail}")
|
| return in_path
|
| finally:
|
| if badge_png and os.path.exists(badge_png):
|
| try:
|
| os.unlink(badge_png)
|
| except OSError:
|
| pass
|
|
|
