bluefoxbox/stealth-mark

1import cv22import numpy as np3import gradio as gr4from imwatermark import WatermarkEncoder, WatermarkDecoder5import re6import os7 8# ⚙️ System Configuration9FIXED_BYTES = 1610BIT_LENGTH = FIXED_BYTES * 811MAX_IMAGE_SIZE = 1920  # Limit to FHD to prevent CPU overload12WM_ALGO = 'dwtDctSvd'  # Strongest algorithm against compression13 14def encode_watermark(img, secret_text):15    # 1. Validation16    if img is None or not secret_text:17        return None, None, "Error: Image and Secret Text are both required."18    19    height, width = img.shape[:2]20    if height > MAX_IMAGE_SIZE or width > MAX_IMAGE_SIZE:21        return None, None, f"Error: Image resolution too high. Max: {MAX_IMAGE_SIZE}px."22 23    if not re.match(r'^[a-zA-Z0-9_]+$', secret_text):24        return None, None, "Error: Use only alphanumeric characters and underscores."25 26    if len(secret_text.encode('utf-8')) > FIXED_BYTES:27        return None, None, f"Error: Text is too long (Max {FIXED_BYTES} bytes)."28 29    # 2. Data Preparation (Null Padding)30    padded_text = secret_text.ljust(FIXED_BYTES, '\x00')31    watermark_bytes = padded_text.encode('utf-8')32 33    # 3. Encoding Process34    # Gradio provides RGB, OpenCV needs BGR35    bgr_img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)36 37    encoder = WatermarkEncoder()38    encoder.set_watermark('bytes', watermark_bytes)39    40    try:41        bgr_encoded = encoder.encode(bgr_img, WM_ALGO)42        43        # 🚨 [CRITICAL] Save to actual PNG file to prevent browser WebP conversion44        output_filename = "watermarked_output.png"45        cv2.imwrite(output_filename, bgr_encoded)46        47        # Convert back to RGB for Gradio preview48        rgb_encoded = cv2.cvtColor(bgr_encoded, cv2.COLOR_BGR2RGB)49        50        return rgb_encoded, output_filename, "✅ Success! Download the lossless PNG file below."51    except Exception as e:52        return None, None, f"❌ Encoding Failed: {str(e)}"53 54def decode_watermark(img):55    if img is None:56        return "Error: Upload an image to decode."57 58    height, width = img.shape[:2]59    if height > 3000 or width > 3000:60        return "Error: Image too large for security reasons."61 62    bgr_wm = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)63    decoder = WatermarkDecoder('bytes', BIT_LENGTH)64 65    try:66        watermark_bytes = decoder.decode(bgr_wm, WM_ALGO)67        try:68            # Strip null bytes and decode69            result = watermark_bytes.decode('utf-8').rstrip('\x00')70            if not result:71                return "❌ No watermark detected or data is corrupted."72            return f"✅ Decoded Text: [{result}]"73        except UnicodeDecodeError:74            # Recovery mode for corrupted bits75            forced_result = watermark_bytes.decode('utf-8', errors='replace').rstrip('\x00')76            return f"⚠️ Bit corruption detected! Recovered text: [{forced_result}]"77    except Exception as e:78        return f"❌ Decoding Failed: Watermark is missing or destroyed."79 80# 🎨 UI Construction (Gradio Blocks)81with gr.Blocks(theme=gr.themes.Monochrome()) as demo:82    gr.Markdown("# 💧 Invisible Watermark & Steganography Tool")83    gr.Markdown("Hide secret text inside images using the robust DWT-DCT-SVD algorithm.")84    85    with gr.Tabs():86        # --- ENCODE TAB ---87        with gr.TabItem("🔒 Encode Watermark"):88            with gr.Row():89                with gr.Column():90                    img_input = gr.Image(label="Source Image", type="numpy")91                    txt_input = gr.Textbox(label="Secret Text", placeholder="e.g. user_01")92                    enc_btn = gr.Button("Embed Watermark", variant="primary")93                with gr.Column():94                    img_output = gr.Image(label="Preview (Browser-rendered)")95                    # 🚨 Lossless file download component96                    file_output = gr.File(label="Download Lossless PNG File")97                    status_enc = gr.Textbox(label="Status")98            99            enc_btn.click(100                fn=encode_watermark, 101                inputs=[img_input, txt_input], 102                outputs=[img_output, file_output, status_enc]103            )104 105        # --- DECODE TAB ---106        with gr.TabItem("🔓 Decode Watermark"):107            with gr.Row():108                with gr.Column():109                    img_input_dec = gr.Image(label="Watermarked Image", type="numpy")110                    dec_btn = gr.Button("Extract Watermark", variant="primary")111                with gr.Column():112                    txt_output = gr.Textbox(label="Extracted Result")113            114            dec_btn.click(115                fn=decode_watermark, 116                inputs=[img_input_dec], 117                outputs=[txt_output]118            )119 120    gr.Markdown("---")121    gr.Markdown("⚠️ **Note:** To ensure the watermark survives, download the file using the **Download** link, not 'Save Image As' in your browser.")122 123# 🚨 API Queue for DoS Protection124if __name__ == "__main__":125    demo.queue(max_size=10).launch()