make clear that compression_ratio is int

guofei9987 · Dec 24, 2023 · 4aa9940 · 4aa9940
1 parent 01306ca
commit 4aa9940
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/blind_watermark/blind_watermark.py b/blind_watermark/blind_watermark.py
@@ -52,7 +52,7 @@ def read_wm(self, wm_content, mode='img'):
 
         self.bwm_core.read_wm(self.wm_bit)
 
-    def embed(self, filename=None, compression_ratio=None):
+    def embed(self, filename=None, compression_ratio: int = 100):
         """
         :param filename: string
             Save the image file as filename
@@ -62,6 +62,8 @@ def embed(self, filename=None, compression_ratio=None):
         :return:
         """
         embed_img = self.bwm_core.embed()
+        if compression_ratio == 100:
+            compression_ratio = None
         if filename is not None:
             if compression_ratio is None:
                 cv2.imwrite(filename=filename, img=embed_img)