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, WithBoard,
WithMetadata, 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.image_records.image_records_common import ImageCategory
from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.app.services.shared.invocation_context import InvocationContext
from invokeai.backend.image_util.invisible_watermark import InvisibleWatermark from invokeai.backend.image_util.invisible_watermark import InvisibleWatermark
@ -542,6 +542,25 @@ class ImageWatermarkInvocation(BaseInvocation, WithMetadata, WithBoard):
return ImageOutput.build(image_dto) 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( @invocation(
"mask_edge", "mask_edge",
title="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 cv2
import numpy as np import numpy as np
from imwatermark import WatermarkEncoder from imwatermark import WatermarkDecoder, WatermarkEncoder
from PIL import Image from PIL import Image
import invokeai.backend.util.logging as logger import invokeai.backend.util.logging as logger
@ -28,3 +28,10 @@ class InvisibleWatermark:
encoder.set_watermark("bytes", watermark_text.encode("utf-8")) encoder.set_watermark("bytes", watermark_text.encode("utf-8"))
bgr_encoded = encoder.encode(bgr, "dwtDct") bgr_encoded = encoder.encode(bgr, "dwtDct")
return Image.fromarray(cv2.cvtColor(bgr_encoded, cv2.COLOR_BGR2RGB)).convert("RGBA") 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")