Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 import os
 import uuid
 import hashlib
@@ -10,97 +11,80 @@ 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 = 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
+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)
+    h = hashlib.sha256(str(int(key)).encode()).digest()
+    return int.from_bytes(h[:8], "little") % (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)
+    digest = hashlib.sha256(
+        f"AudioShield-v3:{int(key)}".encode()
+    ).digest()
+    return np.unpackbits(
+        np.frombuffer(digest[:4], dtype=np.uint8)
+    ).astype(np.int8)
 
 
 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]
+    upper = min(HIGH_HZ, sr / 2 - 300)
+    idx = np.where((freqs >= LOW_HZ) & (freqs <= upper))[0]
     if len(idx) < 20:
-        raise ValueError("Échantillonnage trop faible pour la bande de watermark.")
+        raise ValueError("Fréquence d'échantillonnage trop faible.")
     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)
+        kernel = np.array([1, 2, 3, 2, 1], dtype=np.float32)
+        kernel /= kernel.sum()
+        p = np.apply_along_axis(
+            lambda row: np.convolve(row, kernel, mode="same"),
+            1, p
+        )
+
+    p /= np.sqrt(np.mean(p * p) + 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)
+    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)
+    phase = np.angle(stft)
 
     _, fidx = _freq_bins(TARGET_SR)
-    block_frames = int(BLOCK_SECONDS * TARGET_SR / HOP)
+    block_frames = max(1, int(BLOCK_SECONDS * TARGET_SR / HOP))
     n_blocks = max(1, int(np.ceil(mag.shape[1] / block_frames)))
     payload = _payload_from_key(key)
 
@@ -110,21 +94,28 @@ def _embed_mono(y, sr, key, alpha):
     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)
+        ref = np.maximum(
+            ref,
+            np.percentile(local, 25, axis=1, keepdims=True)
+        )
+        strength = np.clip(
+            ref / (np.median(ref) + 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)
+            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(
@@ -132,24 +123,37 @@ def _embed_mono(y, sr, key, alpha):
         hop_length=HOP,
         win_length=N_FFT,
         window="hann",
-        length=len(y16),
+        length=len(y16)
     )
+
     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 = 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")
+    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)
+    block_frames = max(1, int(BLOCK_SECONDS * TARGET_SR / HOP))
     n_blocks = max(1, int(np.ceil(mag.shape[1] / block_frames)))
     payload = _payload_from_key(key)
 
@@ -158,47 +162,80 @@ def _detect_mono(y, sr, key):
     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]
+        x = mag[fidx, a:z]
+        med = np.median(x, axis=1, keepdims=True)
+        mad = np.median(np.abs(x - med), axis=1, keepdims=True) + EPS
+        x = np.clip((x - med) / (4.0 * mad), -3.0, 3.0)
+
         for bit_i, bit in enumerate(payload):
-            p = _make_keyed_pattern(len(fidx), z - a, key, bit_i, b)
-            # Normalize before correlation.
+            p = _make_keyed_pattern(
+                len(fidx), z - a, key, bit_i, b
+            )
+
             xx = x - np.mean(x)
             pp = p - np.mean(p)
-            denom = (np.linalg.norm(xx) * np.linalg.norm(pp)) + MAG_EPS
+
+            denom = (
+                np.linalg.norm(xx) *
+                np.linalg.norm(pp)
+            ) + 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."
+        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:]]
+        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
+    aligned = np.asarray(scores) * expected
+
     confidence = float(np.mean(aligned))
