Update utils.py

utils.py

@@ -63,7 +63,10 @@ class WatermarkProcessor:
             return im_name, f"Error adding watermark: {str(e)}"
 
     def encode(self, image_path, watermark_text, metadata=None):
-        """Encode watermark using simple LSB steganography"""
+        """Encode watermark using simple LSB steganography with header
+        
+        헤더(32비트)는 watermark 데이터(UTF-8 문자열)의 길이(문자수)를 이진 문자열로 저장합니다.
+        """
         try:
             image = cv2.imread(image_path)
             if image is None:
@@ -76,18 +79,20 @@ class WatermarkProcessor:
                 'metadata': metadata or {}
             }
             
-            # Convert data to string with end marker
+            # Convert data to string (UTF-8)
             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)
+            # 헤더: 32비트에 데이터 길이(문자수)를 저장
+            data_length = len(json_str)
+            header = format(data_length, '032b')
+            # 본문: 각 문자를 8비트 이진수로 변환
+            body = ''.join(format(ord(char), '08b') for char in json_str)
+            binary_data = header + body
             
             # Check capacity
             if len(binary_data) > image.shape[0] * image.shape[1] * 3:
                 return image_path, "Error: Image too small for watermark data"
 
-            # Embed data
+            # Embed data into LSB of each pixel
             data_index = 0
             for i in range(image.shape[0]):
                 for j in range(image.shape[1]):
@@ -95,8 +100,8 @@ class WatermarkProcessor:
                         if data_index < len(binary_data):
                             pixel = int(image[i, j, k])
                             # Clear the LSB
-                            pixel = pixel & 0xFE  # Clear last bit
-                            # Set the LSB according to our data
+                            pixel = pixel & 0xFE
+                            # Set the LSB according to our data bit
                             pixel = pixel | (int(binary_data[data_index]) & 1)
                             image[i, j, k] = np.uint8(pixel)
                             data_index += 1
@@ -107,7 +112,7 @@ class WatermarkProcessor:
                 if data_index >= len(binary_data):
                     break
 
-            # Save result
+            # Save result image (PNG: lossless)
             output_path = f"watermarked_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
             cv2.imwrite(output_path, image)
             return output_path, "Watermark added successfully"
@@ -116,9 +121,12 @@ class WatermarkProcessor:
             return image_path, f"Error in encoding: {str(e)}"
 
     def decode(self, image_path):
-        """Decode watermark using simple LSB steganography"""
+        """Decode watermark using simple LSB steganography with header
+        
+        먼저 32비트를 읽어 watermark 데이터의 길이(문자수)를 구한 후, 해당 길이만큼의 데이터를 추출합니다.
+        """
         try:
-            # Try PNG metadata first
+            # PNG 메타데이터 먼저 확인
             try:
                 im = Image.open(image_path)
                 if "TXT" in im.info:
@@ -130,34 +138,28 @@ class WatermarkProcessor:
             if image is None:
                 raise ValueError("Could not read image file")
 
-            # Extract binary data
+            # 모든 픽셀의 LSB를 읽어 이진 문자열 생성
             binary_data = ''
             for i in range(image.shape[0]):
                 for j in range(image.shape[1]):
                     for k in range(3):
                         binary_data += str(image[i, j, k] & 1)
 
-            # Convert binary to text
+            # 먼저 헤더(32비트)를 읽어 데이터 길이(문자수)를 결정
+            header = binary_data[:32]
+            data_length = int(header, 2)
+            total_bits = data_length * 8
+            # 본문 데이터 읽기
+            message_bits = binary_data[32:32 + total_bits]
             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"
-
+            for i in range(0, len(message_bits), 8):
+                byte = message_bits[i:i+8]
+                text += chr(int(byte, 2))
+            try:
+                data = json.loads(text)
+                return json.dumps(data, ensure_ascii=False, indent=2)
+            except json.JSONDecodeError:
+                return text
         except Exception as e:
             return f"Error in decoding: {str(e)}"