Merge pull request #13 from hobostay/fix/multiple-bugs

Fix multiple bugs in extraction pipeline

src/extraction/synthid_bypass.py

@@ -17,11 +17,15 @@ Based on insights from:
 """
 
 import os
+import sys
 import io
 import numpy as np
 import cv2
 from scipy.fft import fft2, ifft2, fftshift, ifftshift
 from scipy import ndimage
+
+# Ensure same-directory modules are importable regardless of cwd
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 from typing import Optional, Dict, List, Tuple
 from dataclasses import dataclass
 from PIL import Image
@@ -118,12 +122,8 @@ class SynthIDBypass:
     ) -> np.ndarray:
         """Edge-preserving bilateral filter denoising."""
         img_uint8 = (image * 255).clip(0, 255).astype(np.uint8)
-        
+        denoised = cv2.bilateralFilter(img_uint8, d, sigma_color, sigma_space)
-        if len(image.shape) == 3:
+
-            denoised = cv2.bilateralFilter(img_uint8, d, sigma_color, sigma_space)
-        else:
-            denoised = cv2.bilateralFilter(img_uint8, d, sigma_color, sigma_space)
-        
         return denoised.astype(np.float32) / 255.0
     
     def denoise_nlm(
@@ -183,12 +183,14 @@ class SynthIDBypass:
         Similar to multiple KSampler passes in the diffusion workflow.
         """
         current = image.copy()
-        
+
         for i in range(passes):
-            # Decrease noise sigma slightly each pass
+            # Decrease noise sigma slightly each pass, clamp to avoid negative
-            sigma = noise_sigma * (1 - i * 0.2)
+            sigma = noise_sigma * max(0, 1 - i * 0.2)
+            if sigma <= 0:
+                break
             current = self.noise_replacement_pass(current, noise_sigma=sigma)
-        
+
         return current
     
     # ================================================================
@@ -1247,40 +1249,45 @@ class SynthIDBypass:
         current = img_f.copy()
         stages_applied = []
         s = params['base']
-        
+
-        # Single pass through transform categories
+        for iteration in range(iterations):
-        # Each attacks a different dimension of the watermark embedding
+            # Diminish strength slightly on subsequent iterations
-        
+            iter_s = s * max(0.5, 1.0 - iteration * 0.15)
-        # Stage 1: Spatial disruption — only in 'maximum' mode
+            iter_jpeg_q = min(95, params['jpeg_q'] + iteration * 5)
-        # (causes significant pixel misalignment affecting SSIM, but
+
-        #  targets SynthID's weakest category at 52% TPR worst-case)
+            # Pass through transform categories
-        if strength == 'maximum':
+            # Each attacks a different dimension of the watermark embedding
-            current = self._spatial_disruption(current, strength=s)
+
-            stages_applied.append('spatial')
+            # Stage 1: Spatial disruption — only in 'maximum' mode
-        
+            # (causes significant pixel misalignment affecting SSIM, but
-        # Stage 2: Quality degradation (JPEG/WebP/resize cycling)
+            #  targets SynthID's weakest category at 52% TPR worst-case)
-        current = self._quality_degradation(
+            if strength == 'maximum':
-            current, jpeg_quality=params['jpeg_q'], strength=s
+                current = self._spatial_disruption(current, strength=iter_s)
-        )
+                stages_applied.append(f'spatial_{iteration}')
-        stages_applied.append('quality')
+
-        
+            # Stage 2: Quality degradation (JPEG/WebP/resize cycling)
-        # Stage 3: Noise injection + denoising
+            current = self._quality_degradation(
-        current = self._noise_disruption(
+                current, jpeg_quality=iter_jpeg_q, strength=iter_s
-            current, sigma=params['noise_sigma'], strength=s
+            )
-        )
+            stages_applied.append(f'quality_{iteration}')
-        stages_applied.append('noise')
+
-        
+            # Stage 3: Noise injection + denoising
-        # Stage 4: Color manipulation
+            current = self._noise_disruption(
-        current = self._color_disruption(current, strength=s)
+                current, sigma=params['noise_sigma'], strength=iter_s
-        stages_applied.append('color')
+            )
-        
+            stages_applied.append(f'noise_{iteration}')
-        # Stage 5: Overlay disruption
+
-        current = self._overlay_disruption(current, strength=s)
+            # Stage 4: Color manipulation
-        stages_applied.append('overlay')
+            current = self._color_disruption(current, strength=iter_s)
-        
+            stages_applied.append(f'color_{iteration}')
+
+            # Stage 5: Overlay disruption
+            current = self._overlay_disruption(current, strength=iter_s)
+            stages_applied.append(f'overlay_{iteration}')
+
         # Clamp output to valid [0,1] range
         current = np.clip(current, 0, 1)
-        
+
         # Quantize to uint8 for consistent quality metrics
         cleaned_uint8 = (current * 255).clip(0, 255).astype(np.uint8)
         original_uint8 = (img_f * 255).clip(0, 255).astype(np.uint8)
@@ -1321,9 +1328,9 @@ class SynthIDBypass:
             }
         
         # Determine success
-        # Note: Internal SSIM computation is bugged on Python 3.14 (returns ~0)
+        # Rely on PSNR > 28 dB as primary quality gate; SSIM is computed
-        # while external computation is correct. We rely on PSNR > 28 dB which
+        # for reporting but heavy multi-pass transforms can depress it below
-        # strongly correlates with SSIM > 0.90 for these types of distortions.
+        # the 0.90 threshold even when visual quality is acceptable.
         success = psnr > 28
         if detection_before and detection_after:
             conf_drop = detection_before['confidence'] - detection_after['confidence']

src/extraction/watermark_remover.py

@@ -23,6 +23,9 @@ from scipy.ndimage import zoom
 from dataclasses import dataclass, field
 from typing import Optional, Dict, Tuple
 
+# Ensure same-directory modules are importable regardless of cwd
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
 
 @dataclass
 class RemovalResult: