Update utils.py

utils.py

@@ -16,10 +16,9 @@ def png_encode(im_name, extra):
 def to_bin(data):
     """Convert data to binary format as string"""
     if isinstance(data, str):
-        # UTF-8로 인코딩하여 바이트로 변환 후 바이너리 문자열로 변환
-        return ''.join(format(b, '08b') for b in data.encode('utf-8'))
+        return ''.join(format(x, '08b') for x in data.encode('utf-8'))
     elif isinstance(data, bytes):
-        return ''.join(format(i, "08b") for i in data)
+        return ''.join(format(x, '08b') for x in data)
     elif isinstance(data, np.ndarray):
         return [format(i, "08b") for i in data]
     elif isinstance(data, int) or isinstance(data, np.uint8):
@@ -46,6 +45,7 @@ def decode(image_name, txt=None):
         image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
         binary_data = ""
         
+        # Extract binary data from image
         for row in image:
             for pixel in row:
                 r, g, b = to_bin(pixel)
@@ -53,27 +53,50 @@ def decode(image_name, txt=None):
                 binary_data += g[-1]
                 binary_data += b[-1]
 
-        # 8비트씩 묶어서 바이트 배열 생성
-        bytes_array = bytearray()
+        # Convert binary string to bytes
+        bytes_data = b''
         for i in range(0, len(binary_data), 8):
             byte = binary_data[i:i+8]
             if len(byte) == 8:
                 try:
-                    bytes_array.append(int(byte, 2))
+                    byte_val = int(byte, 2)
+                    bytes_data += bytes([byte_val])
                 except ValueError:
                     continue
 
-        # UTF-8로 디코딩
+        # Try to find the end marker in raw bytes
         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"
+            end_marker = b'====='
+            if end_marker in bytes_data:
+                bytes_data = bytes_data.split(end_marker)[0]
+            
+            # Try to find the delimiter in raw bytes
+            delimiter = b'#####'
+            if delimiter in bytes_data:
+                bytes_data = bytes_data.split(delimiter)[0]
+            
+            # Decode the bytes to string
+            decoded_text = bytes_data.decode('utf-8', errors='ignore')
+            return decoded_text.strip()
+            
+        except Exception as e:
+            print(f"Decoding error: {e}")
+            # If regular decoding fails, try character by character
+            result = ""
+            current_bytes = bytearray()
+            
+            for byte in bytes_data:
+                current_bytes.append(byte)
+                try:
+                    char = current_bytes.decode('utf-8')
+                    result += char
+                    current_bytes = bytearray()
+                except UnicodeDecodeError:
+                    continue
+                    
+            if result:
+                return result.strip()
+            return "Error: Could not decode watermark"
 
     except Exception as e:
         return f"Error detecting watermark: {str(e)}"
@@ -90,31 +113,37 @@ 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
         
-        # UTF-8로 인코딩하여 바이트 수 계산
-        secret_bytes = secret_data.encode('utf-8')
-        delimiter_bytes = "#####".encode('utf-8')
-        end_marker_bytes = "=====".encode('utf-8')
+        # Prepare the data with delimiters
+        secret_data = str(secret_data)
+        complete_data = secret_data + "#####" + "====="
         
-        total_bytes = len(secret_bytes) + len(delimiter_bytes) + len(end_marker_bytes)
+        # Convert to binary
+        binary_secret_data = to_bin(complete_data)
         
-        if total_bytes >= n_bytes:
+        # Check if the data can fit in the image
+        if len(binary_secret_data) > n_bytes * 8:
             return image_name, "Watermark is too large for Image Size"
-            
-        # Prepare binary data
-        binary_secret_data = to_bin(secret_data) + to_bin("#####") + to_bin("=====")
-        data_len = len(binary_secret_data)
+        
         data_index = 0
+        binary_len = len(binary_secret_data)
         
-        # Embed the secret data
+        # Embed the data
         for i in range(image.shape[0]):
             for j in range(image.shape[1]):
-                if data_index < data_len:
-                    # Modify the least significant bit of each color channel
-                    for k in range(3):  # For each color channel (RGB)
-                        if data_index < data_len:
-                            image[i, j, k] = int(to_bin(image[i, j, k])[:-1] + binary_secret_data[data_index], 2)
+                if data_index < binary_len:
+                    pixel = image[i, j]
+                    for color_channel in range(3):
+                        if data_index < binary_len:
+                            # Get the binary value of the pixel
+                            binary_value = format(pixel[color_channel], '08b')
+                            # Replace the least significant bit
+                            binary_value = binary_value[:-1] + binary_secret_data[data_index]
+                            # Update the pixel value
+                            image[i, j, color_channel] = int(binary_value, 2)
                             data_index += 1
-                            
+                else:
+                    break
+                    
         # Save the result
         output_path = "watermarked_image.png"
         cv2.imwrite(output_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))