frozencherry/Forgery-Localization-App

1import numpy as np2from PIL import Image, ImageChops, ImageEnhance, UnidentifiedImageError3import tensorflow as tf4import gradio as gr5import os6import traceback7import io8import cv2 # Added missing import for cv29 10# Enable XLA and set libdevice path programmatically11tf.config.optimizer.set_jit(True)12tf.config.optimizer.set_experimental_options({13    'xla_auto_jit': True,14    'xla_gpu_cuda_data_dir': '/home/ubuntu/miniconda3/pkgs/cuda-nvvm-tools-12.5.40-h59595ed_0/nvvm/libdevice'15})16 17# ==== CUSTOM OBJECTS ====18from custom_layer import CBAM  # Your custom attention module19from metrics import *  # Your loss/metric functions20 21# ==== LOAD MODEL ====22custom_objects = {23    'CBAM': CBAM,24    'dice_coef_loss': dice_coef_loss,25    'dice_coefficient': dice_coefficient,26    'iou': iou,27    'accuracy': accuracy,28    'weighted_dice_bce_loss': weighted_dice_bce_loss29}30 31model = tf.keras.models.load_model("final_finetune.keras", custom_objects=custom_objects)32input_height, input_width = 256, 25633print(model.input_shape)34print(model.output_shape)35 36ELA_QUALITY = 95 # Define ELA_QUALITY as it was missing37 38# ==== ELA CALCULATION ====39def calculate_ela(original_image, quality=ELA_QUALITY):40    with io.BytesIO() as output:41        original_image.save(output, format='JPEG', quality=quality)42        jpeg_data = output.getvalue()43    resaved_image = Image.open(io.BytesIO(jpeg_data))44    ela_image = ImageChops.difference(original_image, resaved_image)45    ela_image = ImageEnhance.Brightness(ela_image).enhance(6.0) # Changed to 6.046    return np.array(ela_image) # Return as NumPy array47 48# ==== PREPROCESS FUNCTION ====49def preprocess(image):50    try:51        if image is None:52            raise ValueError("No image provided.")53        54        if isinstance(image, tuple):55            image = image[0]  # Gradio might send (PIL.Image, dict) tuple56 57        if not isinstance(image, Image.Image):58            image = Image.open(image)59        60        # Convert image to RGB and then to JPEG in memory for consistent processing61        if image.mode != 'RGB':62            image = image.convert('RGB')63        with io.BytesIO() as output_jpeg:64            image.save(output_jpeg, format='JPEG')65            jpeg_data = output_jpeg.getvalue()66        image = Image.open(io.BytesIO(jpeg_data)) # Re-open as JPEG for consistent handling67 68        original_image = image.resize((input_width, input_height))69        ela_array = calculate_ela(original_image) # ela_array is now a NumPy array directly70 71        # Ensure original_image is also a numpy array for consistent handling72        original_array = np.array(original_image).astype(np.float32) / 255.073        ela_array = cv2.resize(ela_array, (input_width, input_height))74        ela_array = ela_array.astype(np.float32) / 255.075 76        # Ensure ela_array has 3 channels if it's grayscale (as per model.input_shape)77        if ela_array.ndim == 2:78            ela_array = np.stack([ela_array, ela_array, ela_array], axis=-1)79        elif ela_array.shape[-1] == 1:80            ela_array = np.concatenate([ela_array, ela_array, ela_array], axis=-1)81        82        return np.expand_dims(original_array, axis=0), np.expand_dims(ela_array, axis=0)83    84    except (UnidentifiedImageError, ValueError, TypeError, AttributeError) as e:85        print(f"Preprocessing error: {str(e)}")86        return None, None # Changed to return None, None for prediction to handle87 88# ==== PREDICT FUNCTION ====89def predict(image):90    original_preprocessed, ela_preprocessed = preprocess(image)91    if original_preprocessed is None or ela_preprocessed is None:92        return "Preprocessing failed", None93    try:94        prediction = model.predict([original_preprocessed, ela_preprocessed])[0]95        prediction_mask = (prediction > 0.5).astype(np.uint8) * 25596        prediction_image = Image.fromarray(prediction_mask.squeeze()).convert("L")97        return "Prediction successful", prediction_image98    except Exception as e:99        print(f"Prediction error: {str(e)}")100        return f"Error during prediction: {str(e)}", None101 102# ==== GRADIO APP ====103# Get sample image paths104sample_image_paths = [["samples/IMG_0002364.jpg"]]105 106demo = gr.Interface(107    fn=predict,108    inputs=gr.Image(type="pil", label="Upload Image"),109    outputs=[gr.Text(label="Status"), gr.Image(label="Forgery Mask")],110    title="Forgery Localization",111    description="Upload any image format (e.g., JPG, PNG, TIF) to detect forged regions using ELA + CNN.",112    examples=sample_image_paths # Add sample images here113)114 115if __name__ == "__main__":116    demo.launch()117