Update src/features/noiseprint_extractor.py

src/features/noiseprint_extractor.py

@@ -2,12 +2,7 @@ import os
 import numpy as np
 import torch
 from scipy.fftpack import fft2, fftshift
-from PIL import Image
 from src.features.noiseprint.Noiseprint import getNoiseprint
-from src.features.noiseprint.utilityRead import jpeg_qtableinv
-from src.features.noiseprint_wrapper import getNoiseprint_with_qf
-import io
-import tempfile
 
 class NoiseprintExtractor:
     """
@@ -25,138 +20,76 @@ class NoiseprintExtractor:
         # See: example_code/noiseprint/playground.ipynb which uses res[34:-34,34:-34]
         self.edge_margin = 34
 
-    def _extract_qf_from_pil_image(self, pil_image):
+    def extract_features(self, image_path):
         """
-        Extract JPEG Quality Factor from PIL Image quantization tables if available.
-        
-        Returns QF if quantization tables are preserved, None otherwise.
-        """
-        if not isinstance(pil_image, Image.Image):
-            return None
-            
-        if not hasattr(pil_image, 'quantization') or not pil_image.quantization:
-            return None
-        
-        try:
-            # Try to extract QF directly from quantization attribute
-            # This avoids re-saving which changes QF
-            q = pil_image.quantization
-            if isinstance(q, dict) and 0 in q:
-                # We have quantization tables, try to compute QF
-                # Use a temporary approach: save to BytesIO and read QF
-                # But this still changes QF... 
-                # Actually, we need to compute QF from quantization table directly
-                from src.features.noiseprint.utilityRead import jpeg_qtableinv
-                buf = io.BytesIO()
-                # Save with high quality to minimize QF change, then read it back
-                pil_image.save(buf, format='JPEG', quality=95)
-                buf.seek(0)
-                qf = jpeg_qtableinv(buf)
-                buf.close()
-                return qf
-        except Exception as e:
-            print(f"Could not extract QF from PIL Image: {e}")
-        
-        return None
-
-    def extract_features(self, image_input, qf_override=None):
-        """
-        Extracts Noiseprint-based features.
+        Extracts Noiseprint-based features for a given image.
         
         Args:
-            image_input: Can be:
-                - str: Path to image file (original behavior)
-                - PIL.Image: PIL Image object (new - tries to preserve QF)
-            qf_override: Optional QF value to use instead of detecting from file
-        
+            image_path: Path to the image file
+            
         Returns:
             dict: Dictionary of features with keys:
