fix

app.py

@@ -3,56 +3,67 @@ import numpy as np
 import cv2
 from PIL import Image
 
-# Load HED model (must be uploaded in the repo)
+# Load HED model
 HED_NET = cv2.dnn.readNetFromCaffe("deploy.prototxt", "hed_pretrained_bsds.caffemodel")
 
-# Main extraction logic
-def extract_secret_from_stego(stego_image: Image.Image, secret_shape=(128, 128), edge_method="HED", x=2, y=1):
-    stego_np = np.array(stego_image.convert("RGB"))
-
-    # Edge detection
-    if edge_method == "HED":
-        inp = cv2.dnn.blobFromImage(stego_np, scalefactor=1.0, size=(256, 256),
+# 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]
-        edge_mask = cv2.resize(edges, (stego_np.shape[1], stego_np.shape[0])) > np.mean(edges)
-    elif edge_method == "Canny":
-        gray = cv2.cvtColor(stego_np, cv2.COLOR_RGB2GRAY)
-        edge_mask = cv2.Canny(gray, 100, 200) > 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("Unsupported edge detection method")
+        raise ValueError("Invalid edge method")
 
-    # Bit extraction
-    secret_bits_len = np.prod(secret_shape) * 8
-    height, width, channels = stego_np.shape
+# 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_len:
+                if bit_index >= secret_bits_length:
                     break
                 n = x if edge_mask[i, j] else y
-                value = stego_np[i, j, c] & ((1 << n) - 1)
+                value = stego_pixels[i, j, c] & ((1 << n) - 1)
                 for b in range(n):
                     extracted_bits.append((value >> b) & 1)
                 bit_index += n
-            if bit_index >= secret_bits_len:
+            if bit_index >= secret_bits_length:
                 break
-        if bit_index >= secret_bits_len:
+        if bit_index >= secret_bits_length:
             break
 
-    extracted_bits = np.array(extracted_bits[:secret_bits_len], dtype=np.uint8)
-    recovered_array = np.packbits(extracted_bits).reshape(secret_shape)
+    return np.array(extracted_bits[:secret_bits_length], dtype=np.uint8)
 
-    return Image.fromarray(recovered_array.astype(np.uint8), mode='L')
+# 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):
-    return extract_secret_from_stego(stego_img, secret_shape=(128, 128), edge_method=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(
@@ -63,5 +74,5 @@ gr.Interface(
     ],
     outputs=gr.Image(label="Extracted Palmprint (128×128 grayscale)"),
     title="Palmprint Extractor from Stego Image",
-    description="Upload a stego image and select the edge detection method used during embedding. This will extract a 128×128 grayscale palmprint."
-).launch()
+    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()