Update utils.py

utils.py

@@ -1,211 +1,153 @@
-# utils.py
-
 import cv2
 import numpy as np
-from PIL import Image, PngImagePlugin, ImageDraw
-import json
-from datetime import datetime
-from cryptography.fernet import Fernet
-import base64
-import hashlib
+from PIL import Image, PngImagePlugin
 
-class WatermarkProcessor:
-    def __init__(self, encryption_key=None):
-        """Initialize with optional encryption key"""
-        if encryption_key:
-            self.fernet = Fernet(encryption_key)
-        else:
-            key = Fernet.generate_key()
-            self.fernet = Fernet(key)
+def png_encode(im_name, extra):
+    """Encode watermark using PNG metadata"""
+    try:
+        im = Image.open(im_name)
+        info = PngImagePlugin.PngInfo()
+        info.add_text("TXT", extra)
+        im.save("test.png", pnginfo=info)
+        return "test.png", "Watermark added successfully"
+    except Exception as e:
+        return im_name, f"Error adding watermark: {str(e)}"
 
-    def to_bin(self, data):
-        """Convert data to binary format as string"""
-        if isinstance(data, str):
-            return ''.join(format(x, '08b') for x in data.encode('utf-8'))
-        elif isinstance(data, bytes):
-            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):
-            return format(data, "08b")
-        else:
-            raise TypeError("Type not supported.")
+def to_bin(data):
+    """Convert data to binary format as string"""
+    if isinstance(data, str):
+        return ''.join(format(x, '08b') for x in data.encode('utf-8'))
+    elif isinstance(data, bytes):
+        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):
+        return format(data, "08b")
+    else:
+        raise TypeError("Type not supported.")
 
-    def create_preview(self, image_path, watermark_text, opacity=0.3):
-        """Create a preview of watermark on image"""
+def decode(image_name, txt=None):
+    """Decode watermark from image"""
+    try:
+        # First try PNG metadata method
         try:
-            image = Image.open(image_path)
-            txt_layer = Image.new('RGBA', image.size, (255, 255, 255, 0))
-            draw = ImageDraw.Draw(txt_layer)
-            
-            # Calculate text position
-            text_width = draw.textlength(watermark_text)
-            text_x = (image.width - text_width) // 2
-            text_y = image.height // 2
-            
-            # Add watermark text
-            draw.text((text_x, text_y), watermark_text, 
-                     fill=(255, 255, 255, int(255 * opacity)))
-            
-            # Combine layers
-            preview = Image.alpha_composite(image.convert('RGBA'), txt_layer)
-            return preview
-        except Exception as e:
-            return None
+            im = Image.open(image_name)
+            if "TXT" in im.info:
+                return im.info["TXT"]
+        except:
+            pass
 
-    def png_encode(self, im_name, extra):
-        """Encode watermark using PNG metadata"""
-        try:
-            im = Image.open(im_name)
-            info = PngImagePlugin.PngInfo()
-            info.add_text("TXT", extra)
-            im.save("test.png", pnginfo=info)
-            return "test.png", "Watermark added successfully"
-        except Exception as e:
-            return im_name, f"Error adding watermark: {str(e)}"
+        # Steganography method
+        image = cv2.imread(image_name)
+        if image is None:
+            raise ValueError("Could not read image file")
+            
+        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)
+                binary_data += r[-1]
+                binary_data += g[-1]
+                binary_data += b[-1]
 
-    def encode(self, image_path, watermark_text, metadata=None):
-        """Encode watermark using steganography with encryption"""
-        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 = {
-                '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)
-            
-            # 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"
-            
-            # Embed 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))
-            return output_path, "Watermark added successfully"
-            
-        except Exception as e:
-            return image_path, f"Error in encoding: {str(e)}"
+        # 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:
+                    byte_val = int(byte, 2)
+                    bytes_data += bytes([byte_val])
+                except ValueError:
+                    continue
 
-    def decode(self, image_path):
-        """Decode watermark with decryption"""
+        # Try to find the end marker in raw bytes
         try:
-            # Try PNG metadata method first
-            try:
-                im = Image.open(image_path)
-                if "TXT" in im.info:
-                    return im.info["TXT"]
-            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]
+            end_marker = b'====='
+            if end_marker in bytes_data:
+                bytes_data = bytes_data.split(end_marker)[0]
             
