Create ela_hybrid.py

forensics/ela_hybrid.py

@@ -0,0 +1,61 @@
+# ela_hybrid.py
+import numpy as np
+import cv2 as cv
+from PIL import Image
+import matplotlib.pyplot as plt
+
+
+def compress_jpg(image, quality=75):
+    """Compress image using JPEG compression (shared from ela.py)."""
+    encode_param = [int(cv.IMWRITE_JPEG_QUALITY), quality]
+    _, buffer = cv.imencode('.jpg', image, encode_param)
+    return cv.imdecode(buffer, cv.IMREAD_COLOR)
+
+
+def generate_ela_hybrid(image_path: str, quality: int = 95, scale_factor: int = 150):
+    """
+    Generate a 6-channel hybrid image combining RGB and ELA (3 channels each).
+    
+    Args:
+        image_path (str): Path to the input image.
+        quality (int): JPEG compression quality (1-100).
+        scale_factor (int): Scale factor for ELA contrast.
+        
+    Returns:
+        np.ndarray: 6-channel (RGB + ELA) image with shape (H, W, 6), dtype float32, normalized [0,1]
+    """
+    # Load original image
+    original = cv.imread(image_path, cv.IMREAD_COLOR)
+    original = original.astype(np.float32) / 255  # Normalize to [0, 1]
+
+    # Compress and reload image
+    compressed = compress_jpg(original, quality)
+    compressed = compressed.astype(np.float32) / 255  # Normalize to [0, 1]
+
+    # Generate ELA as absolute difference between original and compressed
+    ela = cv.absdiff(original, compressed)
+
+    # Apply scale factor to enhance ELA differences
+    ela = cv.convertScaleAbs(ela, alpha=scale_factor / 100, beta=0)
+    ela = ela.astype(np.float32) / 255  # Normalize back to [0, 1]
+
+    # Stack RGB and ELA (3 channels each) into 6-channel input
+    hybrid_image = np.concatenate([original, ela], axis=-1)  # Shape: H×W×6
+
+    return hybrid_image
+
+
+def save_hybrid_image(hybrid_array, save_path: str):
+    """Save the 6-channel hybrid image as a .npy file for model input."""
+    np.save(save_path, hybrid_array)
+    print(f"Saved hybrid ELA image to {save_path}")
+
+
+def visualize_hybrid(hybrid_array: np.ndarray, split_visualize: bool = True):
+    """Split the 6 channels into RGB and ELA for visualization."""
+    rgb_image = (hybrid_array[:, :, :3] * 255).astype(np.uint8)
+    ela_image = (hybrid_array[:, :, 3:] * 255).astype(np.uint8)
+    return {
+        "rgb": Image.fromarray(rgb_image),
+        "ela": Image.fromarray(ela_image),
+    }