Update app.py

app.py

@@ -1,321 +1,350 @@
-import uuid
 import os
-import matplotlib
-matplotlib.use("Agg")
-
-import gradio as gr
+import uuid
+import hashlib
+import numpy as np
 import librosa
 import librosa.display
-import matplotlib.pyplot as plt
-import numpy as np
 import soundfile as sf
-import spaces
-
-# ==========================================================
-# AudioShield v2 — Robust PN Spread-Spectrum Watermark
-# ==========================================================
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import gradio as gr
 
+# ============================================================
+# AudioShield v3 — Keyed Robust Audio Watermark
+# ============================================================
+# Design goals:
+# - No fixed ultrasonic tone.
+# - Keyed pseudo-random spread-spectrum payload.
+# - Repeated blocks + synchronization marker.
+# - Detection compares the expected keyed sequence against
+#   differential spectral energy, reducing naive false positives.
+# - WAV output preserves the processing result.
+#
+# IMPORTANT:
+# This is a research prototype, not a claim of "indestructible"
+# watermarking. For state-of-the-art learned watermarking,
+# AudioSeal/WavMark-style trained models are stronger candidates.
+# ============================================================
+
+TARGET_SR = 16000
 N_FFT = 2048
-HOP_LENGTH = 512
-MIN_FREQ = 500.0
-MAX_GUARD_HZ = 1500.0
-
-def _generate_pn(shape, seed=42, chip_hold=4, smooth_taps=5):
-    n_freq, n_time = shape
-    rng = np.random.default_rng(int(seed))
-    n_chips = int(np.ceil(n_time / chip_hold)) + 1
-    chips = rng.choice([-1.0, 1.0], size=(n_freq, n_chips))
-    pattern = np.repeat(chips, chip_hold, axis=1)[:, :n_time]
-
-    if smooth_taps > 1:
-        kernel = np.bartlett(smooth_taps)
-        if np.sum(kernel) > 0:
-            kernel = kernel / np.sum(kernel)
-            pattern = np.apply_along_axis(
-                lambda row: np.convolve(row, kernel, mode="same"),
-                axis=1,
-                arr=pattern
-            )
-    return pattern
-
-def _band_mask(sr, n_fft):
-    freqs = librosa.fft_frequencies(sr=sr, n_fft=n_fft)
-    max_freq = sr / 2.0 - MAX_GUARD_HZ
-    return (freqs >= MIN_FREQ) & (freqs <= max_freq)
-
-def _local_mask(magnitude, sr):
-    # Energy-adaptive psychoacoustic mask.
-    # Robust floor prevents the watermark from disappearing in quiet bins.
-    band = _band_mask(sr, N_FFT)
-    energy = magnitude / (np.percentile(magnitude[band], 95) + 1e-8)
-    energy = np.clip(energy, 0.0, 1.0)
-
-    # Do not inject into very weak bins.
-    active = energy > 0.08
-    mask = np.sqrt(np.maximum(energy, 0.0)) * active
-    mask *= band[:, None]
-    return mask
-
-def _embed_channel(y, sr, key, alpha, masking=True):
-    stft = librosa.stft(
-        y, n_fft=N_FFT, hop_length=HOP_LENGTH,
-        win_length=N_FFT, window="hann"
-    )
-    magnitude = np.abs(stft)
-    phase = np.angle(stft)
-
-    pn = _generate_pn(magnitude.shape, seed=key)
-    mask = _local_mask(magnitude, sr) if masking else _band_mask(sr, N_FFT)[:, None]
-
-    # Differential, relative embedding:
-    # watermark amplitude follows local signal energy.
-    wm = alpha * (magnitude + 1e-8) * mask * pn
-    wm_mag = np.maximum(magnitude * (1.0 + wm), 1e-12)
+HOP = 512
+BLOCK_SECONDS = 2.0
+PAYLOAD_BITS = 32
+REPEATS_PER_BIT = 8
+ALPHA = 0.018
+LOW_HZ = 700.0
+HIGH_HZ = 7000.0
+DETECT_THRESHOLD = 0.22
+MIN_BLOCKS = 3
+
+MAG_EPS = 1e-8
+
+
+def _seed_from_key(key):
+    h = hashlib.sha256(str(int(key)).encode("utf-8")).digest()
+    return int.from_bytes(h[:8], "little", signed=False) % (2**32 - 1)
+
+
+def _payload_from_key(key):
+    """32-bit deterministic identifier derived from secret key."""
+    digest = hashlib.sha256(f"AudioShield-v3:{int(key)}".encode()).digest()
+    bits = np.unpackbits(np.frombuffer(digest[:4], dtype=np.uint8))
+    return bits.astype(np.int8)
+
+
+def _bits_to_symbols(bits):
+    return np.where(bits > 0, 1.0, -1.0)
+
+
+def _freq_bins(sr):
+    freqs = librosa.fft_frequencies(sr=sr, n_fft=N_FFT)
+    mask = (freqs >= LOW_HZ) & (freqs <= min(HIGH_HZ, sr / 2 - 300))
+    idx = np.where(mask)[0]
+    if len(idx) < 20:
+        raise ValueError("Échantillonnage trop faible pour la bande de watermark.")
+    return freqs, idx
+
+
+def _frame_strength(mag, freq_idx):
+    """Robust local normalization; removes absolute loudness dependence."""
+    x = mag[freq_idx, :]
+    med = np.median(x, axis=1, keepdims=True)
+    mad = np.median(np.abs(x - med), axis=1, keepdims=True) + MAG_EPS
+    z = (x - med) / (4.0 * mad)
+    return np.clip(z, -3.0, 3.0)
+
+
+def _make_keyed_pattern(n_freq, n_frames, key, bit_index, block_index):
+    seed = (
+        _seed_from_key(key)
+        ^ ((bit_index + 1) * 0x9E3779B1)
+        ^ ((block_index + 1) * 0x85EBCA77)
+    ) & 0xFFFFFFFF
+    rng = np.random.default_rng(seed)
+    # Zero-mean random chips across frequency and time.
+    p = rng.choice([-1.0, 1.0], size=(n_freq, n_frames))
+    # Temporal smoothing prevents a tonal line from appearing.
+    if n_frames >= 5:
+        k = np.array([1, 2, 3, 2, 1], dtype=np.float32)
+        k /= k.sum()
+        p = np.apply_along_axis(lambda r: np.convolve(r, k, mode="same"), 1, p)
+    p /= np.sqrt(np.mean(p * p) + MAG_EPS)
+    return p
+
+
+def _embed_mono(y, sr, key, alpha):
+    y16 = librosa.resample(y.astype(np.float32), orig_sr=sr, target_sr=TARGET_SR)
+    stft = librosa.stft(y16, n_fft=N_FFT, hop_length=HOP, win_length=N_FFT, window="hann")
+    mag, phase = np.abs(stft), np.angle(stft)
+
+    _, fidx = _freq_bins(TARGET_SR)
+    block_frames = int(BLOCK_SECONDS * TARGET_SR / HOP)
+    n_blocks = max(1, int(np.ceil(mag.shape[1] / block_frames)))
+    payload = _payload_from_key(key)
+
+    wm_mag = mag.copy()
+    used_blocks = 0
+
+    for b in range(n_blocks):
+        a = b * block_frames
+        z = min((b + 1) * block_frames, mag.shape[1])
+        if z - a < max(12, block_frames // 3):
+            continue
+
+        local = mag[fidx, a:z]
+        # Local psychoacoustic strength: stronger where audio already has energy.
+        ref = np.median(local, axis=1, keepdims=True)
+        ref = np.maximum(ref, np.percentile(local, 25, axis=1, keepdims=True))
+        strength = np.clip(ref / (np.median(ref) + MAG_EPS), 0.25, 2.5)
+
+        for bit_i, bit in enumerate(payload):
+            # Spread every bit over a different keyed pattern.
+            p = _make_keyed_pattern(len(fidx), z - a, key, bit_i, b)
+            symbol = 1.0 if bit else -1.0
+            delta = alpha * symbol * p * strength
+            wm_mag[fidx, a:z] *= np.exp(delta)
+        used_blocks += 1
 
     out = librosa.istft(
         wm_mag * np.exp(1j * phase),
-        hop_length=HOP_LENGTH,
+        hop_length=HOP,
         win_length=N_FFT,
         window="hann",
-        length=len(y)
-    )
-    return out.astype(np.float32), magnitude, wm_mag
-
-def _score_channel(y, sr, key):
-    stft = librosa.stft(
-        y, n_fft=N_FFT, hop_length=HOP_LENGTH,
-        win_length=N_FFT, window="hann"
+        length=len(y16),
     )
-    magnitude = np.abs(stft)
-    band = _band_mask(sr, N_FFT)
-    pn = _generate_pn(magnitude.shape, seed=key)
-
-    # Work on relative spectral fluctuations rather than raw magnitude.
-    x = np.log1p(magnitude)
-    x = x - np.median(x, axis=1, keepdims=True)
-
-    # Robust normalization per frequency bin.
-    scale = np.median(np.abs(x), axis=1, keepdims=True) + 1e-8
-    x = x / scale
-
-    p = pn.copy()
-    p = p - np.mean(p, axis=1, keepdims=True)
-    pscale = np.sqrt(np.mean(p ** 2, axis=1, keepdims=True)) + 1e-8
-    p = p / pscale
-
-    x = x[band]
-    p = p[band]
-
-    # Correlation-like normalized projection.
-    numerator = np.sum(x * p)
-    denominator = np.sqrt(np.sum(x * x) * np.sum(p * p)) + 1e-12
-    return float(numerator / denominator)
-
-@spaces.GPU
-def embed_watermark(audio_path, watermark_key=42, alpha=0.015, apply_masking=True):
+    out = np.clip(out, -0.999, 0.999)
+    # Return at original SR.
+    if sr != TARGET_SR:
+        out = librosa.resample(out, orig_sr=TARGET_SR, target_sr=sr)
+        out = out[:len(y)]
+        if len(out) < len(y):
+            out = np.pad(out, (0, len(y) - len(out)))
+    return out, used_blocks
+
+
+def _detect_mono(y, sr, key):
+    y16 = librosa.resample(y.astype(np.float32), orig_sr=sr, target_sr=TARGET_SR)
+    stft = librosa.stft(y16, n_fft=N_FFT, hop_length=HOP, win_length=N_FFT, window="hann")
+    mag = np.abs(stft)
+    _, fidx = _freq_bins(TARGET_SR)
+    block_frames = int(BLOCK_SECONDS * TARGET_SR / HOP)
+    n_blocks = max(1, int(np.ceil(mag.shape[1] / block_frames)))
+    payload = _payload_from_key(key)
+
+    bit_scores = [[] for _ in range(PAYLOAD_BITS)]
+
+    for b in range(n_blocks):
+        a = b * block_frames
+        z = min((b + 1) * block_frames, mag.shape[1])
+        if z - a < max(12, block_frames // 3):
+            continue
+
+        x = _frame_strength(mag, fidx)[:, a:z]
+        for bit_i, bit in enumerate(payload):
+            p = _make_keyed_pattern(len(fidx), z - a, key, bit_i, b)
+            # Normalize before correlation.
+            xx = x - np.mean(x)
+            pp = p - np.mean(p)
+            denom = (np.linalg.norm(xx) * np.linalg.norm(pp)) + MAG_EPS
+            corr = float(np.sum(xx * pp) / denom)
+            bit_scores[bit_i].append(corr)
+
+    if not all(bit_scores):
+        return 0.0, 0, [], "Pas assez de blocs exploitables."
+
+    # For each expected bit, average only the strongest half of blocks.
+    scores = []
+    for vals in bit_scores:
+        vals = np.asarray(vals, dtype=np.float32)
+        k = max(1, len(vals) // 2)
+        strongest = vals[np.argsort(np.abs(vals))[-k:]]
+        scores.append(float(np.mean(strongest)))
+
+    expected = np.where(payload > 0, 1.0, -1.0)
+    aligned = np.array(scores) * expected
+    confidence = float(np.mean(aligned))
+    positive_bits = int(np.sum(aligned > 0.0))
+    status = "WATERMARK DÉTECTÉ" if (
+        len(bit_scores[0]) >= MIN_BLOCKS
+        and confidence >= DETECT_THRESHOLD
+        and positive_bits >= int(PAYLOAD_BITS * 0.75)
+    ) else "WATERMARK NON CONFIRMÉ"
+    return confidence, positive_bits, scores, status
+
+
+def _load_audio(path):
+    y, sr = librosa.load(path, sr=None, mono=False, duration=300)
+    return y.astype(np.float32), sr
+
+
+def embed_watermark(audio_path, watermark_key=42, alpha=ALPHA):
     if not audio_path:
         return None, None, "Veuillez fournir un fichier audio."
 
     try:
-        y, sr = librosa.load(
-            audio_path, sr=None, mono=False, duration=300
-        )
-    except Exception as e:
-        return None, None, f"❌ Erreur de lecture : {e}"
+        y, sr = _load_audio(audio_path)
+        key = int(watermark_key)
+        alpha = float(alpha)
 
-    try:
         if y.ndim == 1:
-            y_wm, orig_mag, wm_mag = _embed_channel(
-                y, sr, watermark_key, alpha, apply_masking
-            )
-            y_out = y_wm
+            out, blocks = _embed_mono(y, sr, key, alpha)
+            out_sf = out
+            original_for_plot = y
         else:
             channels = []
-            orig_mag = wm_mag = None
             for ch in range(y.shape[0]):
-                y_ch, om, wmm = _embed_channel(
-                    y[ch], sr, watermark_key, alpha, apply_masking
-                )
-                channels.append(y_ch)
-                if orig_mag is None:
-                    orig_mag, wm_mag = om, wmm
-            y_out = np.vstack(channels).T
-
-        # Prevent clipping.
-        peak = np.max(np.abs(y_out)) + 1e-12
-        if peak > 0.999:
-            y_out = y_out / peak * 0.999
+                wm, _ = _embed_mono(y[ch], sr, key, alpha)
+                channels.append(wm)
+            out_sf = np.vstack(channels).T
+            original_for_plot = y[0]
 
         uid = uuid.uuid4().hex[:8]
-        output_path = f"audio_watermarked_{uid}.wav"
-        sf.write(output_path, y_out, sr, subtype="PCM_24")
-
-        D_orig = librosa.amplitude_to_db(orig_mag + 1e-10, ref=np.max)
-        D_wm = librosa.amplitude_to_db(wm_mag + 1e-10, ref=np.max)
-
-        fig, ax = plt.subplots(2, 1, figsize=(10, 6), sharex=True)
-        librosa.display.specshow(
-            D_orig, sr=sr, hop_length=HOP_LENGTH,
-            x_axis="time", y_axis="hz", ax=ax[0]
+        output_path = f"audio_watermarked_v3_{uid}.wav"
+        sf.write(output_path, out_sf, sr, subtype="PCM_24")
+
+        # Spectrogram comparison on first channel.
+        wm_plot = out_sf if out_sf.ndim == 1 else out_sf[:, 0]
+        D0 = librosa.amplitude_to_db(
+            np.abs(librosa.stft(original_for_plot, n_fft=N_FFT, hop_length=HOP)),
+            ref=np.max,
         )
-        ax[0].set_title("Spectrogramme original")
-
-        librosa.display.specshow(
-            D_wm - D_orig, sr=sr, hop_length=HOP_LENGTH,
-            x_axis="time", y_axis="hz", ax=ax[1], cmap="magma"
+        D1 = librosa.amplitude_to_db(
+            np.abs(librosa.stft(wm_plot, n_fft=N_FFT, hop_length=HOP)),
+            ref=np.max,
         )
-        ax[1].set_title("Empreinte différentielle du watermark")
+        diff = D1 - D0
 
+        fig, ax = plt.subplots(2, 1, figsize=(11, 7), sharex=True)
+        librosa.display.specshow(D0, sr=sr, hop_length=HOP, x_axis="time", y_axis="hz", ax=ax[0])
+        ax[0].set_title("Original — spectrogramme")
+        librosa.display.specshow(diff, sr=sr, hop_length=HOP, x_axis="time", y_axis="hz", ax=ax[1])
+        ax[1].set_title("Différence spectrale — watermark v3")
         plt.tight_layout()
-        spec_path = f"spectrogram_{uid}.png"
-        plt.savefig(spec_path, dpi=150)
+        plot_path = f"spectrogram_v3_{uid}.png"
+        plt.savefig(plot_path, dpi=140)
         plt.close(fig)
 
-        score = _score_channel(y_out if y_out.ndim == 1 else y_out[:, 0], sr, watermark_key)
-
-        return (
-            output_path,
-            spec_path,
-            f"✅ Watermark v2 appliqué.\n"
-            f"Clé : {int(watermark_key)}\n"
-            f"Alpha : {alpha:.3f}\n"
-            f"Score de contrôle interne : {score:.5f}\n"
-            f"Fichier exporté en WAV PCM 24-bit."
+        return output_path, plot_path, (
+            f"✅ Watermark v3 injecté. Blocs utilisés: {blocks}. "
+            f"Clé: {key}. Alpha: {alpha:.3f}. "
+            f"Le watermark est réparti dans le spectre, sans porteuse ultrasonique fixe."
         )
     except Exception as e:
-        return None, None, f"❌ Erreur pendant l'injection : {e}"
+        return None, None, f"❌ Erreur: {e}"
+
 
-@spaces.GPU
 def detect_watermark(audio_path, watermark_key=42):
     if not audio_path:
         return "Veuillez fournir un fichier audio."
 
     try:
-        y, sr = librosa.load(
-            audio_path, sr=None, mono=False, duration=300
-        )
-    except Exception as e:
-        return f"❌ Erreur de lecture : {e}"
+        y, sr = _load_audio(audio_path)
+        key = int(watermark_key)
 
-    try:
         if y.ndim == 1:
-            scores = [_score_channel(y, sr, watermark_key)]
+            conf, pos, scores, status = _detect_mono(y, sr, key)
         else:
-            scores = [
-                _score_channel(y[ch], sr, watermark_key)
-                for ch in range(y.shape[0])
-            ]
-
-        # Use the strongest channel, while reporting all channels.
-        score = max(scores)
-
-        # Conservative first-pass threshold.
-        # This must be calibrated with genuine non-watermarked audio.
-        threshold = 0.020
-
-        detected = score >= threshold
-        confidence = min(99.9, max(0.0, abs(score) / threshold * 50.0))
-
-        result = (
-            "✅ WATERMARK AUDIOSHIELD V2 DÉTECTÉ."
-            if detected else
-            "❌ Watermark non détecté avec cette clé."
-        )
+            results = [_detect_mono(y[ch], sr, key) for ch in range(y.shape[0])]
+            conf = float(np.mean([r[0] for r in results]))
+            pos = int(np.mean([r[1] for r in results]))
+            scores = results[0][2]
+            status = "WATERMARK DÉTECTÉ" if all(r[3] == "WATERMARK DÉTECTÉ" for r in results) else "WATERMARK NON CONFIRMÉ"
 
         return (
-            f"{result}\n\n"
-            f"Score : {score:.5f}\n"
-            f"Seuil provisoire : {threshold:.5f}\n"
-            f"Canaux : {', '.join(f'{s:.5f}' for s in scores)}\n"
-            f"Confiance indicative : {confidence:.1f}%\n\n"
-            f"⚠️ Le seuil doit être calibré sur un corpus "
-            f"d'audios non tatoués pour mesurer les faux positifs."
+            f"{'🟢' if 'DÉTECTÉ' in status else '🔴'} {status}\n\n"
+            f"Confiance normalisée : {conf:.3f}\n"
+            f"Bits cohérents : {pos}/{PAYLOAD_BITS}\n"
+            f"Seuil : {DETECT_THRESHOLD:.3f}\n"
+            f"Clé testée : {key}\n\n"
+            "⚠️ Le résultat est une décision statistique de ce prototype ; "
+            "il ne constitue pas à lui seul une preuve cryptographique de provenance."
         )
     except Exception as e:
-        return f"❌ Erreur pendant la détection : {e}"
+        return f"❌ Erreur: {e}"
+
 
-# ==========================================================
-# Interface Gradio
-# ==========================================================
+with gr.Blocks(title="AudioShield v3 — Robust Watermark") as demo:
+    gr.Markdown(
+        """
+# 🛡️ AudioShield v3 — Watermarking audio robuste
 
-with gr.Blocks(title="AudioShield v2 — Watermarking") as demo:
-    gr.Markdown("""
-    # 🛡️ AudioShield v2 — Watermarking robuste
-    Watermark audio invisible par étalement pseudo-aléatoire,
-    injection relative au spectre et détection différentielle.
+Watermark invisible **à spectre étalé et clé secrète**, sans tonalité ultrasonique fixe.
 
-    **Entrées :** MP3, WAV, FLAC, OGG, M4A, AAC, etc.
-    **Sortie watermarkée :** WAV PCM 24-bit.
-    """)
+- MP3 / WAV / FLAC / OGG / M4A / AAC / AIFF, selon les codecs disponibles
+- Mono et stéréo
+- Payload déterministe de 32 bits
+- Synchronisation par blocs
+- Détection par corrélation multi-blocs
+- Analyse spectrale Original / Watermark
+"""
+    )
 
-    with gr.Tab("1. Injecter le Watermark"):
+    with gr.Tab("1. Injecter"):
         with gr.Row():
             with gr.Column():
-                audio_in = gr.Audio(
-                    type="filepath",
-                    label="Audio source"
-                )
-                key_in = gr.Number(
-                    value=42, label="Clé secrète (Seed)", precision=0
-                )
+                audio_in = gr.Audio(type="filepath", label="Audio source")
+                key_in = gr.Number(value=42, label="Clé secrète", precision=0)
                 alpha_in = gr.Slider(
-                    minimum=0.003, maximum=0.05, value=0.015,
-                    step=0.001, label="Force d'injection (Alpha)"
-                )
-                mask_in = gr.Checkbox(
-                    value=True, label="Masque psychoacoustique"
+                    minimum=0.006, maximum=0.030, value=ALPHA, step=0.001,
+                    label="Force d'injection"
                 )
-                btn_embed = gr.Button(
-                    "Appliquer le Watermark v2", variant="primary"
-                )
-
+                btn_embed = gr.Button("Appliquer le watermark v3", variant="primary")
             with gr.Column():
-                audio_out = gr.Audio(
-                    label="Audio tatoué (WAV PCM 24-bit)"
-                )
-                plot_out = gr.Image(
-                    label="Empreinte spectrale"
-                )
-                text_out = gr.Textbox(
-                    label="Statut", lines=6
-                )
+                audio_out = gr.Audio(label="Audio watermarké — WAV PCM 24-bit")
+                plot_out = gr.Image(label="Analyse spectrale")
+                text_out = gr.Textbox(label="Statut", lines=4)
 
         btn_embed.click(
             embed_watermark,
-            inputs=[audio_in, key_in, alpha_in, mask_in],
-            outputs=[audio_out, plot_out, text_out]
+            inputs=[audio_in, key_in, alpha_in],
+            outputs=[audio_out, plot_out, text_out],
         )
 
-    with gr.Tab("2. Vérifier / Détecter"):
+    with gr.Tab("2. Détecter"):
         with gr.Row():
             with gr.Column():
-                audio_verify = gr.Audio(
-                    type="filepath",
-                    label="Audio à analyser"
-                )
-                key_verify = gr.Number(
-                    value=42, label="Clé secrète (Seed)", precision=0
-                )
-                btn_detect = gr.Button(
-                    "Détecter le Watermark v2", variant="secondary"
-                )
-
+                audio_verify = gr.Audio(type="filepath", label="Audio à vérifier")
+                key_verify = gr.Number(value=42, label="Clé secrète", precision=0)
+                btn_detect = gr.Button("Vérifier le watermark", variant="secondary")
             with gr.Column():
-                detect_out = gr.Textbox(
-                    label="Résultat de détection", lines=8
-                )
+                detect_out = gr.Textbox(label="Résultat", lines=8)
 
         btn_detect.click(
             detect_watermark,
             inputs=[audio_verify, key_verify],
-            outputs=[detect_out]
+            outputs=[detect_out],
         )
 
-    gr.Markdown("""
-    ### ⚠️ Validation scientifique
-    Un vrai benchmark doit comparer des fichiers tatoués et non tatoués
-    et mesurer TPR/FNR/FPR après MP3, bruit, resampling, filtrage,
-    changement de volume et autres attaques.
-    """)
+    gr.Markdown(
+        """
+### ⚠️ Validation recommandée
+Avant toute affirmation de robustesse, tester séparément :
+**original**, **watermarké**, MP3, bruit, resampling, variation de volume,
+low-pass/high-pass, time-stretch et pitch-shift, puis calculer les faux positifs,
+faux négatifs et BER.
+"""
+    )
 
 demo.queue().launch()