wehub-resource-sync
cddb07a176
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
50 lines
2.0 KiB
Python
50 lines
2.0 KiB
Python
"""
|
|
This module defines a singleton object, "invisible_watermark" that
|
|
wraps the invisible watermark model. It respects the global "invisible_watermark"
|
|
configuration variable, that allows the watermarking to be supressed.
|
|
"""
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
import invokeai.backend.util.logging as logger
|
|
from invokeai.backend.image_util.imwatermark.vendor import WatermarkDecoder, WatermarkEncoder
|
|
|
|
|
|
class InvisibleWatermark:
|
|
"""
|
|
Wrapper around InvisibleWatermark module.
|
|
"""
|
|
|
|
@classmethod
|
|
def add_watermark(cls, image: Image.Image, watermark_text: str) -> Image.Image:
|
|
logger.debug(f'Applying invisible watermark "{watermark_text}"')
|
|
bgr = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)
|
|
encoder = WatermarkEncoder()
|
|
encoder.set_watermark("bytes", watermark_text.encode("utf-8"))
|
|
bgr_encoded = encoder.encode(bgr, "dwtDct")
|
|
return Image.fromarray(cv2.cvtColor(bgr_encoded, cv2.COLOR_BGR2RGB)).convert("RGBA")
|
|
|
|
@classmethod
|
|
def decode_watermark(cls, image: Image.Image, length: int = 8) -> str:
|
|
"""Attempt to decode an invisible watermark from an image.
|
|
|
|
Args:
|
|
image: The PIL Image to decode the watermark from.
|
|
length: The expected watermark length in bytes. Must match the length used when encoding.
|
|
The WatermarkDecoder requires the length in bits; this value is multiplied by 8 internally.
|
|
|
|
Returns:
|
|
The decoded watermark text, or an empty string if no watermark is detected or decoding fails.
|
|
"""
|
|
logger.debug("Attempting to decode invisible watermark")
|
|
try:
|
|
bgr = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)
|
|
decoder = WatermarkDecoder("bytes", length * 8)
|
|
watermark_bytes = decoder.decode(bgr, "dwtDct")
|
|
return watermark_bytes.decode("utf-8", errors="ignore").rstrip("\x00")
|
|
except Exception:
|
|
logger.debug("Failed to decode invisible watermark")
|
|
return ""
|