version 0.3.1: now we can choose different chr_type to get better per…

…formance

README.md

@@ -36,30 +36,55 @@ install
 ### embed message into text:
 
 ```python
-from text_blind_watermark import TextBlindWatermarkThin
+from text_blind_watermark import TextBlindWatermark2
 
 password = '20190808'
+text = '这句话中有盲水印，你能提取出来吗？'
 watermark = 'github.com/guofei9987'
-text_blind_wm = TextBlindWatermarkThin(password=password)
 
-wm = text_blind_wm.embed(watermark=watermark)
-# This is example，you can put wm everywhere
-text_embed = '这句话中有盲' + wm + '水印，你能提取出来吗？'
-print(text_embed)
+text_blind_wm = TextBlindWatermark2(password=password)
+
+text_with_wm = text_blind_wm.embed(text=text, watermark=watermark)
+print(text_with_wm)
 ```
 
 
 ### extract message from text
 
 ```python
-text_blind_wm_new = TextBlindWatermarkThin(password=password)
-wm_extract = text_blind_wm_new.extract(text_embed)
+text_blind_wm2 = TextBlindWatermark2(password=password)
+wm_extract = text_blind_wm2.extract(text_with_wm)
 print('提取内容：', wm_extract)
 ```
 
 >github.com/guofei9987
 
-## Method 2 is more robust
+
+
+### `chr_type`
+
+You can choose different type to get better performance.
+
+
+```python
+from text_blind_watermark import TextBlindWatermark2
+
+password = '20190808'
+text = '这句话中有盲水印，你能提取出来吗？'
+watermark = 'github.com/guofei9987'
+
+text_blind_wm = TextBlindWatermark2(password=password, chr_type=(3, 4))
+
+text_with_wm = text_blind_wm.embed(text=text, watermark=watermark)
+print(text_with_wm)
+
+text_blind_wm2 = TextBlindWatermark2(password=password, chr_type=(3, 4))
+wm_extract = text_blind_wm2.extract(text_with_wm)
+print('提取内容：', wm_extract)
+assert watermark == wm_extract
+```
+
+## Another algorithm is more robust
 
 ### Alice Put her text watermark into a text:
 

README_cn.md

@@ -31,31 +31,53 @@
 ### 把信息不可见地嵌入到文本中
 
 ```python
-from text_blind_watermark import TextBlindWatermarkThin
+from text_blind_watermark import TextBlindWatermark2
 
-# 密码
 password = '20190808'
-# 要嵌入的信息
+text = '这句话中有盲水印，你能提取出来吗？'
 watermark = 'github.com/guofei9987'
-text_blind_wm = TextBlindWatermarkThin(password=password)
 
-wm = text_blind_wm.embed(watermark=watermark)
-# This is example，you can put wm everywhere
-text_embed = '这句话中有盲' + wm + '水印，你能提取出来吗？'
-print(text_embed)
+text_blind_wm = TextBlindWatermark2(password=password)
+
+text_with_wm = text_blind_wm.embed(text=text, watermark=watermark)
+print(text_with_wm)
 ```
 
 
 ### 从文本中提取不可见的信息
 
 ```python
-text_blind_wm_new = TextBlindWatermarkThin(password=password)
-wm_extract = text_blind_wm_new.extract(text_embed)
+text_blind_wm2 = TextBlindWatermark2(password=password)
+wm_extract = text_blind_wm2.extract(text_with_wm)
 print('提取内容：', wm_extract)
 ```
 
 >github.com/guofei9987
 
+### `chr_type`
+
+可以指定 `chr_type` 使其在不同的系统环境中有更好的表现
+
+```python
+from text_blind_watermark import TextBlindWatermark2
+
+password = '20190808'
+text = '这句话中有盲水印，你能提取出来吗？'
+watermark = 'github.com/guofei9987'
+
+text_blind_wm = TextBlindWatermark2(password=password, chr_type=(3, 4))
+
+text_with_wm = text_blind_wm.embed(text=text, watermark=watermark)
+print(text_with_wm)
+
+text_blind_wm2 = TextBlindWatermark2(password=password, chr_type=(3, 4))
+wm_extract = text_blind_wm2.extract(text_with_wm)
+print('提取内容：', wm_extract)
+assert watermark == wm_extract
+```
+
+
+
 ## 更稳定的版本
 ### 张三：把隐秘消息嵌入到另一段文本中
 

examples/example_wm2.py