-            # Convert to bytes
-            bytes_data = bytearray()
-            for i in range(0, len(binary_data), 8):
-                if i + 8 <= len(binary_data):
-                    byte = int(binary_data[i:i+8], 2)
-                    bytes_data.append(byte)
+            # Try to find the delimiter in raw bytes
+            delimiter = b'#####'
+            if delimiter in bytes_data:
+                bytes_data = bytes_data.split(delimiter)[0]
             
-            # Process the data
-            decoded_data = bytes(bytes_data).decode('utf-8')
-            if "=====" in decoded_data:
-                decoded_data = decoded_data.split("=====")[0]
-            if "#####" in decoded_data:
-                decoded_data = decoded_data.split("#####")[0]
-                
-            try:
-                # Parse JSON data
-                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 json.dumps(watermark_data, indent=2, ensure_ascii=False)
-            except:
-                return decoded_data
+            # Decode the bytes to string
+            decoded_text = bytes_data.decode('utf-8', errors='ignore')
+            return decoded_text.strip()
             
         except Exception as e:
-            return f"Error in decoding: {str(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"
 
-    def analyze_quality(self, original_path, watermarked_path):
-        """Analyze watermark quality"""
-        try:
-            original = cv2.imread(original_path)
-            watermarked = cv2.imread(watermarked_path)
-            
-            if original is None or watermarked is None:
-                raise ValueError("Could not read image files")
-            
-            # Calculate PSNR
-            mse = np.mean((original - watermarked) ** 2)
-            if mse == 0:
-                psnr = float('inf')
-            else:
-                psnr = 20 * np.log10(255.0 / np.sqrt(mse))
-            
-            # Calculate histogram similarity
-            hist_original = cv2.calcHist([original], [0], None, [256], [0, 256])
-            hist_watermarked = cv2.calcHist([watermarked], [0], None, [256], [0, 256])
-            hist_correlation = cv2.compareHist(hist_original, hist_watermarked, cv2.HISTCMP_CORREL)
-            
-            # Count modified pixels
-            diff = cv2.bitwise_xor(original, watermarked)
-            modified_pixels = np.count_nonzero(diff)
-            
-            report = {
-                'psnr': round(psnr, 2),
-                'histogram_similarity': round(hist_correlation, 4),
-                'modified_pixels': modified_pixels,
-                'image_size': original.shape,
-                'quality_score': round((psnr / 50) * 100, 2) if psnr != float('inf') else 100
-            }
-            
-            return json.dumps(report, indent=2)
-            
-        except Exception as e:
-            return f"Error in quality analysis: {str(e)}"
-            
+    except Exception as e:
+        return f"Error detecting watermark: {str(e)}"
+
+def encode(image_name, secret_data, txt=None):
+    """Encode watermark using steganography"""
+    try:
+        image = cv2.imread(image_name)
+        if image is None:
+            raise ValueError("Could not read image file")
+            
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        
+        # Calculate maximum bytes that can be encoded
+        n_bytes = image.shape[0] * image.shape[1] * 3 // 8
+        
+        # Prepare the data with delimiters
+        secret_data = str(secret_data)
+        complete_data = secret_data + "#####" + "====="
+        
+        # Convert to binary
+        binary_secret_data = to_bin(complete_data)
+        
+        # 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"
+        
+        data_index = 0
+        binary_len = len(binary_secret_data)
+        
+        # Embed the data
+        for i in range(image.shape[0]):
+            for j in range(image.shape[1]):
+                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))
+        return output_path, "Watermark embedded successfully"
+        
+    except Exception as e:
+        return image_name, f"Error encoding watermark: {str(e)}"