Update app.py

app.py

@@ -5,9 +5,9 @@ import io
 import os
 import cv2
 import numpy as np
+import tempfile
 
 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))
@@ -19,6 +19,26 @@ def highlight_watermark(image, coords=(0, 0, 100, 50), color="red", width=5):
         print(f"Error highlighting: {e}")
         return image
 
+def create_output_file(image, original_filename):
+    try:
+        # Create temp directory if it doesn't exist
+        temp_dir = "temp"
+        os.makedirs(temp_dir, exist_ok=True)
+        
+        # Generate output filename
+        base_name = os.path.splitext(os.path.basename(original_filename))[0]
+        output_path = os.path.join(temp_dir, f"{base_name}_watermarked.png")
+        
+        # Save image
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+        image.save(output_path, format='PNG')
+        
+        return output_path
+    except Exception as e:
+        print(f"Error creating output file: {e}")
+        return None
+
 def choose_encode(inp_im, inp_mark, cho):
     try:
         if not inp_im:
@@ -31,23 +51,19 @@ 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"
+            output_path = create_output_file(highlighted_image, inp_im)
+            if not output_path:
+                return None, "Error saving watermarked image.", None
             
-            return highlighted_image, "Steganography watermark added successfully.", (download_name, img_byte_arr.read())
+            return highlighted_image, "Steganography watermark added successfully.", output_path
 
         elif cho == "pnginfo":
             out_im_path, error_msg = utils.png_encode(inp_im, inp_mark)
@@ -58,13 +74,12 @@ def choose_encode(inp_im, inp_mark, cho):
                 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())
+                
+                output_path = create_output_file(highlighted_image, inp_im)
+                if not output_path:
+                    return None, "Error saving watermarked image.", None
+                
+                return highlighted_image, "PNG metadata watermark added successfully.", output_path
             except Exception as e:
                 return None, f"Error processing PNG: {e}", None
 
@@ -111,4 +126,4 @@ with gr.Blocks() as app:
     det_btn.click(detect_watermark, inputs=[det_im], outputs=[det_msg])
 
 if __name__ == "__main__":
-    app.launch()
+    app.launch()