Update app.py

app.py

@@ -7,6 +7,7 @@ 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))
@@ -30,25 +31,29 @@ def choose_encode(inp_im, inp_mark, cho):
             return None, "Invalid image format. Use PNG or JPG.", None
 
         if cho == "stegan":
+            # 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.getvalue())
+            return highlighted_image, "Steganography watermark added successfully.", (download_name, img_byte_arr.read())
 
         elif cho == "pnginfo":
             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)
@@ -59,7 +64,7 @@ def choose_encode(inp_im, inp_mark, cho):
                 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.getvalue())
+                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
 
@@ -80,6 +85,7 @@ def detect_watermark(det_im):
 
     return detected_text
 
+# Create Gradio interface
 with gr.Blocks() as app:
     with gr.Tab("Add Watermark"):
         cho = gr.Radio(choices=["stegan", "pnginfo"], value="stegan", label="Watermark Method")
@@ -105,4 +111,4 @@ with gr.Blocks() as app:
     det_btn.click(detect_watermark, inputs=[det_im], outputs=[det_msg])
 
 if __name__ == "__main__":
-    app.launch()
+    app.launch()