forensics/ela_hybrid.py

# 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):
    """Return a list of 2 images: [RGB Image, ELA Map (grayscale)], as PIL Images."""
    from PIL import Image

    # Extract RGB (3 channels at front)
    rgb_image = Image.fromarray((hybrid_array[:, :, :3] * 255).astype(np.uint8))
    
    # Extract ELA channels (last 3)
    ela_image = Image.fromarray((hybrid_array[:, :, 3:] * 255).astype(np.uint8))

    return [rgb_image, ela_image]

# Function to generate hybrid ELA as list of PIL Images for run_forensic_task
def generate_hybrid_ela_func(img):
    # Convert to numpy array if needed (e.g., from OpenCV BGR to PIL RGB)
    img_array = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return visualize_hybrid(generate_ela_hybrid(img_array))