Update utils.py

utils.py

@@ -22,44 +22,43 @@ 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
         
-        # Prepare secret data
+        # 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)
-        
-        # Create a copy of the image
-        watermarked_image = image.copy()
         data_index = 0
         
-        # 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
+        # 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 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)
+                        # 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
-                    else:
-                        break
+                image[i, j] = pixel
                 if data_index >= data_len:
                     break
             if data_index >= data_len:
                 break
-        
-        return watermarked_image, None
+                
+        return image, None
     except Exception as e:
         return None, f"Error during encoding: {str(e)}"
 
@@ -69,30 +68,23 @@ def decode(image_path):
         image = cv2.imread(image_path)
         if image is None:
             return "Failed to read image"
-        
+            
         binary_data = ""
-        # 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]
+        for row in image:
+            for pixel in row:
+                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 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"
+        for byte in all_bytes:
+            decoded_data += chr(int(byte, 2))
+            if decoded_data[-5:] == "=====":
+                break
         
+        # Remove delimiters
+        return decoded_data[:-5].split("#####")[0]
     except Exception as e:
         return f"Error during decoding: {str(e)}"
 
@@ -117,4 +109,4 @@ def png_encode(im_name, extra):
         return new_filename, None
 
     except Exception as e:
-        return None, str(e)
+        return None, str(e)