app.py

"""
StegNet — Gradio web interface.

Entry point for HuggingFace Spaces (or any Gradio-compatible host).
Set the HF_API_KEY environment variable (Spaces secret) to enable
the "Generate Image" tab.

Run locally:
    python app.py
"""

import os
import sys

# Make sure the package root is on the path when run directly
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import gradio as gr

from app.models.DEEP_STEGO.hide_image import hide_image
from app.models.DEEP_STEGO.reveal_image import reveal_image
from app.models.ESRGAN.upscale_image import upscale_image
from app.models.encryption import aes, blowfish
from app.models.StableDiffusionAPI.StableDiffusionV2 import generate


# ---------------------------------------------------------------------------
# Handler functions
# ---------------------------------------------------------------------------

def hide_fn(cover, secret):
    if cover is None or secret is None:
        return None, "Please upload both a cover image and a secret image."
    try:
        output_path = hide_image(cover, secret)
        return output_path, "Secret image hidden successfully!"
    except Exception as e:
        return None, f"Error: {e}"


def reveal_fn(steg):
    if steg is None:
        return None, "Please upload a steg image."
    try:
        output_path = reveal_image(steg)
        return output_path, "Secret image revealed successfully!"
    except Exception as e:
        return None, f"Error: {e}"


def encrypt_fn(image, key, algo):
    if image is None:
        return None, "Please upload an image."
    if not key:
        return None, "Please enter a secret key."
    try:
        if algo == "AES":
            aes.encrypt(image, key)
        else:
            blowfish.encrypt(image, key)
        enc_path = image + ".enc"
        return enc_path, f"Image encrypted with {algo}. Download the file below."
    except Exception as e:
        return None, f"Error: {e}"


def decrypt_fn(enc_file, key, algo):
    if enc_file is None:
        return None, "Please upload an encrypted .enc file."
    if not key:
        return None, "Please enter the secret key."
    try:
        if algo == "AES":
            code, filename = aes.decrypt(enc_file, key)
        else:
            code, filename = blowfish.decrypt(enc_file, key)
        if code == -1:
            return None, "Decryption failed — wrong key."
        return filename, "Image decrypted successfully!"
    except Exception as e:
        return None, f"Error: {e}"


def upscale_fn(image):
    if image is None:
        return None, "Please upload a low-resolution image."
    try:
        output_path = upscale_image(image)
        return output_path, "Image upscaled 4× successfully!"
    except Exception as e:
        return None, f"Error: {e}"


def generate_fn(prompt):
    if not prompt or not prompt.strip():
        return None, "Please enter a text prompt."
    if not os.environ.get("HF_API_KEY"):
        return None, "HF_API_KEY is not set. Add it as a Space secret to enable image generation."
    try:
        img = generate(prompt.strip())
        return img, "Image generated successfully!"
    except Exception as e:
        return None, f"Error: {e}"


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

DESCRIPTION = """
# StegNet
### AI-Based Secure Image Steganography Platform

Hide secret images inside cover images using deep learning, encrypt the result,
reveal hidden images, enhance resolution with ESRGAN, or generate new images
with Stable Diffusion — all in one place.
"""