-    positive_bits = int(np.sum(aligned > 0.0))
-    status = "WATERMARK DÉTECTÉ" if (
+    positive_bits = int(np.sum(aligned > 0))
+
+    detected = (
         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
+    )
+
+    return (
+        confidence,
+        positive_bits,
+        "WATERMARK DÉTECTÉ"
+        if detected else
+        "WATERMARK NON CONFIRMÉ"
+    )
 
 
 def _load_audio(path):
-    y, sr = librosa.load(path, sr=None, mono=False, duration=300)
+    y, sr = librosa.load(
+        path, sr=None, mono=False, duration=300
+    )
     return y.astype(np.float32), sr
 
 
+# ------------------------------------------------------------
+# ZeroGPU functions
+# ------------------------------------------------------------
+# The @spaces.GPU decorator is required by Hugging Face
+# ZeroGPU. Keep it on the OUTER processing functions.
+# ------------------------------------------------------------
+
+@spaces.GPU(duration=120)
 def embed_watermark(audio_path, watermark_key=42, alpha=ALPHA):
     if not audio_path:
         return None, None, "Veuillez fournir un fichier audio."
@@ -211,50 +248,100 @@ def embed_watermark(audio_path, watermark_key=42, alpha=ALPHA):
         if y.ndim == 1:
             out, blocks = _embed_mono(y, sr, key, alpha)
             out_sf = out
-            original_for_plot = y
+            original = y
         else:
             channels = []
+            blocks = 0
+
             for ch in range(y.shape[0]):
-                wm, _ = _embed_mono(y[ch], sr, key, alpha)
+                wm, b = _embed_mono(
+                    y[ch], sr, key, alpha
+                )
                 channels.append(wm)
+                blocks = max(blocks, b)
+
             out_sf = np.vstack(channels).T
-            original_for_plot = y[0]
+            original = y[0]
 
         uid = uuid.uuid4().hex[:8]
-        output_path = f"audio_watermarked_v3_{uid}.wav"
-        sf.write(output_path, out_sf, sr, subtype="PCM_24")
+        output_path = (
+            f"audio_watermarked_v3_{uid}.wav"
+        )
+
+        sf.write(
+            output_path,
+            out_sf,
+            sr,
+            subtype="PCM_24"
+        )
+
+        wm_plot = (
+            out_sf if out_sf.ndim == 1
+            else out_sf[:, 0]
+        )
 
-        # 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,
+            np.abs(librosa.stft(
+                original,
+                n_fft=N_FFT,
+                hop_length=HOP
+            )),
+            ref=np.max
         )
+
         D1 = librosa.amplitude_to_db(
-            np.abs(librosa.stft(wm_plot, n_fft=N_FFT, hop_length=HOP)),
-            ref=np.max,
+            np.abs(librosa.stft(
+                wm_plot,
+                n_fft=N_FFT,
+                hop_length=HOP
+            )),
+            ref=np.max
+        )
+
+        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(
+            D1 - D0, sr=sr, hop_length=HOP,
+            x_axis="time", y_axis="hz",
+            ax=ax[1]
+        )
+        ax[1].set_title(
+            "Différence spectrale — Watermark v3"
         )
-        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()
+
         plot_path = f"spectrogram_v3_{uid}.png"
         plt.savefig(plot_path, dpi=140)
         plt.close(fig)
 
-        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."
+        return (
+            output_path,
+            plot_path,
+            "✅ Watermark v3 injecté.\n"
+            f"Blocs utilisés : {blocks}\n"
+            f"Clé : {key}\n"
+            f"Alpha : {alpha:.3f}\n\n"
+            "Watermark réparti dans le spectre "
+            "sans porteuse ultrasonique fixe."
         )
+
     except Exception as e:
-        return None, None, f"❌ Erreur: {e}"
+        return None, None, f"❌ Erreur : {e}"
 
 
+@spaces.GPU(duration=120)
 def detect_watermark(audio_path, watermark_key=42):
     if not audio_path:
         return "Veuillez fournir un fichier audio."
@@ -264,87 +351,189 @@ def detect_watermark(audio_path, watermark_key=42):
         key = int(watermark_key)
 
         if y.ndim == 1:
-            conf, pos, scores, status = _detect_mono(y, sr, key)
+            conf, bits, status = _detect_mono(
+                y, sr, key
+            )
         else:
