Update utils.py

utils.py

@@ -1,87 +1,119 @@
-import cv2 
+import cv2
 import numpy as np
 from PIL import Image, PngImagePlugin
 
-def png_encode(im_name,extra):
-        # Check text roundtripping
-
+def png_encode(im_name, extra):
+    """Encode watermark using PNG metadata"""
+    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","")
-
+        # Encode the text as UTF-8 before adding to metadata
+        info.add_text("TXT", extra.encode('utf-8').decode('utf-8'))
+        im.save("test.png", pnginfo=info)
+        return "test.png", "Watermark added successfully"
+    except Exception as e:
+        return im_name, f"Error adding watermark: {str(e)}"
 
 def to_bin(data):
-    """Convert `data` to binary format as string"""
+    """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):
+    """Decode watermark from image"""
     try:
-        BGRimage = cv2.imread(f"{image_name}")
-        image = cv2.cvtColor(BGRimage, cv2.COLOR_BGR2RGB)   
+        # First try PNG metadata method
+        try:
+            im = Image.open(image_name)
+            if "TXT" in im.info:
+                return im.info["TXT"]
+        except:
+            pass
+
+        # If PNG metadata fails, try steganography method
+        image = cv2.imread(image_name)
+        if image is None:
+            raise ValueError("Could not read image file")
+            
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
         binary_data = ""
+        
         for row in image:
             for pixel in row:
                 r, g, b = to_bin(pixel)
                 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) ]
+
+        # Process binary data in chunks of 8 bits
         decoded_data = ""
-        for byte in all_bytes:
-            decoded_data += chr(int(byte, 2))
-            if decoded_data[-5:] == "=====":
+        for i in range(0, len(binary_data), 8):
+            byte = binary_data[i:i+8]
+            if len(byte) == 8:  # Ensure we have a complete byte
+                try:
+                    decoded_data += chr(int(byte, 2))
+                except ValueError:
+                    continue
+
+            if decoded_data.endswith("====="):
                 break
-        this = decoded_data[:-5].split("#####",1)[0]
-        #this = eval(this)
+
+        if "#####" in decoded_data:
+            decoded_data = decoded_data.split("#####")[0]
+            
+        return decoded_data.strip().rstrip("=")
+
     except Exception as e:
-        this = e
-    return this
+        return f"Error detecting watermark: {str(e)}"
 
-def encode(image_name, secret_data,txt=None):
-    msg=""
-    BGRimage = cv2.imread(f"{image_name}")
-    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
-    while True:
-        if len(secret_data1)+5 < (n_bytes):
-            secret_data1 = f'{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
-    secret_data += "====="
-    data_index = 0
-    binary_secret_data = to_bin(secret_data)
-    data_len = len(binary_secret_data)
-    for row in image:
-        for pixel in row:
-            r, g, b = to_bin(pixel)
-            if data_index < data_len:
-                pixel[0] = int(r[:-1] + binary_secret_data[data_index], 2)
-                data_index += 1
-            if data_index < data_len:
-                pixel[1] = int(g[:-1] + binary_secret_data[data_index], 2)
-                data_index += 1
-            if data_index < data_len:
-                pixel[2] = int(b[:-1] + binary_secret_data[data_index], 2)
-                data_index += 1
-            if data_index >= data_len:
-                break
-    return image,msg
+def encode(image_name, secret_data, txt=None):
+    """Encode watermark using steganography"""
+    try:
+        image = cv2.imread(image_name)
+        if image is None:
+            raise ValueError("Could not read image file")
+            
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        
+        # Calculate maximum bytes that can be encoded
+        n_bytes = image.shape[0] * image.shape[1] * 3 // 8
+        
+        # Prepare the secret data
+        secret_data = str(secret_data).encode('utf-8').decode('utf-8')  # Ensure proper encoding
+        secret_data_with_delimiter = secret_data + "#####"
+        
+        if len(secret_data_with_delimiter) + 5 >= n_bytes:
+            return image_name, "Watermark is too large for Image Size"
+            
+        secret_data_with_delimiter += "====="
+        
+        # Convert secret data to binary
+        binary_secret_data = to_bin(secret_data_with_delimiter)
+        data_len = len(binary_secret_data)
+        data_index = 0
+        
+        # Embed the secret data
+        for i in range(image.shape[0]):
+            for j in range(image.shape[1]):
+                if data_index < data_len:
+                    # Modify the least significant bit of each color channel
+                    for k in range(3):  # For each color channel (RGB)
+                        if data_index < data_len:
+                            image[i, j, k] = int(to_bin(image[i, j, k])[:-1] + binary_secret_data[data_index], 2)
+                            data_index += 1
+                            
+        # Save the result
+        output_path = "watermarked_image.png"
+        cv2.imwrite(output_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
+        return output_path, "Watermark embedded successfully"
+        
+    except Exception as e:
+        return image_name, f"Error encoding watermark: {str(e)}"