import gradio as gr
import numpy as np
import cv2
from PIL import Image

# Load HED model
HED_NET = cv2.dnn.readNetFromCaffe("deploy.prototxt", "hed_pretrained_bsds.caffemodel")

# Edge detection
def detect_edges(image_np, method):
    if method == "HED":
        inp = cv2.dnn.blobFromImage(image_np, scalefactor=1.0, size=(256, 256),
                                    mean=(104.00698793, 116.66876762, 122.67891434),
                                    swapRB=False, crop=False)
        HED_NET.setInput(inp)
        edges = HED_NET.forward()[0, 0]
        edges = cv2.resize(edges, (image_np.shape[1], image_np.shape[0]))
        return edges > np.mean(edges)
    elif method == "Canny":
        gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
        edges = cv2.Canny(gray, 100, 200)
        return edges > 0
    else:
        raise ValueError("Invalid edge method")

# LSB extraction
def extract_lsb(stego_pixels, edge_mask, secret_bits_length, x=2, y=1):
    height, width, channels = stego_pixels.shape
    extracted_bits = []
    bit_index = 0

    for i in range(height):
        for j in range(width):
            for c in range(channels):
                if bit_index >= secret_bits_length:
                    return np.array(extracted_bits, dtype=np.uint8)
                n = x if edge_mask[i, j] else y
                if n <= 0:
                    continue
                value = stego_pixels[i, j, c] & ((1 << n) - 1)
                for b in range(n):
                    extracted_bits.append((value >> b) & 1)
                bit_index += n

    return np.array(extracted_bits[:secret_bits_length], dtype=np.uint8)


# Convert bits to image
def from_bit_array(bit_array, shape):
    expected_bits = np.prod(shape) * 8
    bit_array = bit_array[:expected_bits]
    return np.packbits(bit_array).reshape(shape)

# Gradio interface function
def extract_interface(stego_img, edge_method):
    secret_shape = (128, 128)  # Must match embedding shape
    stego_np = np.array(stego_img.convert("RGB"))

    edge_mask = detect_edges(stego_np, edge_method)
    secret_bits_len = np.prod(secret_shape) * 8

    extracted_bits = extract_lsb(stego_np, edge_mask, secret_bits_len, x=2, y=1)
    recovered = from_bit_array(extracted_bits, secret_shape)

    return Image.fromarray(recovered.astype(np.uint8), mode='L')

# Launch app
gr.Interface(
    fn=extract_interface,
    inputs=[
        gr.Image(label="Stego Image (RGB)", type="pil"),
        gr.Dropdown(["HED", "Canny"], label="Edge Detection Used", value="HED")
    ],
    outputs=gr.Image(label="Extracted Palmprint (128×128 grayscale)"),
    title="Palmprint Extractor from Stego Image",
    description="Upload the stego image and choose the same edge detection method used in embedding. The system will extract a 128×128 grayscale palmprint."
).launch()