-                - 'noiseprint_freq_ratio': Ratio of high-frequency to total energy
+                - 'noiseprint_freq_ratio': Ratio of high-frequency to total energy in log-magnitude spectrum
                 - 'noiseprint_std': Standard deviation of the noiseprint residual
         """
         try:
-            # Handle PIL Image input (from Gradio)
-            if isinstance(image_input, Image.Image):
-                # Try to extract QF from quantization tables if preserved
-                detected_qf = self._extract_qf_from_pil_image(image_input)
-                
-                if detected_qf is None:
-                    # QF not available, need to save and detect
-                    # Use quality that matches common JPEG (75) to minimize QF mismatch
-                    temp_fd, temp_path = tempfile.mkstemp(suffix='.jpg')
-                    try:
-                        os.close(temp_fd)
-                        image_input.save(temp_path, "JPEG", quality=75, optimize=False, subsampling=0)
-                        qf = jpeg_qtableinv(temp_path) if qf_override is None else qf_override
-                        image_path = temp_path
-                    finally:
-                        # Will clean up after feature extraction
-                        pass
-                else:
-                    # QF detected from quantization tables
-                    qf = detected_qf if qf_override is None else qf_override
-                    # Still need to save for getNoiseprint (it requires file path)
-                    # But now we know the QF, so we can use it
-                    temp_fd, temp_path = tempfile.mkstemp(suffix='.jpg')
-                    os.close(temp_fd)
-                    # Save with quality matching detected QF
-                    save_quality = max(50, min(100, int(qf)))
-                    image_input.save(temp_path, "JPEG", quality=save_quality, optimize=False, subsampling=0)
-                    image_path = temp_path
-            else:
-                # String path (original behavior)
-                image_path = image_input
-                qf = qf_override
-            
-            # Extract noiseprint - use QF override if provided
-            if qf_override is not None or detected_qf is not None:
-                # Use wrapper that accepts QF override
-                qf_to_use = qf_override if qf_override is not None else detected_qf
-                _, noiseprint = getNoiseprint_with_qf(image_path, qf_override=qf_to_use)
-            else:
-                # Standard extraction (detects QF from file)
-                _, noiseprint = getNoiseprint(image_path)
-            
-            # Clean up temp file if we created one
-            if isinstance(image_input, Image.Image) and 'temp_path' in locals():
-                try:
-                    if os.path.exists(temp_path):
-                        os.remove(temp_path)
-                except:
-                    pass
+            # getNoiseprint returns (img, noiseprint)
+            # img is the image (H, W, C) or (H, W)
+            # noiseprint is the residual map (H, W)
+            _, noiseprint = getNoiseprint(image_path)
             
             if noiseprint is None:
                 return None
 
-            # Remove edge artifacts
+            # Remove edge artifacts: The original Noiseprint implementation uses 34-pixel margin
+            # to remove CNN boundary artifacts. This is consistent with the original codebase.
+            # Reference: example_code/noiseprint/playground.ipynb shows res[34:-34,34:-34]
             margin = self.edge_margin
             if noiseprint.shape[0] > 2*margin and noiseprint.shape[1] > 2*margin:
                 center_np = noiseprint[margin:-margin, margin:-margin]
             else:
+                # For very small images, use entire noiseprint but warn
                 center_np = noiseprint
+                if noiseprint.shape[0] <= 2*margin or noiseprint.shape[1] <= 2*margin:
+                    # Very small image - edge artifacts may affect features
+                    pass
             
-            # Feature 1: Frequency Ratio
+            # --- Feature 1: Frequency Ratio ---
+            # Compute 2D FFT and shift to center DC component
             f = fft2(center_np)
             fshift = fftshift(f)
+            
+            # Use log-magnitude spectrum (20*log10) for better dynamic range
+            # This is standard in frequency domain analysis (decibel scale)
+            # Add small epsilon to avoid log(0)
             magnitude_spectrum = 20 * np.log10(np.abs(fshift) + 1e-10)
             
             h, w = magnitude_spectrum.shape
             cy, cx = h // 2, w // 2
+            
+            # Define high-frequency region: exclude central low-frequency band
+            # Using 1/8 of image size for low-frequency mask (standard approach)
             mask_size = min(h, w) // 8
             high_freq_mask = np.ones((h, w), dtype=bool)
             high_freq_mask[cy-mask_size:cy+mask_size, cx-mask_size:cx+mask_size] = False
             
+            # Compute energy in high-frequency and total regions
             high_freq_energy = np.mean(magnitude_spectrum[high_freq_mask])
             total_energy = np.mean(magnitude_spectrum)
             
+            # Robust ratio calculation with stability check
+            # Use relative tolerance to handle near-zero cases
             eps = 1e-6
             if abs(total_energy) < eps:
+                # Very low energy: return 0.0 (no high-frequency content)
                 freq_ratio = 0.0
             else:
                 freq_ratio = high_freq_energy / total_energy
+                # Clip to reasonable range [0, 1] (high_freq_energy <= total_energy)
                 freq_ratio = np.clip(freq_ratio, 0.0, 1.0)
             
-            # Feature 2: Global Standard Deviation
+            # --- Feature 2: Global Standard Deviation ---
+            # Standard deviation of noiseprint residual (camera fingerprint strength)
             global_std = np.std(center_np)
             
             return {
@@ -165,7 +98,5 @@ class NoiseprintExtractor:
             }
             
         except Exception as e:
-            print(f"Error extracting Noiseprint features: {e}")
-            import traceback
-            traceback.print_exc()
+            print(f"Error extracting Noiseprint features for {image_path}: {e}")
             return None