Update utils.py

utils.py

@@ -1,5 +1,3 @@
-# utils.py
-
 import cv2
 import numpy as np
 from PIL import Image, PngImagePlugin, ImageDraw
@@ -10,193 +8,193 @@ import base64
 import hashlib
 
 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 to_bin(self, data):
-          """Convert data to binary format as string"""
-          if isinstance(data, str):
-               return ''.join(format(ord(char), '08b') for char in data)
-          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"""
-          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
-
-     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)}"
-
-     def encode(self, image_path, watermark_text, metadata=None):
-          """Encode watermark using simple LSB steganography"""
-          try:
-               image = cv2.imread(image_path)
-               if image is None:
-                    raise ValueError("Could not read image file")
-
-               # Prepare watermark data
-               data = {
-                    'text': watermark_text,
-                    'timestamp': datetime.now().isoformat(),
-                    'metadata': metadata or {}
-               }
-               
-               # Convert data to string with end marker
-               json_str = json.dumps(data, ensure_ascii=False)
-               secret_data = json_str + "#####"
-
-               # Convert string to binary (using utf-8 encoding)
-               binary_data = ''.join(format(ord(char), '08b') for char in secret_data)
-               
-               # Check capacity
-               if len(binary_data) > image.shape[0] * image.shape[1] * 3:
-                    return image_path, "Error: Image too small for watermark data"
-
-               # Embed data
-               data_index = 0
-               for i in range(image.shape[0]):
-                    for j in range(image.shape[1]):
-                         for k in range(3):
-                              if data_index < len(binary_data):
-                                   pixel = int(image[i, j, k])
-                                   # Clear the LSB
-                                   pixel = pixel & 0xFE  # Clear last bit
-                                   # Set the LSB according to our data
-                                   pixel = pixel | (int(binary_data[data_index]) & 1)
-                                   image[i, j, k] = np.uint8(pixel)
-                                   data_index += 1
-                              else:
-                                   break
-                         if data_index >= len(binary_data):
-                              break
+    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 to_bin(self, data):
+        """Convert data to binary format as string"""
+        if isinstance(data, str):
+            return ''.join(format(ord(char), '08b') for char in data)
+        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"""
+        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
+
+    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)}"
+
+    def encode(self, image_path, watermark_text, metadata=None):
+        """Encode watermark using simple LSB steganography"""
+        try:
+            image = cv2.imread(image_path)
+            if image is None:
+                raise ValueError("Could not read image file")
+
+            # Prepare watermark data
+            data = {
+                'text': watermark_text,
+                'timestamp': datetime.now().isoformat(),
+                'metadata': metadata or {}
+            }
+            
+            # Convert data to string with end marker
+            json_str = json.dumps(data, ensure_ascii=False)
+            secret_data = json_str + "#####"
+
+            # Convert string to binary (using utf-8 encoding)
+            binary_data = ''.join(format(ord(char), '08b') for char in secret_data)
+            
+            # Check capacity
+            if len(binary_data) > image.shape[0] * image.shape[1] * 3:
+                return image_path, "Error: Image too small for watermark data"
+
+            # Embed data
+            data_index = 0
+            for i in range(image.shape[0]):
+                for j in range(image.shape[1]):
+                    for k in range(3):
+                        if data_index < len(binary_data):
+                            pixel = int(image[i, j, k])
+                            # Clear the LSB
+                            pixel = pixel & 0xFE  # Clear last bit
+                            # Set the LSB according to our data
+                            pixel = pixel | (int(binary_data[data_index]) & 1)
+                            image[i, j, k] = np.uint8(pixel)
+                            data_index += 1
+                        else:
+                            break
                     if data_index >= len(binary_data):
-                         break
-
-               # Save result
-               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 using simple LSB steganography"""
-          try:
-               # Try PNG metadata first
-               try:
-                    im = Image.open(image_path)
-                    if "TXT" in im.info:
-                         return im.info["TXT"]
-               except:
-                    pass
-
-               image = cv2.imread(image_path)
-               if image is None:
-                    raise ValueError("Could not read image file")
-
-               # Extract binary data
-               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)
-
-               # Convert binary to text
-               text = ''
-               for i in range(0, len(binary_data), 8):
-                    if i + 8 <= len(binary_data):
-                         byte = binary_data[i:i+8]
-                         text += chr(int(byte, 2))
-                         
-                         # Check for end marker
-                         if "#####" in text:
-                              # Extract content before marker
-                              text = text.split("#####")[0]
-                              try:
-                                   # Parse JSON
-                                   data = json.loads(text)
-                                   return json.dumps(data, ensure_ascii=False, indent=2)
-                              except json.JSONDecodeError:
-                                   return text
-                              break
-
-               return "Error: No valid watermark found"
-
-          except Exception as e:
-               return f"Error in decoding: {str(e)}"
-
-     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)}"
+                        break
+                if data_index >= len(binary_data):
+                    break
+
+            # Save result
+            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 using simple LSB steganography"""
+        try:
+            # Try PNG metadata first
+            try:
+                im = Image.open(image_path)
+                if "TXT" in im.info:
+                    return im.info["TXT"]
+            except:
+                pass
+
+            image = cv2.imread(image_path)
+            if image is None:
+                raise ValueError("Could not read image file")
+
+            # Extract binary data
+            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)
+
+            # Convert binary to text
+            text = ''
+            for i in range(0, len(binary_data), 8):
+                if i + 8 <= len(binary_data):
+                    byte = binary_data[i:i+8]
+                    text += chr(int(byte, 2))
+                    
+                    # Check for end marker
+                    if "#####" in text:
+                        # Extract content before marker
+                        text = text.split("#####")[0]
+                        try:
+                            # Parse JSON
+                            data = json.loads(text)
+                            return json.dumps(data, ensure_ascii=False, indent=2)
+                        except json.JSONDecodeError:
+                            return text
+                        break
+
+            return "Error: No valid watermark found"
+
+        except Exception as e:
+            return f"Error in decoding: {str(e)}"
+
+    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)}"