Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

add node to read watermarks

Browse Source
This commit is contained in:
Lincoln Stein 2024-07-17 15:03:00 -04:00
parent 0583101c1c
commit c70437d120
2 changed files with 28 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
invokeai/app/invocations/image.py
View File

@ -16,7 +16,7 @@ from invokeai.app.invocations.fields import (
WithBoard,
WithMetadata,
)
from invokeai.app.invocations.primitives import ImageOutput
from invokeai.app.invocations.primitives import ImageOutput, StringOutput
from invokeai.app.services.image_records.image_records_common import ImageCategory
from invokeai.app.services.shared.invocation_context import InvocationContext
from invokeai.backend.image_util.invisible_watermark import InvisibleWatermark
@ -542,6 +542,25 @@ class ImageWatermarkInvocation(BaseInvocation, WithMetadata, WithBoard):
return ImageOutput.build(image_dto)
@invocation(
"read_watermark",
title="Read Invisible Watermark from Image",
tags=["image", "watermark"],
category="image",
version="1.2.2",
)
class ReadWatermarkInvocation(BaseInvocation):
"""Read an invisible watermark from an image"""
image: ImageField = InputField(description="The image to read the watermark from")
watermark_len: int = InputField(default=8, description="length of watermark, in characters")
def invoke(self, context: InvocationContext) -> StringOutput:
image = context.images.get_pil(self.image.image_name)
watermark_text = InvisibleWatermark.read_watermark(image, self.watermark_len * 8)
return StringOutput(value=watermark_text)
@invocation(
"mask_edge",
title="Mask Edge",

9
invokeai/backend/image_util/invisible_watermark.py
View File

@ -6,7 +6,7 @@ configuration variable, that allows the watermarking to be supressed.
import cv2
import numpy as np
from imwatermark import WatermarkEncoder
from imwatermark import WatermarkDecoder, WatermarkEncoder
from PIL import Image
import invokeai.backend.util.logging as logger
@ -28,3 +28,10 @@ class InvisibleWatermark:
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 read_watermark(cls, image: Image.Image, length: int = 64) -> str:
bgr = cv2.cvtColor(np.array(image.convert("RGB")), cv2.COLOR_RGB2BGR)
decoder = WatermarkDecoder("bytes", length)
watermark = decoder.decode(bgr, "dwtDct")
return watermark.decode("utf-8")