utils/ela.py

import numpy as np
import cv2 as cv
from time import time
from PIL import Image

def compress_jpg(image, quality):
    """Compress image using JPEG compression."""
    encode_param = [int(cv.IMWRITE_JPEG_QUALITY), quality]
    _, buffer = cv.imencode('.jpg', image, encode_param)
    return cv.imdecode(buffer, cv.IMREAD_COLOR)

def desaturate(image):
    """Convert image to grayscale."""
    return cv.cvtColor(image, cv.COLOR_BGR2GRAY)

def create_lut(contrast, brightness):
    """Create lookup table for contrast and brightness adjustment."""
    lut = np.arange(256, dtype=np.uint8)
    lut = cv.LUT(lut, lut)
    lut = cv.convertScaleAbs(lut, None, contrast/128, brightness)
    return lut

def elapsed_time(start):
    """Calculate elapsed time since start."""
    return f"{time() - start:.3f}s"

def genELA(img, quality=75, scale=50, contrast=20, linear=False, grayscale=False):
    """
    Perform Error Level Analysis on an image.
    
    Args:
        img: Input image (numpy array)
        quality: JPEG compression quality (1-100)
        scale: Output multiplicative gain (1-100)
        contrast: Output tonality compression (0-100)
        linear: Whether to use linear difference
        grayscale: Whether to output grayscale image
    
    Returns:
        Processed ELA image
    """
    # Convert image to float32 and normalize
    original = img.astype(np.float32) / 255
    
    # Compress image
    compressed = compress_jpg(img, quality)
    compressed = compressed.astype(np.float32) / 255
    
    # Calculate difference based on mode
    if not linear:
        difference = cv.absdiff(original, compressed)
        ela = cv.convertScaleAbs(cv.sqrt(difference) * 255, None, scale / 20)
    else:
        ela = cv.convertScaleAbs(cv.subtract(compressed, img), None, scale)
    
    # Apply contrast adjustment
    contrast_value = int(contrast / 100 * 128)
    ela = cv.LUT(ela, create_lut(contrast_value, contrast_value))
    
    # Convert to grayscale if requested
    if grayscale:
        ela = desaturate(ela)
    
    return Image.fromarray(ela)