Update app.py

app.py

@@ -4,10 +4,13 @@ import utils
 import io
 import os
 import cv2
+import numpy as np
 
 def highlight_watermark(image, coords=(0, 0, 100, 50), color="red", width=5):
     """Highlights watermark area. Adapts rectangle size based on coordinates."""
     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)
@@ -28,24 +31,29 @@ def choose_encode(inp_im, inp_mark, cho):
             return None, "Invalid image format. Use PNG or JPG.", None
 
         if cho == "stegan":
-            image_pil = Image.open(inp_im)
-            image_cv2, msg = utils.encode(inp_im, inp_mark)
-            if msg:
-                return None, msg, None
-
-            image_pil = Image.fromarray(cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB))
+            # Call encode function with the image path and watermark text
+            watermarked_image, error_msg = utils.encode(inp_im, inp_mark)
+            
+            if error_msg:
+                return None, error_msg, None
+                
+            # Convert to PIL Image and highlight
+            image_pil = Image.fromarray(cv2.cvtColor(watermarked_image, cv2.COLOR_BGR2RGB))
             highlighted_image = highlight_watermark(image_pil.copy(), (0, 0, image_pil.width, image_pil.height), color="blue")
-
+            
+            # Prepare for download
             img_byte_arr = io.BytesIO()
             highlighted_image.save(img_byte_arr, format='PNG')
             img_byte_arr.seek(0)
             download_name = os.path.splitext(os.path.basename(inp_im))[0] + "_watermarked.png"
+            
             return highlighted_image, "Steganography watermark added successfully.", (download_name, img_byte_arr.read())
 
         elif cho == "pnginfo":
-            out_im_path, out_msg = utils.png_encode(inp_im, inp_mark)
-            if out_msg:  # Handle potential errors from png_encode
-                return None, out_msg, None
+            out_im_path, error_msg = utils.png_encode(inp_im, inp_mark)
+            if error_msg:
+                return None, error_msg, None
+                
             try:
                 img = Image.open(out_im_path)
                 coords = utils.png_encode_coords if hasattr(utils, 'png_encode_coords') else (50, 50, 200, 50)
@@ -60,15 +68,9 @@ def choose_encode(inp_im, inp_mark, cho):
             except Exception as e:
                 return None, f"Error processing PNG: {e}", None
 
-
-    except FileNotFoundError:
-        return None, "Image file not found.", None
-    except ValueError as e:
-        return None, f"Invalid input: {e}", None
     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."
@@ -77,34 +79,36 @@ def detect_watermark(det_im):
     if not any(det_im.lower().endswith("." + ext) for ext in valid_image_types):
         return "Invalid image format. Use PNG or JPG."
 
-    det_msg = utils.decode(det_im)
-
-    if det_msg is None or det_msg == "":
+    detected_text = utils.decode(det_im)
+    if not detected_text or detected_text.startswith("Error"):
         return "No watermark detected or not encoded with this tool."
 
-    return det_msg
+    return detected_text
 
+# Create Gradio interface
 with gr.Blocks() as app:
     with gr.Tab("Add Watermark"):
-        cho = gr.Radio(choices=["stegan", "pnginfo"], value="stegan")
+        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")
+                inp_mark = gr.Textbox(label="Watermark Text")
                 mark_btn = gr.Button("Add Watermark")
                 msg_box = gr.Textbox(label="System Message")
             with gr.Column():
                 out_im = gr.Image(label="Watermarked Image")
-                file_output = gr.File()
+                file_output = gr.File(label="Download Watermarked Image")
+    
     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")
+                det_btn = gr.Button("Detect Watermark")
             with gr.Column():
-                det_msg = gr.Textbox(label="Detected Watermark", lines=6, max_lines=50)
+                det_msg = gr.Textbox(label="Detected Watermark", lines=6)
 
-    mark_btn.click(choose_encode, [inp_im, inp_mark, cho], [out_im, msg_box, file_output])
-    det_btn.click(utils.decode, [det_im], det_msg)
+    mark_btn.click(choose_encode, inputs=[inp_im, inp_mark, cho], outputs=[out_im, msg_box, file_output])
+    det_btn.click(detect_watermark, inputs=[det_im], outputs=[det_msg])
 
-app.launch()
+if __name__ == "__main__":
+    app.launch()