reverse-synthid/analysis_engine.py at main · dearabhin/reverse-synthid · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
import matplotlib.pyplot as plt
import os
import cv2
import numpy as np
import pywt
from scipy.fft import fft2, fftshift
import matplotlib
matplotlib.use('Agg')  # Backend for server-side plotting

# The "Secret" Carrier Frequencies from the Research
CARRIERS = [
    (14, 14), (-14, -14), (126, 14), (-126, -14),
    (98, -14), (-98, 14), (128, 128), (-128, -128),
    (210, -14), (-210, 14), (238, 14), (-238, -14)
]


def wavelet_denoise(channel, wavelet='db4', level=3):
    """
    Exact denoising logic from the research paper.
    Separates the image structure from the high-frequency watermark.
    """
    coeffs = pywt.wavedec2(channel, wavelet, level=level)
    detail = coeffs[-1][0]
    sigma = np.median(np.abs(detail)) / 0.6745
    threshold = sigma * np.sqrt(2 * np.log(channel.size))

    new_coeffs = [coeffs[0]]
    for details in coeffs[1:]:
        new_details = tuple(pywt.threshold(
            d, threshold, mode='soft') for d in details)
        new_coeffs.append(new_details)

    denoised = pywt.waverec2(new_coeffs, wavelet)
    # Ensure size match (handling odd padding)
    return denoised[:channel.shape[0], :channel.shape[1]]


def process_image(image_path, output_folder):
    """
    Main pipeline: Load -> Denoise -> FFT -> Visualize
    """
    filename = os.path.basename(image_path)
    base_name = os.path.splitext(filename)[0]

    # 1. Load and Resize (Standardize to 512px)
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (512, 512))

    # 2. Extract Noise (The Watermark Layer)
    img_f = img.astype(np.float32) / 255.0
    noise = np.zeros_like(img_f)

    for c in range(3):
        denoised = wavelet_denoise(img_f[:, :, c])
        noise[:, :, c] = img_f[:, :, c] - denoised

    # 3. Generate Noise Visualization (Amplified)
    # Amplify 50x and shift to gray
    noise_vis = np.clip(noise * 50 + 0.5, 0, 1)
    noise_path = os.path.join(output_folder, f"{base_name}_noise.png")
    plt.imsave(noise_path, noise_vis)

    # 4. FFT Analysis (The "Star Pattern")
    # Convert noise to grayscale for FFT
    noise_gray = np.mean(noise, axis=2)
    f = fft2(noise_gray)
    fshift = fftshift(f)
    magnitude = np.abs(fshift)

    # Log scale for visibility
    log_mag = np.log1p(magnitude)

    # Normalize 0-1
    log_mag = (log_mag - np.min(log_mag)) / (np.max(log_mag) - np.min(log_mag))

    # 5. Plot FFT with Green Circles Overlay
    fig, ax = plt.subplots(figsize=(8, 8), dpi=100)
    ax.imshow(log_mag, cmap='inferno')
    ax.axis('off')

    # Draw the known carrier locations
    center = 256
    for dy, dx in CARRIERS:
        # SynthID coordinates are (y, x) offsets from center
        y, x = center + dy, center + dx
        circle = plt.Circle((x, y), 6, color='#00FF00',
                            fill=False, linewidth=1.5)
        ax.add_patch(circle)

    fft_path = os.path.join(output_folder, f"{base_name}_fft.png")
    plt.savefig(fft_path, bbox_inches='tight', pad_inches=0, facecolor='black')
    plt.close()

    return {
        'original': filename,
        'noise': f"{base_name}_noise.png",
        'fft': f"{base_name}_fft.png"
    }