Update utils.py

utils.py

@@ -22,43 +22,44 @@ def encode(image_path, secret_data):
         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
+        n_bytes = (image.shape[0] * image.shape[1] * 3) // 8
         
-        # Check if the secret data is too large
+        # Prepare secret data
         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
+        secret_data_with_delimiter += "====="
         binary_secret_data = to_bin(secret_data_with_delimiter)
         data_len = len(binary_secret_data)
+        
+        # Create a copy of the 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):
+        # Encode the data
+        for i in range(watermarked_image.shape[0]):
+            for j in range(watermarked_image.shape[1]):
+                for k in range(3):  # RGB channels
                     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)
+                        # Get the current pixel value
+                        pixel = watermarked_image[i, j, k]
+                        # Convert to binary and modify the least significant bit
+                        binary_pixel = format(pixel, '08b')
+                        new_pixel = binary_pixel[:-1] + binary_secret_data[data_index]
+                        # Update the pixel value
+                        watermarked_image[i, j, k] = int(new_pixel, 2)
                         data_index += 1
-                image[i, j] = pixel
+                    else:
+                        break
                 if data_index >= data_len:
                     break
             if data_index >= data_len:
                 break
-                
-        return image, None
+        
+        return watermarked_image, None
     except Exception as e:
         return None, f"Error during encoding: {str(e)}"
 
@@ -68,23 +69,30 @@ def decode(image_path):
         image = cv2.imread(image_path)
         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 the least significant bits
+        for i in range(image.shape[0]):
+            for j in range(image.shape[1]):
+                for k in range(3):
+                    pixel = format(image[i, j, k], '08b')
+                    binary_data += 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
+        for i in range(0, len(binary_data), 8):
+            byte = binary_data[i:i+8]
+            if len(byte) == 8:
+                decoded_data += chr(int(byte, 2))
+                # Check for ending delimiter
+                if decoded_data[-5:] == "=====":
+                    break
+        
+        # Remove delimiters and return the message
+        if "#####" in decoded_data:
+            return decoded_data.split("#####")[0]
+        return "No watermark found"
         
-        # Remove delimiters
-        return decoded_data[:-5].split("#####")[0]
     except Exception as e:
         return f"Error during decoding: {str(e)}"
 
@@ -109,4 +117,4 @@ def png_encode(im_name, extra):
         return new_filename, None
 
     except Exception as e:
-        return None, str(e)
+        return None, str(e)