Update app.py

app.py

@@ -1,32 +1,110 @@
 import gradio as gr
+from PIL import Image, ImageDraw
 import utils
+import io
+import os
+import cv2
 
-def choose_encode(inp_im,inp_mark,cho):
-    if cho == "stegan":
-        out_im, out_msg = utils.encode(inp_im,inp_mark)
-        return out_im,out_msg
-    if cho == "pnginfo":
-        out_im, out_msg = utils.png_encode(inp_im,inp_mark)
-        return out_im,out_msg
+def highlight_watermark(image, coords=(0, 0, 100, 50), color="red", width=5):
+    """Highlights watermark area. Adapts rectangle size based on coordinates."""
+    try:
+        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
+
+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
+
+        valid_image_types = ["png", "jpg", "jpeg", "gif", "bmp", "tiff", "webp"]
+        if not any(inp_im.lower().endswith("." + ext) for ext in valid_image_types):
+            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))
+            highlighted_image = highlight_watermark(image_pil.copy(), (0, 0, image_pil.width, image_pil.height), color="blue")
+
+            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
+            try:
+                img = Image.open(out_im_path)
+                coords = utils.png_encode_coords if hasattr(utils, 'png_encode_coords') else (50, 50, 200, 50)
+                highlighted_image = highlight_watermark(img.copy(), coords)
+
+                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, "PNG metadata watermark added successfully.", (download_name, img_byte_arr.read())
+            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."
+
+    valid_image_types = ["png", "jpg", "jpeg", "gif", "bmp", "tiff", "webp"]
+    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 == "":
+        return "No watermark detected or not encoded with this tool."
+
+    return det_msg
 
 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")
         with gr.Row():
             with gr.Column():
-                inp_im = gr.Image(label="Input Image",type="filepath")
+                inp_im = gr.Image(label="Input Image", type="filepath")
                 inp_mark = gr.Textbox(label="Watermark")
-                mark_btn = gr.Button()
+                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()
     with gr.Tab("Detect Watermark"):
         with gr.Row():
             with gr.Column():
-                det_im = gr.Image(label="Watermarked Image",type="filepath")
+                det_im = gr.Image(label="Watermarked Image", type="filepath")
                 det_btn = gr.Button("Detect")
             with gr.Column():
-                det_msg = gr.Textbox(label="Detected Watermark",lines=6, max_lines=50)
-    mark_btn.click(choose_encode,[inp_im,inp_mark,cho],[out_im,msg_box])
-    det_btn.click(utils.decode,[det_im],det_msg)
+                det_msg = gr.Textbox(label="Detected Watermark", lines=6, max_lines=50)
+
+    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)
+
 app.launch()