with gr.Blocks(title="StegNet", theme=gr.themes.Soft()) as demo:
    gr.Markdown(DESCRIPTION)

    # ── Tab 1: Hide Image ──────────────────────────────────────────────────
    with gr.Tab("🔒 Hide Image"):
        gr.Markdown("Upload a **cover image** and a **secret image**. The model embeds the secret inside the cover.")
        with gr.Row():
            cover_input = gr.Image(type="filepath", label="Cover Image")
            secret_input = gr.Image(type="filepath", label="Secret Image")
        hide_btn = gr.Button("Hide Secret", variant="primary")
        with gr.Row():
            steg_output = gr.Image(label="Steg Image (download me)")
            hide_status = gr.Textbox(label="Status", interactive=False, lines=2)
        hide_btn.click(hide_fn, inputs=[cover_input, secret_input], outputs=[steg_output, hide_status])

    # ── Tab 2: Reveal Image ────────────────────────────────────────────────
    with gr.Tab("🔍 Reveal Image"):
        gr.Markdown("Upload a **steg image** to extract the hidden secret image.")
        steg_input = gr.Image(type="filepath", label="Steg Image")
        reveal_btn = gr.Button("Reveal Secret", variant="primary")
        with gr.Row():
            revealed_output = gr.Image(label="Revealed Secret Image")
            reveal_status = gr.Textbox(label="Status", interactive=False, lines=2)
        reveal_btn.click(reveal_fn, inputs=[steg_input], outputs=[revealed_output, reveal_status])

    # ── Tab 3: Encrypt Image ───────────────────────────────────────────────
    with gr.Tab("🛡️ Encrypt Image"):
        gr.Markdown("Encrypt any image file (e.g. your steg image) with AES or Blowfish.")
        enc_image_input = gr.Image(type="filepath", label="Image to Encrypt")
        with gr.Row():
            enc_key_input = gr.Textbox(label="Secret Key", type="password", placeholder="Enter a strong key…")
            enc_algo = gr.Radio(["AES", "Blowfish"], value="AES", label="Algorithm")
        enc_btn = gr.Button("Encrypt", variant="primary")
        with gr.Row():
            enc_file_output = gr.File(label="Encrypted File (.enc) — Download")
            enc_status = gr.Textbox(label="Status", interactive=False, lines=2)
        enc_btn.click(encrypt_fn, inputs=[enc_image_input, enc_key_input, enc_algo], outputs=[enc_file_output, enc_status])

    # ── Tab 4: Decrypt Image ───────────────────────────────────────────────
    with gr.Tab("🔓 Decrypt Image"):
        gr.Markdown("Upload an encrypted **.enc** file and the correct key to recover the image.")
        dec_file_input = gr.File(label="Encrypted File (.enc)", type="filepath")
        with gr.Row():
            dec_key_input = gr.Textbox(label="Secret Key", type="password", placeholder="Enter the key used for encryption…")
            dec_algo = gr.Radio(["AES", "Blowfish"], value="AES", label="Algorithm")
        dec_btn = gr.Button("Decrypt", variant="primary")
        with gr.Row():
            dec_image_output = gr.Image(label="Decrypted Image")
            dec_status = gr.Textbox(label="Status", interactive=False, lines=2)
        dec_btn.click(decrypt_fn, inputs=[dec_file_input, dec_key_input, dec_algo], outputs=[dec_image_output, dec_status])

    # ── Tab 5: Super Resolution ────────────────────────────────────────────
    with gr.Tab("✨ Super Resolution"):
        gr.Markdown("Upscale any image **4×** using ESRGAN.")
        lr_input = gr.Image(type="filepath", label="Low-Resolution Image")
        upscale_btn = gr.Button("Upscale 4×", variant="primary")
        with gr.Row():
            hr_output = gr.Image(label="High-Resolution Output")
            upscale_status = gr.Textbox(label="Status", interactive=False, lines=2)
        upscale_btn.click(upscale_fn, inputs=[lr_input], outputs=[hr_output, upscale_status])

    # ── Tab 6: Generate Image ──────────────────────────────────────────────
    with gr.Tab("🎨 Generate Image"):
        gr.Markdown(
            "Generate an image from a text prompt using Stable Diffusion 2.1.  \n"
            "*(Requires `HF_API_KEY` set as a Space secret.)*"
        )
        prompt_input = gr.Textbox(
            label="Text Prompt",
            placeholder="A futuristic city at sunset, digital art…",
            lines=2,
        )
        gen_btn = gr.Button("Generate", variant="primary")
        with gr.Row():
            gen_output = gr.Image(label="Generated Image")
            gen_status = gr.Textbox(label="Status", interactive=False, lines=2)
        gen_btn.click(generate_fn, inputs=[prompt_input], outputs=[gen_output, gen_status])

demo.launch()