Update utils.py

utils.py

@@ -63,9 +63,9 @@ 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 with header
+        """Encode watermark using simple LSB steganography with header.
         
-        헤더(32비트)는 watermark 데이터(UTF-8 문자열)의 길이(문자수)를 이진 문자열로 저장합니다.
+        헤더(32비트)는 watermark 데이터(JSON 문자열)의 길이(문자수)를 이진 문자열로 저장합니다.
         """
         try:
             image = cv2.imread(image_path)
@@ -121,9 +121,9 @@ class WatermarkProcessor:
             return image_path, f"Error in encoding: {str(e)}"
 
     def decode(self, image_path):
-        """Decode watermark using simple LSB steganography with header
+        """Decode watermark using simple LSB steganography with header.
         
-        먼저 32비트를 읽어 watermark 데이터의 길이(문자수)를 구한 후, 해당 길이만큼의 데이터를 추출합니다.
+        먼저 32비트(헤더)만큼 읽어 데이터 길이(문자수)를 구한 후, 그 길이에 해당하는 본문 비트만 읽어 문자열로 복원합니다.
         """
         try:
             # PNG 메타데이터 먼저 확인
@@ -138,18 +138,33 @@ class WatermarkProcessor:
             if image is None:
                 raise ValueError("Could not read image file")
 
-            # 모든 픽셀의 LSB를 읽어 이진 문자열 생성
-            binary_data = ''
+            bits = []
+            total_needed = None
+            count = 0
+            # 픽셀 순회: 필요한 비트 수만큼 읽어오기
             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)
+                        bits.append(str(image[i, j, k] & 1))
+                        count += 1
+                        # 헤더 32비트를 모두 읽은 경우
+                        if count == 32 and total_needed is None:
+                            header = ''.join(bits[:32])
+                            data_length = int(header, 2)
+                            total_needed = 32 + data_length * 8
+                        if total_needed is not None and count >= total_needed:
+                            break
+                    if total_needed is not None and count >= total_needed:
+                        break
+                if total_needed is not None and count >= total_needed:
+                    break
 
-            # 먼저 헤더(32비트)를 읽어 데이터 길이(문자수)를 결정
+            if total_needed is None:
+                return "Error: Not enough data to read header"
+            binary_data = ''.join(bits[:total_needed])
             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(message_bits), 8):