app.py

import gradio as gr
from PIL import Image, ImageDraw, PngImagePlugin
import io
import os
import cv2
import numpy as np
import uuid


# Helper functions
def to_bin(data):
    if isinstance(data, str):
        return ''.join([format(ord(i), "08b") for i in data])
    elif isinstance(data, bytes):
        return ''.join([format(i, "08b") for i in data])
    elif isinstance(data, np.ndarray):
        return [format(i, "08b") for i in data]
    elif isinstance(data, int) or isinstance(data, np.uint8):
        return format(data, "08b")
    else:
        raise TypeError("Type not supported.")


def encode(image_path, secret_data):
    try:
        # Read the image
        image = cv2.imread(image_path)
        if image is None:
            return None, "Failed to read image"
        
        # Calculate maximum bytes for encoding
        n_bytes = (image.shape[0] * image.shape[1] * 3) // 8
        
        # Prepare secret data
        secret_data_with_delimiter = f'{secret_data}#####'
        if len(secret_data_with_delimiter) >= n_bytes:
            return None, "Watermark is too large for Image Size"
        
        secret_data_with_delimiter += "====="
        binary_secret_data = to_bin(secret_data_with_delimiter)
        data_len = len(binary_secret_data)
        
        # Create a copy of the image
        watermarked_image = image.copy()
        data_index = 0
        
        # Encode the data
        for i in range(watermarked_image.shape[0]):
            for j in range(watermarked_image.shape[1]):
                for k in range(3):  # RGB channels
                    if data_index < data_len:
                        # Get the current pixel value
                        pixel = watermarked_image[i, j, k]
                        # Convert to binary and modify the least significant bit
                        binary_pixel = format(pixel, '08b')
                        new_pixel = binary_pixel[:-1] + binary_secret_data[data_index]
                        # Update the pixel value
                        watermarked_image[i, j, k] = int(new_pixel, 2)
                        data_index += 1
                    else:
                        break
                if data_index >= data_len:
                    break
            if data_index >= data_len:
                break
        
        return watermarked_image, None
    except Exception as e:
        return None, f"Error during encoding: {str(e)}"


def decode(image_path):
    try:
        # Read the image
        image = cv2.imread(image_path)
        if image is None:
            return "Failed to read image"
        
        binary_data = ""
        # Extract the least significant bits
        for i in range(image.shape[0]):
            for j in range(image.shape[1]):
                for k in range(3):
                    pixel = format(image[i, j, k], '08b')
                    binary_data += pixel[-1]
        
        # Convert binary to text
        decoded_data = ""
        for i in range(0, len(binary_data), 8):
            byte = binary_data[i:i+8]
            if len(byte) == 8:
                decoded_data += chr(int(byte, 2))
                # Check for ending delimiter
                if decoded_data[-5:] == "=====":
                    break
        
        # Remove delimiters and return the message
        if "#####" in decoded_data:
            return decoded_data.split("#####")[0]
        return "No watermark found"
        
    except Exception as e:
        return f"Error during decoding: {str(e)}"


def png_encode(im_name, extra):
    try:
        im = Image.open(im_name)
        info = PngImagePlugin.PngInfo()
        info.add_text("TXT", extra)

        unique_id = str(uuid.uuid4())[:8]
        filename, ext = os.path.splitext(im_name)
        new_filename = f"{filename}_{unique_id}{ext}"
        im.save(new_filename, "PNG", pnginfo=info)

        width, height = im.size
        rect_width, rect_height = 200, 50
        x = width - rect_width - 10
        y = height - rect_height - 10
        global png_encode_coords
        png_encode_coords = (x, y, rect_width, rect_height)

        return new_filename, None

    except Exception as e:
        return None, str(e)


def highlight_watermark(image, coords=(0, 0, 100, 50), color="red", width=5):
    try:
        if isinstance(image, np.ndarray):
            image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(image)
        x, y, w, h = coords
        draw.rectangle((max(0, x - 5), max(0, y - 5), x + w + 5, y + h + 5), outline=color, width=width)
        return image
    except Exception as e:
        print(f"Error highlighting: {e}")
        return image


# Gradio function handlers
def choose_encode(inp_im, inp_mark, cho):
    try:
        if not inp_im:
            return None, "Please upload an image.", None
        if not inp_mark:
            return None, "Please enter watermark text.", None

        if cho == "stegan":
            watermarked_image, error_msg = encode(inp_im, inp_mark)
            if error_msg:
                return None, error_msg, None

            output_path = os.path.splitext(inp_im)[0] + "_watermarked.png"
            cv2.imwrite(output_path, watermarked_image)
            return output_path, "Steganography watermark added successfully.", output_path

        elif cho == "pnginfo":
            output_path, error_msg = png_encode(inp_im, inp_mark)
            if error_msg:
                return None, error_msg, None
            return output_path, "PNG metadata watermark added successfully.", output_path
    except Exception as e:
        return None, f"An unexpected error occurred: {e}", None


def detect_watermark(det_im):
    if not det_im:
        return "Please upload an image."

    detected_text = decode(det_im)
    if not detected_text or detected_text.startswith("Error"):
        return "No watermark detected or not encoded with this tool."

    return detected_text


# Gradio app
with gr.Blocks() as app:
    with gr.Tab("Add Watermark"):
        cho = gr.Radio(choices=["stegan", "pnginfo"], value="stegan", label="Watermark Method")
        with gr.Row():
            with gr.Column():
                inp_im = gr.Image(label="Input Image", type="filepath")
                inp_mark = gr.Textbox(label="Watermark Text")
                mark_btn = gr.Button("Add Watermark")
                msg_box = gr.Textbox(label="System Message")
                file_output = gr.File(label="Download Watermarked Image")
            with gr.Column():
                out_im = gr.Image(label="Watermarked Image")

        mark_btn.click(choose_encode, inputs=[inp_im, inp_mark, cho], outputs=[out_im, msg_box, file_output])

    with gr.Tab("Detect Watermark"):
        with gr.Row():
            with gr.Column():
                det_im = gr.Image(label="Watermarked Image", type="filepath")
                det_btn = gr.Button("Detect Watermark")
            with gr.Column():
                det_msg = gr.Textbox(label="Detected Watermark", lines=6)

        det_btn.click(detect_watermark, inputs=[det_im], outputs=[det_msg])

if __name__ == "__main__":
    app.launch()