@@ -0,0 +1,33 @@
+from text_blind_watermark import TextBlindWatermark2
+
+password = '20190808'
+text = '这句话中有盲水印，你能提取出来吗？'
+watermark = 'github.com/guofei9987'
+
+text_blind_wm = TextBlindWatermark2(password=password)
+
+text_with_wm = text_blind_wm.embed(text=text, watermark=watermark)
+print(text_with_wm)
+
+# %%
+
+text_blind_wm2 = TextBlindWatermark2(password=password)
+wm_extract = text_blind_wm2.extract(text_with_wm)
+print('提取内容：', wm_extract)
+
+assert watermark == wm_extract
+
+# %% put watermark by index
+text_blind_wm3 = TextBlindWatermark2(password=password)
+
+text_with_wm = text_blind_wm2.embed(text=text, watermark=watermark, idx=0)
+assert text_blind_wm2.extract(text_with_wm) == wm_extract
+
+text_with_wm = text_blind_wm2.embed(text=text, watermark=watermark, idx=len(text))
+assert text_blind_wm2.extract(text_with_wm) == wm_extract
+
+# %%
+
+text_blind_wm4 = TextBlindWatermark2(password=password, chr_type=(3, 4))
+text_with_wm = text_blind_wm.embed(text=text, watermark=watermark)
+print(text_with_wm)

examples/example_wm2_type.py

@@ -0,0 +1,16 @@
+from text_blind_watermark import TextBlindWatermark2
+
+password = '20190808'
+text = '这句话中有盲水印，你能提取出来吗？'
+watermark = 'github.com/guofei9987'
+
+text_blind_wm = TextBlindWatermark2(password=password, chr_type=(3, 4))
+
+text_with_wm = text_blind_wm.embed(text=text, watermark=watermark)
+print(text_with_wm)
+
+text_blind_wm2 = TextBlindWatermark2(password=password, chr_type=(3, 4))
+wm_extract = text_blind_wm2.extract(text_with_wm)
+print('提取内容：', wm_extract)
+
+assert watermark == wm_extract

setup.py

@@ -29,8 +29,8 @@ def read_requirements(filename):
     author='Guo Fei',
     author_email='guofei9987@foxmail.com',
     license='MIT',
-    packages=find_packages(),  # 也可以是一个列表，例如 ['sko']
+    packages=find_packages(),
     platforms=['linux', 'windows', 'macos'],
-    install_requires=[],  # 指定此包的依赖
-    zip_safe=False,  # 为了兼容性，一般填 False
+    install_requires=[],
+    zip_safe=False,
 )

text_blind_watermark/__init__.py

@@ -1,3 +1,3 @@
-from .text_blind_watermark import TextBlindWatermark, TextBlindWatermarkThin
+from .text_blind_watermark import TextBlindWatermark, TextBlindWatermark2
 
-__version__ = '0.2.1'
+__version__ = '0.3.1'

text_blind_watermark/text_blind_watermark.py

@@ -61,36 +61,49 @@ def extract(self, text_embed):
                       range(len(wm_extract_bin) // 8)]).decode('utf-8')
 
 
-class TextBlindWatermarkThin:
-    def __init__(self, password):
-        self.bit2char_dict = {'0': chr(29), '1': chr(127)}
-        self.char2bit_dict = {chr(29): '0', chr(127): '1'}
+class TextBlindWatermark2:
+    def __init__(self, password, chr_type=(0, 1)):
+        all_chr_wm_hex = ('1d', '7F', '200B', '200C', '200D')
+        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]}
+        self.char2bit_dict = {chr_wm[0]: '0', chr_wm[1]: '1'}
         self.password = password
 
-    def embed(self, watermark):
+    def get_wm(self, watermark: str):
         random.seed(self.password)
         wm_bin = [format(i ^ random.randint(0, 255), '08b') for i in watermark.encode('utf-8')]  # 8位2进制格式
         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):
+        wm = self.get_wm(watermark)
+        if idx is None:
+            idx = random.randint(0, len(text))
+        else:
+            assert idx <= len(text)
+
+        return text[:idx] + wm + text[idx:]
+
     def extract(self, text_embed):
 
-        idx_start, idx_end = 0, 0
-        find_left, find_right = False, False
+        idx_left, idx_right = None, None
+
         for idx, char in enumerate(text_embed):
             if char in self.char2bit_dict:
-                if not find_left:
-                    idx_start, find_left = idx, True
+                if idx_left is None:
+                    idx_left = idx
             else:
-                if find_left and not find_right:
-                    idx_end, find_right = idx, True
-
-            if find_left and find_right:
-                break
+                if idx_left is not None and idx_right is None:
+                    idx_right = idx
+                    break
         else:
-            idx_end = len(text_embed)
+            idx_right = len(text_embed)
+
+        if idx_left is None or idx_right is None:
+            raise IOError("There is no watermark!")
 
-        wm_extract_bin = ''.join(self.char2bit_dict[i] for i in text_embed[idx_start:idx_end])
+        wm_extract_bin = ''.join(self.char2bit_dict[i] for i in text_embed[idx_left:idx_right])
 
         random.seed(self.password)