Update utils.py

utils.py

@@ -116,7 +116,7 @@ class WatermarkProcessor:
             return image_path, f"Error in encoding: {str(e)}"
 
     def decode(self, image_path):
-        """Decode watermark with improved error handling"""
+        """Decode watermark with simplified robust error handling"""
         try:
             # Try PNG metadata method first
             try:
@@ -139,8 +139,8 @@ class WatermarkProcessor:
                 for pixel in row:
                     for value in pixel:
                         binary_data += format(value, '08b')[-1]
-            
-            # Convert to bytes with error checking
+
+            # Convert binary string to bytes
             bytes_data = bytearray()
             for i in range(0, len(binary_data), 8):
                 if i + 8 <= len(binary_data):
@@ -149,32 +149,23 @@ class WatermarkProcessor:
                         bytes_data.append(byte)
                     except ValueError:
                         continue
-            
-            # Process the data with robust decoding
+
+            # Find end marker in raw bytes
+            end_sequence = b'====='
+            if end_sequence in bytes_data:
+                bytes_data = bytes_data[:bytes_data.index(end_sequence)]
+
+            # Find delimiter in raw bytes
+            delimiter = b'#####'
+            if delimiter in bytes_data:
+                bytes_data = bytes_data[:bytes_data.index(delimiter)]
+
             try:
-                # First try UTF-8 decoding with error handling
-                decoded_data = bytes(bytes_data).decode('utf-8', errors='ignore')
-                
-                # Clean up decoded data
-                decoded_data = ''.join(char for char in decoded_data if ord(char) < 128 or ord(char) > 160)
+                # Try direct UTF-8 decoding first
+                decoded_data = bytes_data.decode('utf-8')
                 
-                # Find markers
-                end_marker = "====="
-                delimiter = "#####"
-                
-                if end_marker in decoded_data:
-                    parts = decoded_data.split(end_marker)
-                    decoded_data = parts[0]  # Take the part before =====
-                    
-                if delimiter in decoded_data:
-                    parts = decoded_data.split(delimiter)
-                    decoded_data = parts[0]  # Take the part before #####
-                    
-                # Clean the decoded data
-                decoded_data = decoded_data.strip()
-                
-                # Try to parse as JSON
                 try:
+                    # Try to parse as JSON
                     watermark_data = json.loads(decoded_data)
                     # Verify image hash
                     image_copy = image.copy() & 0xFE
@@ -183,35 +174,29 @@ class WatermarkProcessor:
                     if current_hash != watermark_data.get('image_hash'):
                         return "Warning: Image has been modified after watermarking"
                     
+                    # Return properly formatted JSON
                     return json.dumps(watermark_data, indent=2, ensure_ascii=False)
                 except json.JSONDecodeError:
-                    # If not valid JSON, return cleaned decoded data
-                    return decoded_data
+                    # If not JSON, return decoded text
+                    return decoded_data.strip()
                     
             except UnicodeDecodeError:
-                # Try alternative decoding method
-                try:
-                    fallback_decoded = ""
-                    current_bytes = bytearray()
-                    
-                    for byte in bytes_data:
-                        current_bytes.append(byte)
+                # If UTF-8 decoding fails, try character by character
+                result = []
+                i = 0
+                while i < len(bytes_data):
+                    for j in range(min(4, len(bytes_data) - i), 0, -1):
                         try:
-                            char = current_bytes.decode('utf-8')
-                            fallback_decoded += char
-                            current_bytes = bytearray()
+                            char = bytes_data[i:i+j].decode('utf-8')
+                            result.append(char)
+                            i += j
+                            break
                         except UnicodeDecodeError:
-                            if len(current_bytes) >= 4:  # Maximum UTF-8 character length
-                                current_bytes = bytearray()
+                            if j == 1:  # If even single byte fails
+                                i += 1
                             continue
-                            
-                    if fallback_decoded:
-                        return fallback_decoded.strip()
-                    else:
-                        return "Error: Could not decode watermark data"
-                        
-                except Exception as e:
-                    return f"Error in fallback decoding: {str(e)}"
+                
+                return ''.join(result).strip()
             
         except Exception as e:
             return f"Error in decoding: {str(e)}"