Update utils.py

utils.py

@@ -64,61 +64,60 @@ 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 steganography with encryption"""
+        """Encode watermark using steganography"""
         try:
+            # 이미지 읽기
             image = cv2.imread(image_path)
             if image is None:
                 raise ValueError("Could not read image file")
-                
-            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
             
-            # Prepare watermark data
-            watermark_data = {
+            # 워터마크 데이터 준비
+            data = {
                 'text': watermark_text,
                 'timestamp': datetime.now().isoformat(),
                 'metadata': metadata or {}
             }
             
-            # Add image hash
-            image_copy = image.copy() & 0xFE  # Clear LSB
-            watermark_data['image_hash'] = hashlib.sha256(image_copy.tobytes()).hexdigest()
-            
-            # Encrypt data
-            secret_data = json.dumps(watermark_data)
-            secret_data = f"{secret_data}#####=====" # Add delimiters
-            binary_secret_data = self.to_bin(secret_data)
+            # JSON으로 변환
+            json_data = json.dumps(data, ensure_ascii=False)
+            secret_data = json_data + "#####"  # 구분자 추가
             
-            # Calculate capacity
-            n_bytes = image.shape[0] * image.shape[1] * 3 // 8
-            if len(binary_secret_data) > n_bytes * 8:
-                return image_path, "Watermark is too large for Image Size"
+            # 바이트로 변환
+            byte_data = secret_data.encode('utf-8')
             
-            # Embed data
+            # 이미지 용량 확인
+            max_bytes = (image.shape[0] * image.shape[1] * 3) // 8
+            if len(byte_data) > max_bytes:
+                raise ValueError("Watermark data is too large for this image")
+
+            # 바이너리 데이터로 변환
+            binary_data = ''.join(format(byte, '08b') for byte in byte_data)
             data_index = 0
+
+            # 데이터 임베딩
             for i in range(image.shape[0]):
                 for j in range(image.shape[1]):
-                    if data_index < len(binary_secret_data):
-                        pixel = image[i, j]
-                        for k in range(3):
-                            if data_index < len(binary_secret_data):
-                                binary_value = format(pixel[k], '08b')
-                                binary_value = binary_value[:-1] + binary_secret_data[data_index]
-                                image[i, j, k] = int(binary_value, 2)
-                                data_index += 1
-            
-            # Save result
-            output_path = "watermarked_" + datetime.now().strftime("%Y%m%d_%H%M%S") + ".png"
-            cv2.imwrite(output_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
+                    for k in range(3):  # RGB channels
+                        if data_index < len(binary_data):
+                            # LSB 수정
+                            image[i, j, k] = (image[i, j, k] & 0xFE) | int(binary_data[data_index])
+                            data_index += 1
+
+            # 결과 저장
+            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"
-            
+
         except Exception as e:
             return image_path, f"Error in encoding: {str(e)}"
 
     def decode(self, image_path):
-        """Decode watermark with simplified robust error handling"""
+        """Decode watermark from image"""
         try:
-            # Try PNG metadata method first
+            # PNG metadata method first
             try:
                 im = Image.open(image_path)
                 if "TXT" in im.info:
@@ -126,78 +125,42 @@ class WatermarkProcessor:
             except:
                 pass
 
-            # Steganography method
+            # 이미지 읽기
             image = cv2.imread(image_path)
             if image is None:
                 raise ValueError("Could not read image file")
-                
-            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-            
-            # Extract binary data
-            binary_data = ""
-            for row in image:
-                for pixel in row:
-                    for value in pixel:
-                        binary_data += format(value, '08b')[-1]
-
-            # Convert binary string to bytes
+
+            # 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)
+
+            # 8비트씩 잘라서 ��이트로 변환
             bytes_data = bytearray()
             for i in range(0, len(binary_data), 8):
                 if i + 8 <= len(binary_data):
-                    try:
-                        byte = int(binary_data[i:i+8], 2)
-                        bytes_data.append(byte)
-                    except ValueError:
-                        continue
-
-            # Find end marker in raw bytes
-            end_sequence = b'====='
-            if end_sequence in bytes_data:
-                bytes_data = bytes_data[:bytes_data.index(end_sequence)]
-
-            # Find delimiter in raw bytes
-            delimiter = b'#####'
-            if delimiter in bytes_data:
-                bytes_data = bytes_data[:bytes_data.index(delimiter)]
+                    byte = int(binary_data[i:i+8], 2)
+                    bytes_data.append(byte)
+                    
+                    # 구분자 체크 (#####)
+                    if bytes_data[-5:] == b'#####':
+                        bytes_data = bytes_data[:-5]  # 구분자 제거
+                        break
 
             try:
-                # Try direct UTF-8 decoding first
+                # UTF-8로 디코딩
                 decoded_data = bytes_data.decode('utf-8')
                 
-                try:
-                    # Try to parse as JSON
-                    watermark_data = json.loads(decoded_data)
-                    # Verify image hash
-                    image_copy = image.copy() & 0xFE
-                    current_hash = hashlib.sha256(image_copy.tobytes()).hexdigest()
-                    
-                    if current_hash != watermark_data.get('image_hash'):
-                        return "Warning: Image has been modified after watermarking"
-                    
-                    # Return properly formatted JSON
-                    return json.dumps(watermark_data, indent=2, ensure_ascii=False)
-                except json.JSONDecodeError:
-                    # If not JSON, return decoded text
-                    return decoded_data.strip()
-                    
+                # JSON 파싱
+                watermark_data = json.loads(decoded_data)
+                return json.dumps(watermark_data, ensure_ascii=False, indent=2)
             except UnicodeDecodeError:
-                # If UTF-8 decoding fails, try character by character
-                result = []
-                i = 0
-                while i < len(bytes_data):
-                    for j in range(min(4, len(bytes_data) - i), 0, -1):
-                        try:
-                            char = bytes_data[i:i+j].decode('utf-8')
-                            result.append(char)
-                            i += j
-                            break
-                        except UnicodeDecodeError:
-                            if j == 1:  # If even single byte fails
-                                i += 1
-                            continue
-                
-                return ''.join(result).strip()
-            
+                return "Error: Invalid character encoding"
+            except json.JSONDecodeError:
+                return decoded_data
+
         except Exception as e:
             return f"Error in decoding: {str(e)}"