PRNU.py

import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
import imageio
import tempfile
import os

# --- Import your custom source files ---
# Make sure the 'src' folder is in the same directory as this notebook
import utils.src.Functions as Fu
import utils.src.Filter as Ft
import utils.src.maindir as md

# --- App Description ---
# A detailed description using Markdown. It explains what the tool does and how to use it.
# You can add images from the web using standard markdown syntax.
description = """

This tool analyzes an image to detect potential manipulations using Photo-Response Non-Uniformity (PRNU), a unique noise pattern that acts as a camera's fingerprint.

## How it Works
1.  **Camera Fingerprint**: Every digital camera sensor has a unique, systematic noise pattern called PRNU. We use a pre-extracted fingerprint file (`.dat`) for a specific camera.
2.  **Image Analysis**: The tool extracts the noise residual from the uploaded image.
3.  **Correlation**: It then compares the image's noise with the camera's fingerprint by calculating the Peak-to-Correlation Energy (PCE) across different blocks of the image.
4.  **PCE Map**: The output PCE map visualizes this correlation. High PCE values (brighter areas) suggest that this part of the image was likely taken by the fingerprinted camera. Dark or inconsistent areas could indicate tampering or that the image was taken with a different device.

![Diagram of the PRNU process](https://www.researchgate.net/profile/Feyisetan-Ojerinde/publication/323982885/figure/fig1/AS:607997380104192@1521971778119/Block-diagram-of-PRNU-based-source-camera-identification-process.png)

## Instructions:
1.  Upload the camera's fingerprint file (`.dat` format).
2.  Upload the JPG/PNG image you want to analyze.
3.  Click **Submit** and view the resulting PCE map.
"""

# --- Main Analysis Function ---
# This function contains all the logic from your script.
# It takes a fingerprint file and an image array as input, and returns two plots.
def analyze_image_forgery(fingerprint_file, input_image):
    """
    Processes an image against a camera fingerprint to generate a PCE map.
    
    Args:
        fingerprint_file (gradio.File): The uploaded camera fingerprint .dat file.
        input_image (np.array): The uploaded image as a NumPy array.
        
    Returns:
        (matplotlib.figure, matplotlib.figure): A tuple containing the two output plots.
    """
    # --- 1. Load Camera Fingerprint ---
    print("Loading camera fingerprint...")
    Fingerprint = np.genfromtxt(fingerprint_file.name)
    print(f"Fingerprint loaded. Shape: {Fingerprint.shape}")

    # --- 2. Save uploaded image to a temporary file ---
    # The NoiseExtractFromImage function expects a file path, so we create one.
    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_img_file:
        temp_img_path = temp_img_file.name
        imageio.imwrite(temp_img_path, input_image)

    # --- 3. Extract and filter PRNU noise from the image ---
    print("Extracting noise from image...")
    Noisex = Ft.NoiseExtractFromImage(temp_img_path, sigma=2.)
    Noisex = Fu.WienerInDFT(Noisex, np.std(Noisex))
    print(f"Noise extracted. Shape: {Noisex.shape}")

    # Clean up the temporary image file
    os.remove(temp_img_path)

    # --- 4. Align Fingerprint and PRNU sizes by padding if necessary ---
    if Noisex.shape != Fingerprint.shape:
        print("Shapes do not match. Padding PRNU noise to match fingerprint size.")
        Noisex_padded = np.zeros_like(Fingerprint)
        h = min(Noisex.shape[0], Fingerprint.shape[0])
        w = min(Noisex.shape[1], Fingerprint.shape[1])
        Noisex_padded[:h, :w] = Noisex[:h, :w]
        Noisex = Noisex_padded
    
    # --- 5. Compute PCE Map in blocks ---
    print("Computing PCE map...")
    block_size = 64
    blocks_x = np.arange(0, Noisex.shape[0], block_size)
    blocks_y = np.arange(0, Noisex.shape[1], block_size)
    PCE_map = np.zeros((len(blocks_x), len(blocks_y)))

    for y_idx, y_start in enumerate(blocks_y):
        for x_idx, x_start in enumerate(blocks_x):
            block_Noisex = Noisex[x_start:x_start+block_size, y_start:y_start+block_size]
            block_Fingerprint = Fingerprint[x_start:x_start+block_size, y_start:y_start+block_size]
            
            # Skip if blocks are not of the expected size (can happen at edges)
            if block_Noisex.shape != (block_size, block_size):
                continue

            C = Fu.crosscorr(block_Noisex, block_Fingerprint)
            det, _ = md.PCE(C)
            PCE_map[x_idx, y_idx] = det.get('PCE', 0) # Use .get for safety

    print("PCE map computed successfully.")

    # --- 6. Generate Output Plots ---
    # Plot 1: PCE Map
    fig1, ax1 = plt.subplots(figsize=(8, 6))
    im = ax1.imshow(PCE_map, cmap='viridis')
    ax1.set_title('Detection PCE-map')
    fig1.colorbar(im, ax=ax1, label='PCE Value')
    
    # Plot 2: Original Image
    fig2, ax2 = plt.subplots(figsize=(8, 6))
    ax2.imshow(input_image)
    ax2.set_title('Analyzed Image')
    ax2.axis('off')

    return fig1, fig2

# --- Create and Launch the Gradio Interface ---
def build_demo():
    return gr.Interface(
        fn=analyze_image_forgery,
        inputs=[
            gr.File(label="Upload Camera Fingerprint (.dat file)"),
            gr.Image(type="numpy", label="Upload Image to Analyze")
        ],
        outputs=[
            gr.Plot(label="PCE Map"),
            gr.Plot(label="Analyzed Image")
        ],
        title="📸 PRNU-Based Image Forgery Detector",
        description=description
    )

if __name__ == "__main__":
    iface = build_demo()
    iface.launch()