Update utils.py

utils.py

@@ -1,36 +1,49 @@
-import cv2 
+import cv2
 import numpy as np
 from PIL import Image, PngImagePlugin
+import os
+import uuid
 
-def png_encode(im_name,extra):
-        # Check text roundtripping
-
+def png_encode(im_name, extra):
+    try:
         im = Image.open(im_name)
-
         info = PngImagePlugin.PngInfo()
         info.add_text("TXT", extra)
-        im.save("test.png",pnginfo=info)
-        test = Image.open("test.png")
-        print(test.text)
-        return("test.png","")
+
+        unique_id = str(uuid.uuid4())[:8]
+        filename, ext = os.path.splitext(im_name)
+        new_filename = f"{filename}_{unique_id}{ext}"
+        im.save(new_filename, pnginfo=info)
+
+        width, height = im.size
+        rect_width, rect_height = 200, 50
+        x = width - rect_width - 10
+        y = height - rect_height - 10
+        png_encode_coords = (x, y, rect_width, rect_height)
+
+        return new_filename, "", png_encode_coords
+
+    except Exception as e:
+        return None, str(e), None
 
 
 def to_bin(data):
-    """Convert `data` to binary format as string"""
     if isinstance(data, str):
-        return ''.join([ format(ord(i), "08b") for i in data ])
+        return ''.join([format(ord(i), "08b") for i in data])
     elif isinstance(data, bytes):
-        return ''.join([ format(i, "08b") for i in data ])
+        return ''.join([format(i, "08b") for i in data])
     elif isinstance(data, np.ndarray):
-        return [ format(i, "08b") for i in data ]
+        return [format(i, "08b") for i in data]
     elif isinstance(data, int) or isinstance(data, np.uint8):
         return format(data, "08b")
     else:
         raise TypeError("Type not supported.")
-def decode(image_name,txt=None):
+
+
+def decode(image_name, txt=None):
     try:
         BGRimage = cv2.imread(f"{image_name}")
-        image = cv2.cvtColor(BGRimage, cv2.COLOR_BGR2RGB)   
+        image = cv2.cvtColor(BGRimage, cv2.COLOR_BGR2RGB)
         binary_data = ""
         for row in image:
             for pixel in row:
@@ -38,34 +51,34 @@ def decode(image_name,txt=None):
                 binary_data += r[-1]
                 binary_data += g[-1]
                 binary_data += b[-1]
-        all_bytes = [ binary_data[i: i+8] for i in range(0, len(binary_data), 8) ]
+        all_bytes = [binary_data[i: i + 8] for i in range(0, len(binary_data), 8)]
         decoded_data = ""
         for byte in all_bytes:
             decoded_data += chr(int(byte, 2))
             if decoded_data[-5:] == "=====":
                 break
-        this = decoded_data[:-5].split("#####",1)[0]
-        #this = eval(this)
+        this = decoded_data[:-5].split("#####", 1)[0]
+        return this
     except Exception as e:
-        this = e
-    return this
+        return str(e)
+
 
-def encode(image_name, secret_data,txt=None):
-    msg=""
+def encode(image_name, secret_data, txt=None):
+    msg = ""
     BGRimage = cv2.imread(f"{image_name}")
-    image = cv2.cvtColor(BGRimage, cv2.COLOR_BGR2RGB)    
+    image = cv2.cvtColor(BGRimage, cv2.COLOR_BGR2RGB)
     n_bytes = image.shape[0] * image.shape[1] * 3 // 8
     print("[*] Maximum bytes to encode:", n_bytes)
-    secret_data1=secret_data
+    secret_data1 = secret_data
     while True:
-        if len(secret_data1)+5 < (n_bytes):
+        if len(secret_data1) + 5 < (n_bytes):
             secret_data1 = f'{secret_data1}#####'
-        elif len(secret_data1)+5 >= (n_bytes):
-            break    
-    secret_data = secret_data1    
+        elif len(secret_data1) + 5 >= (n_bytes):
+            break
+    secret_data = secret_data1
     if len(secret_data) > n_bytes:
-        msg ="Watermark is too large for Image Size"
-        return image_name,msg
+        msg = "Watermark is too large for Image Size"
+        return image_name, msg
     secret_data += "====="
     data_index = 0
     binary_secret_data = to_bin(secret_data)
@@ -84,4 +97,4 @@ def encode(image_name, secret_data,txt=None):
                 data_index += 1
             if data_index >= data_len:
                 break
-    return image,msg
+    return image, msg