Update utils.py

utils.py

@@ -116,7 +116,8 @@ class WatermarkProcessor:
           except Exception as e:
                return image_path, f"Error in encoding: {str(e)}"
 
-     def decode(self, image_path):
+
+    def decode(self, image_path):
           """Decode watermark using simple LSB steganography"""
           try:
                # Try PNG metadata method first
@@ -127,72 +128,61 @@ class WatermarkProcessor:
                except:
                     pass
 
+               # Read image
                image = cv2.imread(image_path)
                if image is None:
                     raise ValueError("Could not read image file")
 
-               # Extract binary string
-               binary_data = ''
+               # Extract LSB data
+               binary_text = ''
+               # Get LSB from each byte of each pixel
                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 string
-               decoded_chars = []
-               for i in range(0, len(binary_data), 8):
-                    byte = binary_data[i:i+8]
-                    if len(byte) == 8:
-                         decoded_chars.append(chr(int(byte, 2)))
-                         # Check for delimiter
-                         if ''.join(decoded_chars[-7:]) == "###END###":
-                              decoded_text = ''.join(decoded_chars[:-7])
-                              try:
-                                   # Parse JSON data
-                                   watermark_data = json.loads(decoded_text)
-                                   return json.dumps(watermark_data, ensure_ascii=False, indent=2)
-                              except json.JSONDecodeError:
-                                   return decoded_text
-               
-               return "Error: No valid watermark found"
+                              bit = image[i, j, k] & 1
+                              binary_text += str(bit)
+
+               # Convert binary to text
+               text = ''
+               # Process 8 bits at a time
+               for i in range(0, len(binary_text)-8, 8):
+                    byte = binary_text[i:i+8]
+                    text += chr(int(byte, 2))
+                    # Check for end marker
+                    if text[-7:] == "###END###":
+                         extracted_text = text[:-7]
+                         try:
+                              # Attempt to parse as JSON
+                              data = json.loads(extracted_text)
+                              return json.dumps(data, ensure_ascii=False, indent=2)
+                         except:
+                              # If not valid JSON, return as plain text
+                              return extracted_text.strip()
+
+               # If we got here, try to find any valid JSON in the text
+               try:
+                    # Convert remaining bits
+                    for i in range(0, len(binary_text), 8):
+                         if i + 8 <= len(binary_text):
+                              byte = binary_text[i:i+8]
+                              text += chr(int(byte, 2))
+                    
+                    # Look for JSON markers
+                    start = text.find('{')
+                    end = text.rfind('}')
+                    if start != -1 and end != -1:
+                         json_text = text[start:end+1]
+                         try:
+                              data = json.loads(json_text)
+                              return json.dumps(data, ensure_ascii=False, indent=2)
+                         except:
+                              pass
+
+               except:
+                    pass
+
+               return text if text.strip() else "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)}"