From 9ef61bd51b35e5cbc67061f289553e43297000ad Mon Sep 17 00:00:00 2001 From: guofei9987 Date: Sun, 28 Jan 2024 16:57:37 +0800 Subject: [PATCH] add remove_watermark, run it before embed --- text_blind_watermark/text_blind_watermark.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/text_blind_watermark/text_blind_watermark.py b/text_blind_watermark/text_blind_watermark.py index 2209f58..60c408a 100644 --- a/text_blind_watermark/text_blind_watermark.py +++ b/text_blind_watermark/text_blind_watermark.py @@ -62,8 +62,8 @@ def extract(self, text_embed): class TextBlindWatermark2: - def __init__(self, password, chr_type=(0, 1)): - all_chr_wm_hex = ('1d', '7F', '200B', '200C', '200D') + def __init__(self, password, chr_type=(4, 5)): + all_chr_wm_hex = ('1d', '7F', '200B', '200C', '200D', 'FEFF') chr_wm = [chr(int(all_chr_wm_hex[chr_idx], base=16)) for chr_idx in chr_type] self.bit2char_dict = {'0': chr_wm[0], '1': chr_wm[1]} @@ -76,7 +76,8 @@ def get_wm(self, watermark: str): wm_bin = ''.join(wm_bin) return ''.join(self.bit2char_dict[i] for i in wm_bin) - def embed(self, text: str, watermark: str, idx: int = None): + def embed(self, text: str, watermark: str, idx: int = None) -> str: + text = self.remove_watermark(text) # remove existing watermark before add new one wm = self.get_wm(watermark) if idx is None: idx = random.randint(0, len(text)) @@ -85,7 +86,7 @@ def embed(self, text: str, watermark: str, idx: int = None): return text[:idx] + wm + text[idx:] - def extract(self, text_embed): + def extract(self, text_embed: str) -> str: idx_left, idx_right = None, None @@ -109,3 +110,8 @@ def extract(self, text_embed): return bytes([int(wm_extract_bin[8 * i:8 * i + 8], base=2) ^ random.randint(0, 255) for i in range(len(wm_extract_bin) // 8)]).decode('utf-8') + + def remove_watermark(self, text_embed: str) -> str: + return (text_embed + .replace(self.bit2char_dict["0"], "") + .replace(self.bit2char_dict["1"], ""))