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(x, '08b') for x in data.encode('utf-8'))
+               return ''.join(format(ord(char), '08b') for char in data)
           elif isinstance(data, bytes):
                return ''.join(format(x, '08b') for x in data)
           elif isinstance(data, np.ndarray):
@@ -130,43 +130,33 @@ class WatermarkProcessor:
                if image is None:
                     raise ValueError("Could not read image file")
 
-               # Extract LSB and build binary string
-               binary_data = ""
-               total_bits = image.shape[0] * image.shape[1] * 3
-
+               # Extract LSB data
+               binary_string = ''
                for i in range(image.shape[0]):
                     for j in range(image.shape[1]):
                          for k in range(3):
-                              if len(binary_data) >= total_bits:
-                                   break
-                              binary_data += str(image[i, j, k] & 1)
-
-               # Convert binary to bytes
-               bytes_data = bytearray()
-               for i in range(0, len(binary_data), 8):
-                    if i + 8 <= len(binary_data):
-                         byte = binary_data[i:i+8]
-                         bytes_data.append(int(byte, 2))
-
-                    # Early stopping if we find the end marker
-                    try:
-                         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')
+                              binary_string += str(image[i, j, k] & 1)
+
+               # Convert binary to text
+               chars = []
+               for i in range(0, len(binary_string), 8):
+                    byte = binary_string[i:i+8]
+                    if len(byte) == 8:
+                         chars.append(chr(int(byte, 2)))
+
+               # Join all characters
+               text = ''.join(chars)
+
+               # Find end marker
+               if "###END###" in text:
+                    text = text.split("###END###")[0]
+                    
                     try:
-                         # Try parsing as JSON
-                         json_data = json.loads(decoded_text)
-                         return json.dumps(json_data, ensure_ascii=False, indent=2)
+                         # Try to parse as JSON
+                         data = json.loads(text)
+                         return json.dumps(data, ensure_ascii=False, indent=2)
                     except json.JSONDecodeError:
-                         return decoded_text
-               except UnicodeDecodeError:
-                    return "Error: Invalid character encoding"
+                         return text
 
                return "Error: No valid watermark found"