Update utils.py

utils.py

@@ -7,8 +7,7 @@ def png_encode(im_name, extra):
     try:
         im = Image.open(im_name)
         info = PngImagePlugin.PngInfo()
-        # Encode the text as UTF-8 before adding to metadata
-        info.add_text("TXT", extra.encode('utf-8').decode('utf-8'))
+        info.add_text("TXT", extra)
         im.save("test.png", pnginfo=info)
         return "test.png", "Watermark added successfully"
     except Exception as e:
@@ -17,9 +16,10 @@ def png_encode(im_name, extra):
 def to_bin(data):
     """Convert data to binary format as string"""
     if isinstance(data, str):
-        return ''.join([format(ord(i), "08b") for i in data])
+        # UTF-8로 인코딩하여 바이트로 변환 후 바이너리 문자열로 변환
+        return ''.join(format(b, '08b') for b in data.encode('utf-8'))
     elif isinstance(data, bytes):
-        return ''.join([format(i, "08b") for i in data])
+        return ''.join(format(i, "08b") for i in data)
     elif isinstance(data, np.ndarray):
         return [format(i, "08b") for i in data]
     elif isinstance(data, int) or isinstance(data, np.uint8):
@@ -38,7 +38,7 @@ def decode(image_name, txt=None):
         except:
             pass
 
-        # If PNG metadata fails, try steganography method
+        # Steganography method
         image = cv2.imread(image_name)
         if image is None:
             raise ValueError("Could not read image file")
@@ -53,23 +53,27 @@ def decode(image_name, txt=None):
                 binary_data += g[-1]
                 binary_data += b[-1]
 
-        # Process binary data in chunks of 8 bits
-        decoded_data = ""
+        # 8비트씩 묶어서 바이트 배열 생성
+        bytes_array = bytearray()
         for i in range(0, len(binary_data), 8):
             byte = binary_data[i:i+8]
-            if len(byte) == 8:  # Ensure we have a complete byte
+            if len(byte) == 8:
                 try:
-                    decoded_data += chr(int(byte, 2))
+                    bytes_array.append(int(byte, 2))
                 except ValueError:
                     continue
 
-            if decoded_data.endswith("====="):
-                break
-
-        if "#####" in decoded_data:
-            decoded_data = decoded_data.split("#####")[0]
-            
-        return decoded_data.strip().rstrip("=")
+        # UTF-8로 디코딩
+        try:
+            decoded_data = bytes_array.decode('utf-8')
+            if "=====" in decoded_data:
+                decoded_data = decoded_data.split("=====")[0]
+            if "#####" in decoded_data:
+                decoded_data = decoded_data.split("#####")[0]
+            return decoded_data
+        except UnicodeDecodeError:
+            # UTF-8 디코딩에 실패한 경우
+            return "Error: Invalid character encoding detected"
 
     except Exception as e:
         return f"Error detecting watermark: {str(e)}"
@@ -86,17 +90,18 @@ def encode(image_name, secret_data, txt=None):
         # Calculate maximum bytes that can be encoded
         n_bytes = image.shape[0] * image.shape[1] * 3 // 8
         
-        # Prepare the secret data
-        secret_data = str(secret_data).encode('utf-8').decode('utf-8')  # Ensure proper encoding
-        secret_data_with_delimiter = secret_data + "#####"
+        # UTF-8로 인코딩하여 바이트 수 계산
+        secret_bytes = secret_data.encode('utf-8')
+        delimiter_bytes = "#####".encode('utf-8')
+        end_marker_bytes = "=====".encode('utf-8')
         
-        if len(secret_data_with_delimiter) + 5 >= n_bytes:
+        total_bytes = len(secret_bytes) + len(delimiter_bytes) + len(end_marker_bytes)
+        
+        if total_bytes >= n_bytes:
             return image_name, "Watermark is too large for Image Size"
             
-        secret_data_with_delimiter += "====="
-        
-        # Convert secret data to binary
-        binary_secret_data = to_bin(secret_data_with_delimiter)
+        # Prepare binary data
+        binary_secret_data = to_bin(secret_data) + to_bin("#####") + to_bin("=====")
         data_len = len(binary_secret_data)
         data_index = 0