Update utils.py

utils.py

@@ -64,6 +64,8 @@ class WatermarkProcessor:
           except Exception as e:
                return im_name, f"Error adding watermark: {str(e)}"
 
+# utils.py의 encode와 decode 메서드만 수정
+
      def encode(self, image_path, watermark_text, metadata=None):
           """Encode watermark using simple LSB steganography"""
           try:
@@ -78,14 +80,15 @@ class WatermarkProcessor:
                     'metadata': metadata or {}
                }
                
-               # Convert to string and add delimiter
-               secret_data = json.dumps(data, ensure_ascii=False) + "###END###"
-               
-               # Convert string to bytes then to binary
-               binary_secret = ''.join(format(x, '08b') for x in secret_data.encode('utf-8'))
+               # Convert data to string with end marker
+               json_str = json.dumps(data, ensure_ascii=False)
+               secret_data = json_str + "#####"
+
+               # Convert string to binary (using utf-8 encoding)
+               binary_data = ''.join(format(ord(char), '08b') for char in secret_data)
                
                # Check capacity
-               if len(binary_secret) > image.shape[0] * image.shape[1] * 3:
+               if len(binary_data) > image.shape[0] * image.shape[1] * 3:
                     return image_path, "Error: Image too small for watermark data"
 
                # Embed data
@@ -93,17 +96,15 @@ class WatermarkProcessor:
                for i in range(image.shape[0]):
                     for j in range(image.shape[1]):
                          for k in range(3):
-                              if data_index < len(binary_secret):
-                                   # Clear LSB and add watermark bit
-                                   pixel = image[i, j, k]
-                                   pixel = (pixel & 0xFE) | int(binary_secret[data_index])
-                                   image[i, j, k] = pixel
+                              if data_index < len(binary_data):
+                                   # Replace LSB with watermark bit
+                                   image[i, j, k] = np.uint8((image[i, j, k] & ~1) | int(binary_data[data_index]))
                                    data_index += 1
                               else:
                                    break
-                         if data_index >= len(binary_secret):
+                         if data_index >= len(binary_data):
                               break
-                    if data_index >= len(binary_secret):
+                    if data_index >= len(binary_data):
                          break
 
                # Save result
@@ -117,7 +118,7 @@ class WatermarkProcessor:
      def decode(self, image_path):
           """Decode watermark using simple LSB steganography"""
           try:
-               # Try PNG metadata method first
+               # Try PNG metadata first
                try:
                     im = Image.open(image_path)
                     if "TXT" in im.info:
@@ -125,44 +126,42 @@ class WatermarkProcessor:
                except:
                     pass
 
-               # Read image
                image = cv2.imread(image_path)
                if image is None:
                     raise ValueError("Could not read image file")
 
-               # Extract LSB data
-               binary_string = ''
+               # Extract binary data
+               binary_data = ''
                for i in range(image.shape[0]):
                     for j in range(image.shape[1]):
                          for k in range(3):
-                              binary_string += str(image[i, j, k] & 1)
+                              binary_data += 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 to parse as JSON
-                         data = json.loads(text)
-                         return json.dumps(data, ensure_ascii=False, indent=2)
-                    except json.JSONDecodeError:
-                         return text
+               text = ''
+               for i in range(0, len(binary_data), 8):
+                    if i + 8 <= len(binary_data):
+                         byte = binary_data[i:i+8]
+                         text += chr(int(byte, 2))
+                         
+                         # Check for end marker
+                         if "#####" in text:
+                              # Extract content before marker
+                              text = text.split("#####")[0]
+                              try:
+                                   # Parse JSON
+                                   data = json.loads(text)
+                                   return json.dumps(data, ensure_ascii=False, indent=2)
+                              except json.JSONDecodeError:
+                                   return text
+                              break
 
                return "Error: No valid watermark found"
 
           except Exception as e:
                return f"Error in decoding: {str(e)}"
 
+              
      def analyze_quality(self, original_path, watermarked_path):
           """Analyze watermark quality"""
           try: