Update app.py

app.py

@@ -1,45 +1,139 @@
+# app.py
+
 import gradio as gr
-import utils
-
-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
-
-css = """
-footer {
-    visibility: hidden;
-}
-"""
-
-with gr.Blocks(css=css) as app:
-    gr.HTML("""<a href="https://visitorbadge.io/status?path=https%3A%2F%2Ffantos-watermark3.hf.space">
-               <img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Ffantos-watermark3.hf.space&countColor=%23263759" />
-               </a>""")
-    
-    with gr.Tab("Add Watermark"):
-        cho = gr.Radio(choices=["stegan", "pnginfo"], value="stegan")
-        with gr.Row():
-            with gr.Column():
-                inp_im = gr.Image(label="Input Image", type="filepath")
-                inp_mark = gr.Textbox(label="Watermark")
-                mark_btn = gr.Button("Add Watermark")
-                msg_box = gr.Textbox(label="System Message")
-            with gr.Column():
-                out_im = gr.Image(label="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")
-            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)
-
-app.launch()
+from utils import WatermarkProcessor
+import json
+import tempfile
+import os
+from datetime import datetime
+import cv2
+from PIL import Image
+import numpy as np
+
+class WatermarkGUI:
+    def __init__(self):
+        self.processor = WatermarkProcessor()
+        self.create_interface()
+
+    def process_watermark(self, image, watermark_text, author, purpose, opacity):
+        """Process watermark with metadata"""
+        if image is None or watermark_text.strip() == "":
+            return None, "Please provide both image and watermark text"
+            
+        metadata = {
+            "author": author,
+            "purpose": purpose,
+            "opacity": opacity
+        }
+        
+        # Save temporary image
+        temp_path = tempfile.mktemp(suffix='.png')
+        Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).save(temp_path)
+        
+        # Add watermark
+        result_path, message = self.processor.encode(temp_path, watermark_text, metadata)
+        
+        if "Error" in message:
+            return None, message
+            
+        # Generate quality report
+        quality_report = self.processor.analyze_quality(temp_path, result_path)
+        quality_data = json.loads(quality_report)
+        
+        # Create formatted report
+        report = f"""
+### Watermark Quality Report
+- Quality Score: {quality_data['quality_score']}%
+- PSNR: {quality_data['psnr']} dB
+- Histogram Similarity: {quality_data['histogram_similarity'] * 100:.2f}%
+- Modified Pixels: {quality_data['modified_pixels']:,}
+
+### Metadata
+- Author: {author}
+- Purpose: {purpose}
+- Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
+        """
+        
+        os.remove(temp_path)
+        return cv2.imread(result_path), report
+
+    def detect_watermark(self, image):
+        """Detect and extract watermark"""
+        if image is None:
+            return "Please provide an image"
+            
+        # Save temporary image
+        temp_path = tempfile.mktemp(suffix='.png')
+        Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).save(temp_path)
+        
+        # Extract watermark
+        result = self.processor.decode(temp_path)
+        os.remove(temp_path)
+        
+        try:
+            # Parse JSON result
+            data = json.loads(result)
+            report = f"""
+### Extracted Watermark
+Text: {data['text']}
+
+### Metadata
+- Timestamp: {data['timestamp']}
+- Author: {data['metadata'].get('author', 'N/A')}
+- Purpose: {data['metadata'].get('purpose', 'N/A')}
+            """
+            return report
+        except:
+            return result
+
+    def create_interface(self):
+        """Create Gradio interface"""
+        with gr.Blocks(css="footer {visibility: hidden}") as self.interface:
+            gr.Markdown("# Enhanced Image Watermarking System")
+            
+            with gr.Tabs():
+                # Add Watermark Tab
+                with gr.Tab("Add Watermark"):
+                    with gr.Row():
+                        with gr.Column():
+                            input_image = gr.Image(label="Input Image", type="numpy")
+                            watermark_text = gr.Textbox(label="Watermark Text")
+                            author = gr.Textbox(label="Author", placeholder="Enter author name")
+                            purpose = gr.Textbox(label="Purpose", placeholder="Enter watermark purpose")
+                            opacity = gr.Slider(minimum=0.1, maximum=1.0, value=0.3, 
+                                             label="Watermark Opacity")
+                            
+                            with gr.Row():
+                                process_btn = gr.Button("Add Watermark", variant="primary")
+                                
+                        with gr.Column():
+                            result_image = gr.Image(label="Watermarked Image")
+                            quality_report = gr.Markdown(label="Quality Report")
+                
+                # Detect Watermark Tab
+                with gr.Tab("Detect Watermark"):
+                    with gr.Row():
+                        detect_image = gr.Image(label="Input Image", type="numpy")
+                        detect_result = gr.Markdown(label="Detected Watermark")
+                    detect_btn = gr.Button("Detect Watermark")
+            
+            # Event handlers            
+            process_btn.click(
+                fn=self.process_watermark,
+                inputs=[input_image, watermark_text, author, purpose, opacity],
+                outputs=[result_image, quality_report]
+            )
+            
+            detect_btn.click(
+                fn=self.detect_watermark,
+                inputs=[detect_image],
+                outputs=detect_result
+            )
+
+    def launch(self, *args, **kwargs):
+        """Launch the interface"""
+        self.interface.launch(*args, **kwargs)
+
+if __name__ == "__main__":
+    app = WatermarkGUI()
+    app.launch()