Update utils.py

utils.py

@@ -132,55 +132,50 @@ class WatermarkProcessor:
                if image is None:
                     raise ValueError("Could not read image file")
 
-               # Extract LSB data
-               binary_text = ''
-               # Get LSB from each byte of each pixel
+               # Extract LSB from image
+               binary_data = []
                for i in range(image.shape[0]):
                     for j in range(image.shape[1]):
-                         for k in range(3):
-                              bit = image[i, j, k] & 1
-                              binary_text += str(bit)
-
-               # Convert binary to text
-               text = ''
-               # Process 8 bits at a time
-               for i in range(0, len(binary_text)-8, 8):
-                    byte = binary_text[i:i+8]
-                    text += chr(int(byte, 2))
-                    # Check for end marker
-                    if text[-7:] == "###END###":
-                         extracted_text = text[:-7]
-                         try:
-                              # Attempt to parse as JSON
-                              data = json.loads(extracted_text)
-                              return json.dumps(data, ensure_ascii=False, indent=2)
-                         except:
-                              # If not valid JSON, return as plain text
-                              return extracted_text.strip()
-
-               # If we got here, try to find any valid JSON in the text
-               try:
-                    # Convert remaining bits
-                    for i in range(0, len(binary_text), 8):
-                         if i + 8 <= len(binary_text):
-                              byte = binary_text[i:i+8]
-                              text += chr(int(byte, 2))
-                    
-                    # Look for JSON markers
-                    start = text.find('{')
-                    end = text.rfind('}')
-                    if start != -1 and end != -1:
-                         json_text = text[start:end+1]
+                         for k in range(3):  # RGB channels
+                              binary_data.append(str(image[i, j, k] & 1))
+
+               # Convert bits 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])
+                         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)]
+                    try:
+                         # Try to decode as UTF-8
+                         decoded_text = data.decode('utf-8', errors='ignore')
                          try:
-                              data = json.loads(json_text)
-                              return json.dumps(data, ensure_ascii=False, indent=2)
-                         except:
-                              pass
-
-               except:
-                    pass
-
-               return text if text.strip() else "Error: No valid watermark found"
+                              # 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"
+               
+               return "Error: No valid watermark found"
 
           except Exception as e:
                return f"Error in decoding: {str(e)}"