-            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É"
+            results = [
+                _detect_mono(y[ch], sr, key)
+                for ch in range(y.shape[0])
+            ]
+
+            conf = float(
+                np.mean([r[0] for r in results])
+            )
+            bits = int(
+                np.mean([r[1] for r in results])
+            )
+            status = (
+                "WATERMARK DÉTECTÉ"
+                if all(
+                    r[2] == "WATERMARK DÉTECTÉ"
+                    for r in results
+                )
+                else
+                "WATERMARK NON CONFIRMÉ"
+            )
+
+        icon = (
+            "🟢"
+            if status == "WATERMARK DÉTECTÉ"
+            else "🔴"
+        )
 
         return (
-            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"{icon} {status}\n\n"
+            f"Confiance : {conf:.3f}\n"
+            f"Bits cohérents : {bits}/{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."
+            "⚠️ Résultat statistique du prototype. "
+            "Ce résultat n'est pas une preuve cryptographique "
+            "de provenance."
         )
+
     except Exception as e:
-        return f"❌ Erreur: {e}"
+        return f"❌ Erreur : {e}"
+
 
+# ------------------------------------------------------------
+# Interface
+# ------------------------------------------------------------
+
+with gr.Blocks(
+    title="AudioShield v3 — Robust Watermark"
+) as demo:
 
-with gr.Blocks(title="AudioShield v3 — Robust Watermark") as demo:
     gr.Markdown(
         """
 # 🛡️ AudioShield v3 — Watermarking audio robuste
 
-Watermark invisible **à spectre étalé et clé secrète**, sans tonalité ultrasonique fixe.
+Watermark invisible **à spectre étalé et clé secrète**,
+sans tonalité ultrasonique fixe.
 
-- MP3 / WAV / FLAC / OGG / M4A / AAC / AIFF, selon les codecs disponibles
+- MP3 / WAV / FLAC / OGG / M4A / AAC / AIFF
 - Mono et stéréo
 - Payload déterministe de 32 bits
-- Synchronisation par blocs
-- Détection par corrélation multi-blocs
+- Répétition par blocs
+- Détection multi-blocs
+- Score de confiance
 - Analyse spectrale Original / Watermark
 """
     )
 
-    with gr.Tab("1. Injecter"):
+    with gr.Tab("1. Injecter le Watermark"):
+
         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", 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.006, maximum=0.030, value=ALPHA, step=0.001,
+                    minimum=0.006,
+                    maximum=0.030,
+                    value=ALPHA,
+                    step=0.001,
                     label="Force d'injection"
                 )
-                btn_embed = gr.Button("Appliquer le watermark v3", variant="primary")
+
+                btn_embed = gr.Button(
+                    "Appliquer le Watermark v3",
+                    variant="primary"
+                )
+
             with gr.Column():
-                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)
+
+                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=6
+                )
 
         btn_embed.click(
             embed_watermark,
-            inputs=[audio_in, key_in, alpha_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. Détecter"):
+    with gr.Tab("2. Vérifier / Détecter"):
+
         with gr.Row():
+
             with gr.Column():
-                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")
+
+                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", lines=8)
+
+                detect_out = gr.Textbox(
+                    label="Résultat de détection",
+                    lines=9
+                )
 
         btn_detect.click(
             detect_watermark,
-            inputs=[audio_verify, key_verify],
-            outputs=[detect_out],
+            inputs=[
+                audio_verify,
+                key_verify
+            ],
+            outputs=[detect_out]
         )
 
     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.
+### ⚠️ Validation
+
+Tester séparément :
+
+**ORIGINAL → doit rester NON CONFIRMÉ**
+
+**WATERMARKÉ → doit être DÉTECTÉ**
+
+Puis tester MP3, bruit, resampling, variation de volume,
+low-pass/high-pass, time-stretch et pitch-shift.
 """
     )
 
+
+# Required for ZeroGPU request handling.
 demo.queue().launch()