add: jpeg compression estimator

app.py

@@ -26,6 +26,7 @@ from forensics.minmax import minmax_process
 from forensics.ela import ELA
 from forensics.wavelet import noise_estimation
 from forensics.bitplane import bit_plane_extractor
+from forensics.jpeg_compression import estimate_qf
 from utils.hf_logger import log_inference_data
 from utils.load import load_image
 from agents.ensemble_team import EnsembleMonitorAgent, WeightOptimizationAgent, SystemHealthAgent

forensics/__init__.py

@@ -4,6 +4,7 @@ from .ela import ELA
 from .gradient import gradient_processing
 from .minmax import minmax_process
 from .wavelet import noise_estimation
+from .jpeg_compression import estimate_qf
 
 __all__ = [
     'bit_plane_extractor',
@@ -11,5 +12,6 @@ __all__ = [
      # 'exif_full_dump',
     'gradient_processing',
     'minmax_process',
-    'noise_estimation'
-] 
+    'noise_estimation',
+    'estimate_qf'
+]

forensics/jpeg_compression.py

@@ -0,0 +1,157 @@
+import cv2 as cv
+import numpy as np
+from PIL import Image
+
+DCT_SIZE = 8
+TABLE_SIZE = DCT_SIZE ** 2
+ZIG_ZAG = [
+    [0, 0],
+    [0, 1],
+    [1, 0],
+    [2, 0],
+    [1, 1],
+    [0, 2],
+    [0, 3],
+    [1, 2],
+    [2, 1],
+    [3, 0],
+    [4, 0],
+    [3, 1],
+    [2, 2],
+    [1, 3],
+    [0, 4],
+    [0, 5],
+    [1, 4],
+    [2, 3],
+    [3, 2],
+    [4, 1],
+    [5, 0],
+    [6, 0],
+    [5, 1],
+    [4, 2],
+    [3, 3],
+    [2, 4],
+    [1, 5],
+    [0, 6],
+    [0, 7],
+    [1, 6],
+    [2, 5],
+    [3, 4],
+    [4, 4],
+    [5, 3],
+    [6, 2],
+    [7, 1],
+    [7, 2],
+    [6, 3],
+    [5, 4],
+    [4, 5],
+    [3, 5],
+    [2, 6],
+    [1, 7],
+    [2, 7],
+    [3, 6],
+    [4, 5],
+    [5, 4],
+    [6, 3],
+    [7, 2],
+    [7, 3],
+    [6, 4],
+    [5, 5],
+    [4, 6],
+    [3, 7],
+    [4, 7],
+    [5, 6],
+    [6, 5],
+    [7, 4],
+    [7, 5],
+    [6, 6],
+    [5, 7],
+    [6, 7],
+    [7, 6],
+    [7, 7],
+]
+
+
+def compress_jpg(image: Image.Image, quality, color=True):
+    """Compress a PIL image to JPEG format with specified quality.
+    
+    Args:
+        image: Input PIL image (RGB format)
+        quality: JPEG compression quality (1-100)
+        color: Whether to preserve color (BGR format)
+    
+    Returns:
+        np.ndarray: Decompressed image in BGR or grayscale format
+    """
+    # Convert PIL image to OpenCV BGR format
+    img_np = np.array(image)
+    if color:
+        img_np = cv.cvtColor(img_np, cv.COLOR_RGB2BGR)
+    
+    _, buffer = cv.imencode(".jpg", img_np, [cv.IMWRITE_JPEG_QUALITY, quality])
+    return cv.imdecode(buffer, cv.IMREAD_COLOR if color else cv.IMREAD_GRAYSCALE)
+
+
+def loss_curve(image: Image.Image, qualities=tuple(range(1, 101)), normalize=True):
+    """Calculate JPEG compression loss curve for quality estimation.
+    
+    Args:
+        image: Input PIL image (RGB format)
+        qualities: Quality values to test (1-100)
+        normalize: Whether to normalize the output curve
+    
+    Returns:
+        np.ndarray: Mean absolute difference values across quality levels
+    """
+    # Convert input image to grayscale BGR for compression testing
+    img_np = np.array(image)
+    if len(img_np.shape) == 3:
+        x = cv.cvtColor(img_np, cv.COLOR_RGB2GRAY)
+    else:
+        x = img_np
+    
+    c = np.array(
+        [cv.mean(cv.absdiff(compress_jpg(x, q, False), x))[0] for q in qualities]
+    )
+    if normalize:
+        c = cv.normalize(c, None, 0, 1, cv.NORM_MINMAX).flatten()
+    return c
+
+
+def estimate_qf(image):
+    return np.argmin(loss_curve(image))
+
+
+def get_tables(quality):
+    luma = np.array(
+        [
+            [16, 11, 10, 16, 24, 40, 51, 61],
+            [12, 12, 14, 19, 26, 58, 60, 55],
+            [14, 13, 16, 24, 40, 57, 69, 56],
+            [14, 17, 22, 29, 51, 87, 80, 62],
+            [18, 22, 37, 56, 68, 109, 103, 77],
+            [24, 35, 55, 64, 81, 104, 113, 92],
+            [49, 64, 78, 87, 103, 121, 120, 101],
+            [72, 92, 95, 98, 112, 100, 103, 99],
+        ]
+    )
+    chroma = np.array(
+        [
+            [17, 18, 24, 47, 99, 99, 99, 99],
+            [18, 21, 26, 66, 99, 99, 99, 99],
+            [24, 26, 56, 99, 99, 99, 99, 99],
+            [47, 66, 99, 99, 99, 99, 99, 99],
+            [99, 99, 99, 99, 99, 99, 99, 99],
+            [99, 99, 99, 99, 99, 99, 99, 99],
+            [99, 99, 99, 99, 99, 99, 99, 99],
+            [99, 99, 99, 99, 99, 99, 99, 99],
+        ]
+    )
+    quality = np.clip(quality, 1, 100)
+    if quality < 50:
+        quality = 5000 / quality
+    else:
+        quality = 200 - quality * 2
+    tables = np.concatenate((luma[:, :, np.newaxis], chroma[:, :, np.newaxis]), axis=2)
+    tables = (tables * quality + 50) / 100
+    return np.clip(tables, 1, 255).astype(int)