Update utils.py

utils.py

@@ -21,7 +21,7 @@ class WatermarkProcessor:
      def to_bin(self, data):
           """Convert data to binary format as string"""
           if isinstance(data, str):
-               return ''.join(format(ord(char), '08b') for char in data)
+               return ''.join(format(x, '08b') for x in data.encode('utf-8'))
           elif isinstance(data, bytes):
                return ''.join(format(x, '08b') for x in data)
           elif isinstance(data, np.ndarray):
@@ -81,8 +81,8 @@ class WatermarkProcessor:
                # Convert to string and add delimiter
                secret_data = json.dumps(data, ensure_ascii=False) + "###END###"
                
-               # Convert to binary string
-               binary_secret = ''.join(format(ord(char), '08b') for char in secret_data)
+               # Convert string to bytes then to binary
+               binary_secret = ''.join(format(x, '08b') for x in secret_data.encode('utf-8'))
                
                # Check capacity
                if len(binary_secret) > image.shape[0] * image.shape[1] * 3:
@@ -94,12 +94,10 @@ class WatermarkProcessor:
                     for j in range(image.shape[1]):
                          for k in range(3):
                               if data_index < len(binary_secret):
-                                   # Get current pixel value
+                                   # Clear LSB and add watermark bit
                                    pixel = image[i, j, k]
-                                   # Clear the LSB by setting it to 0 and add the watermark bit
-                                   pixel = (pixel & 0xFE) | (int(binary_secret[data_index]) & 0x01)
-                                   # Ensure value stays within uint8 range
-                                   image[i, j, k] = np.uint8(pixel)
+                                   pixel = (pixel & 0xFE) | int(binary_secret[data_index])
+                                   image[i, j, k] = pixel
                                    data_index += 1
                               else:
                                    break
@@ -132,49 +130,44 @@ class WatermarkProcessor:
                if image is None:
                     raise ValueError("Could not read image file")
 
-               # Extract LSB from image
-               binary_data = []
+               # Extract LSB and build binary string
+               binary_data = ""
+               total_bits = image.shape[0] * image.shape[1] * 3
+
                for i in range(image.shape[0]):
                     for j in range(image.shape[1]):
-                         for k in range(3):  # RGB channels
-                              binary_data.append(str(image[i, j, k] & 1))
+                         for k in range(3):
+                              if len(binary_data) >= total_bits:
+                                   break
+                              binary_data += str(image[i, j, k] & 1)
 
-               # Convert bits to bytes
+               # Convert binary to bytes
                bytes_data = bytearray()
                for i in range(0, len(binary_data), 8):
                     if i + 8 <= len(binary_data):
-                         byte = ''.join(binary_data[i:i+8])
+                         byte = binary_data[i:i+8]
                          bytes_data.append(int(byte, 2))
 
-               # Find end marker in raw bytes
-               end_marker = b'###END###'
-               if end_marker in bytes_data:
-                    # Extract data before end marker
-                    data = bytes_data[:bytes_data.index(end_marker)]
+                    # Early stopping if we find the end marker
                     try:
-                         # Try to decode as UTF-8
-                         decoded_text = data.decode('utf-8', errors='ignore')
-                         try:
-                              # Try to parse as JSON
-                              json_data = json.loads(decoded_text)
-                              return json.dumps(json_data, ensure_ascii=False, indent=2)
-                         except json.JSONDecodeError:
-                              return decoded_text
-                    except UnicodeDecodeError:
-                         # If UTF-8 decoding fails, try decoding parts
-                         result = ''
-                         buffer = bytearray()
-                         for byte in data:
-                              buffer.append(byte)
-                              try:
-                                   result += buffer.decode('utf-8')
-                                   buffer.clear()
-                              except UnicodeDecodeError:
-                                   if len(buffer) >= 4:  # Max UTF-8 character length
-                                        buffer.clear()
-                                   continue
-                         return result if result else "Error: Invalid character encoding"
-               
+                         if b'###END###' in bytes_data:
+                              bytes_data = bytes_data[:bytes_data.index(b'###END###')]
+                              break
+                    except ValueError:
+                         continue
+
+               try:
+                    # Try UTF-8 decoding
+                    decoded_text = bytes_data.decode('utf-8')
+                    try:
+                         # Try parsing as JSON
+                         json_data = json.loads(decoded_text)
+                         return json.dumps(json_data, ensure_ascii=False, indent=2)
+                    except json.JSONDecodeError:
+                         return decoded_text
+               except UnicodeDecodeError:
+                    return "Error: Invalid character encoding"
+
                return "Error: No valid watermark found"
 
           except Exception as e: