Update utils.py

utils.py

@@ -5,90 +5,107 @@ import os
 import uuid
 
 def to_bin(data):
+    """Convert data to binary format"""
     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")
+        return format(data, '08b')
     else:
         raise TypeError("Type not supported.")
 
 def encode(image_path, secret_data):
+    """Encode watermark into image"""
     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
+        max_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 += "====="
+        # Add delimiters to secret data
+        secret_data = f"START{secret_data}END"
+        
+        # Check if the data can fit in the image
+        if len(secret_data) * 8 > max_bytes:
+            return None, "Watermark is too large for image size"
         
         # Convert data to binary
-        binary_secret_data = to_bin(secret_data_with_delimiter)
+        binary_secret_data = to_bin(secret_data)
         data_len = len(binary_secret_data)
+        
+        # Create output image
+        watermarked_image = image.copy()
         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
+        # Embed the data
+        for i in range(watermarked_image.shape[0]):
             if data_index >= data_len:
                 break
-                
-        return image, None
+            for j in range(watermarked_image.shape[1]):
+                if data_index >= data_len:
+                    break
+                for k in range(3):
+                    if data_index >= data_len:
+                        break
+                    # Get binary pixel value
+                    binary_pixel = to_bin(watermarked_image[i, j, k])
+                    # Replace least significant bit
+                    new_binary = binary_pixel[:-1] + binary_secret_data[data_index]
+                    watermarked_image[i, j, k] = int(new_binary, 2)
+                    data_index += 1
+        
+        return watermarked_image, None
     except Exception as e:
         return None, f"Error during encoding: {str(e)}"
 
 def decode(image_path):
+    """Decode watermark from image"""
     try:
         # Read the image
-        image = cv2.imread(image_path)
+        if isinstance(image_path, str):
+            image = cv2.imread(image_path)
+        else:
+            # Handle PIL Image input
+            image = cv2.cvtColor(np.array(Image.open(image_path)), cv2.COLOR_RGB2BGR)
+            
         if image is None:
             return "Failed to read image"
-            
+        
         binary_data = ""
-        for row in image:
-            for pixel in row:
-                for color_channel in pixel:
-                    binary_data += to_bin(color_channel)[-1]
+        
+        # Extract LSB from each pixel
+        for i in range(image.shape[0]):
+            for j in range(image.shape[1]):
+                for k in range(3):
+                    binary_pixel = to_bin(image[i, j, k])
+                    binary_data += binary_pixel[-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
+        text_data = ""
+        for i in range(0, len(binary_data), 8):
+            byte = binary_data[i:i+8]
+            if len(byte) == 8:
+                text_data += chr(int(byte, 2))
+                # Check for end marker
+                if "END" in text_data:
+                    # Extract the actual message between START and END
+                    start_idx = text_data.find("START")
+                    end_idx = text_data.find("END")
+                    if start_idx != -1 and end_idx != -1:
+                        return text_data[start_idx+5:end_idx]
         
-        # Remove delimiters
-        return decoded_data[:-5].split("#####")[0]
+        return "No watermark found"
     except Exception as e:
         return f"Error during decoding: {str(e)}"
 
 def png_encode(im_name, extra):
+    """Encode watermark into PNG metadata"""
     try:
         im = Image.open(im_name)
         info = PngImagePlugin.PngInfo()
@@ -96,7 +113,7 @@ def png_encode(im_name, extra):
 
         unique_id = str(uuid.uuid4())[:8]
         filename, ext = os.path.splitext(im_name)
-        new_filename = f"{filename}_{unique_id}{ext}"
+        new_filename = f"{filename}_{unique_id}.png"
         im.save(new_filename, "PNG", pnginfo=info)
 
         width, height = im.size
@@ -107,6 +124,5 @@ def png_encode(im_name, extra):
         png_encode_coords = (x, y, rect_width, rect_height)
 
         return new_filename, None
-
     except Exception as e:
-        return None, str(e)
+        return None, str(e)