picture_watermark.py 1.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  1. """
  2. 图片嵌入水印工具类
  3. """
  4. from watermark_generate.blind_watermark import WaterMark
  5. class PictureWatermarkEmbeder:
  6. def __init__(self, secret):
  7. """
  8. 初始化图片盲水印嵌入器
  9. :param secret: 密码标签
  10. """
  11. self.bwm = WaterMark(password_img=1, password_wm=1)
  12. self.bwm.read_wm(secret, mode='str')
  13. self.secret = secret
  14. self.dest_img = None
  15. def embed(self, src_img, dest_img):
  16. """
  17. 盲水印嵌入方法
  18. :param src_img: 原始图片位置
  19. :param dest_img: 嵌入图片存放位置
  20. """
  21. self.bwm.read_img(src_img)
  22. self.bwm.embed(dest_img)
  23. self.dest_img = dest_img
  24. def verify(self) -> bool:
  25. """
  26. 验证嵌入结果,检查嵌入的字符串与提取字符串是否一致
  27. :return: 嵌入结果
  28. """
  29. wm_extract = self.bwm.extract(self.dest_img, wm_shape=self.bwm.wm_size, mode='str')
  30. print(wm_extract)
  31. return wm_extract == self.secret
  32. def extract(embed_img, secret_len=512):
  33. bwm = WaterMark(password_img=1, password_wm=1)
  34. # todo 根据生成的密码标签长度修改此处计算水印长度函数
  35. secret = bwm.extract(embed_img, wm_shape=get_wm_bit(secret_len), mode='str')
  36. return secret
  37. def get_wm_bit(len):
  38. return len * 8 - 2