Update utils.py

utils.py

@@ -4,29 +4,6 @@ from PIL import Image, PngImagePlugin
 import os
 import uuid
 
-def png_encode(im_name, extra):
-    try:
-        im = Image.open(im_name)
-        info = PngImagePlugin.PngInfo()
-        info.add_text("TXT", extra)
-
-        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):
     if isinstance(data, str):
         return ''.join([format(ord(i), "08b") for i in data])
@@ -39,62 +16,97 @@ def to_bin(data):
     else:
         raise TypeError("Type not supported.")
 
+def encode(image_path, secret_data):
+    try:
+        # Read the image
+        image = cv2.imread(image_path)
+        if image is None:
+            return None, "Failed to read image"
+            
+        # Calculate maximum bytes for encoding
+        n_bytes = image.shape[0] * image.shape[1] * 3 // 8
+        
+        # Check if the secret data is too large
+        secret_data_with_delimiter = f'{secret_data}#####'
+        if len(secret_data_with_delimiter) + 5 >= n_bytes:
+            return None, "Watermark is too large for Image Size"
+            
+        # Add ending delimiter
+        secret_data_with_delimiter += "====="
+        
+        # Convert data to binary
+        binary_secret_data = to_bin(secret_data_with_delimiter)
+        data_len = len(binary_secret_data)
+        data_index = 0
+        
+        # Encode the data into the image
+        for i in range(image.shape[0]):
+            for j in range(image.shape[1]):
+                pixel = image[i, j]
+                for color_channel in range(3):
+                    if data_index < data_len:
+                        # Get the binary value of the pixel
+                        binary_value = to_bin(pixel[color_channel])
+                        # Replace the least significant bit
+                        new_binary = binary_value[:-1] + binary_secret_data[data_index]
+                        # Update the pixel with the new value
+                        pixel[color_channel] = int(new_binary, 2)
+                        data_index += 1
+                image[i, j] = pixel
+                if data_index >= data_len:
+                    break
+            if data_index >= data_len:
+                break
+                
+        return image, None
+    except Exception as e:
+        return None, f"Error during encoding: {str(e)}"
 
-def decode(image_name, txt=None):
+def decode(image_path):
     try:
-        BGRimage = cv2.imread(f"{image_name}")
-        image = cv2.cvtColor(BGRimage, cv2.COLOR_BGR2RGB)
+        # Read the image
+        image = cv2.imread(image_path)
+        if image is None:
+            return "Failed to read image"
+            
         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]
+                for color_channel in pixel:
+                    binary_data += to_bin(color_channel)[-1]
+        
+        # Convert binary to text
         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]
-        return this
+        
+        # Remove delimiters
+        return decoded_data[:-5].split("#####")[0]
     except Exception as e:
-        return str(e)
+        return f"Error during decoding: {str(e)}"
 
+def png_encode(im_name, extra):
+    try:
+        im = Image.open(im_name)
+        info = PngImagePlugin.PngInfo()
+        info.add_text("TXT", extra)
 
-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
+        unique_id = str(uuid.uuid4())[:8]
+        filename, ext = os.path.splitext(im_name)
+        new_filename = f"{filename}_{unique_id}{ext}"
+        im.save(new_filename, "PNG", pnginfo=info)
+
+        width, height = im.size
+        rect_width, rect_height = 200, 50
+        x = width - rect_width - 10
+        y = height - rect_height - 10
+        global png_encode_coords
+        png_encode_coords = (x, y, rect_width, rect_height)
+
+        return new_filename, None
+
+    except Exception as e:
+        return None, str(e)