Added watermark

app.py

@@ -46,30 +46,24 @@ def base64_to_image(base64_str):
     if not base64_str:
         return None
     
-    # Remove data URI prefix if present (e.g., "data:image/png;base64,")
     if isinstance(base64_str, str) and "base64," in base64_str:
         base64_str = base64_str.split("base64,", 1)[1]
     
     try:
-        # Strip any whitespace that might be in the base64 string
         if isinstance(base64_str, str):
             base64_str = base64_str.strip()
         
-        # Decode the base64 data
         image_data = base64.b64decode(base64_str)
         
-        # Check if we have data
         if not image_data:
             print("Decoded base64 data is empty")
             return None
         
-        # Attempt to open the image
         image = Image.open(io.BytesIO(image_data))
         
-        # Convert the image to ensure it's valid
         return image.copy()
         
-    except base64.binascii.Error as e:
+    except binascii.Error as e:
         print(f"Base64 decoding error: {str(e)}")
         if isinstance(base64_str, str):
             preview = base64_str[:30] + "..." if len(base64_str) > 30 else base64_str
@@ -79,12 +73,10 @@ def base64_to_image(base64_str):
     except Exception as e:
         print(f"Error converting base64 to image: {str(e)}")
         
-        # Print preview of the base64 string for debugging
         if isinstance(base64_str, str):
             preview = base64_str[:30] + "..." if len(base64_str) > 30 else base64_str
             print(f"Base64 preview: {preview}")
         
-        # Additional debug information
         if 'image_data' in locals() and image_data:
             try:
                 magic_bytes = image_data[:12].hex()
@@ -99,21 +91,28 @@ def url_to_base64(url):
     return base64.b64encode(response.content).decode()
 
 def gradio_remove_background(image):
+    """
+    Attempt to remove watermarks from an image using contrast and brightness adjustment.
+    
+    Args:
+        image_input: PIL Image object or base64 string
+        alpha: Contrast control (1.0-3.0, default 2.0). Higher values increase contrast.
+        beta: Brightness control (-255 to 255, default -160). Negative values decrease brightness.
+        
+    Returns:
+        PIL Image object with watermark removal attempted
+    """
     if image is None:
         return None
     base64_img = image_to_base64(image)
     result = remove_background(f"data:image/png;base64,{base64_img}")
     
-    # Check if the result is directly a base64 string or has an image_data key
     if isinstance(result, str):
         return base64_to_image(result)
     elif isinstance(result, dict) and "image_data" in result:
-        # If image_data contains a data URI prefix
         if isinstance(result["image_data"], str) and result["image_data"].startswith("data:"):
-            # The response already contains the full data URI
             return base64_to_image(result["image_data"])
         else:
-            # Try to process it as a regular base64 string
             try:
                 return base64_to_image(result["image_data"])
             except Exception as e:
@@ -175,14 +174,19 @@ def toggle_intensity_slider(filter_type):
     return gr.Slider(interactive=filter_type in intensity_filters)
 
 def gradio_add_watermark(image, watermark_text, opacity=0.5):
+    """
+    Gradio wrapper for add_watermark function
+    """
     if image is None:
         return None
+    if not watermark_text or watermark_text.strip() == "":
+        return image
+    
     try:
-        base64_img = image_to_base64(image)
-        result = add_watermark(base64_img, watermark_text, opacity)
-        return base64_to_image(result)
+        result = add_watermark(image, watermark_text, opacity)
+        return result
     except Exception as e:
-        print(f"Error adding watermark: {e}")
+        print(f"Error in gradio_add_watermark: {e}")
         return image
 
 def gradio_remove_watermark(image):

src/utils/watermark.py

@@ -3,22 +3,36 @@ import os
 from typing import Dict, Any
 import cv2
 import numpy as np
+import binascii
+from typing import Union
+import base64
+import io
 
-def add_watermark(image_path: str, watermark_text: str) -> Dict[str, Any]:
+def add_watermark(image_input: Union[str, Image.Image], watermark_text: str, opacity: float = 0.5) -> Union[str, Image.Image]:
     """
     Add a semi-transparent text watermark to an image.
     
     Args:
-        image_path: The path to the input image file.
-        watermark_text: The text to be used as watermark.
+        image_input: PIL Image object or base64 string
+        watermark_text: The text to be used as watermark
+        opacity: Opacity of the watermark (0.1 to 1.0)
         
     Returns:
-        A dictionary containing success status, file paths, and operation details.
-        On success: success=True, input_path, output_path, output_size_bytes, watermark_text, message.
-        On failure: success=False, error message, input_path, output_path=None.
+        PIL Image object with watermark applied
     """
     try:
-        image = Image.open(image_path)
+        if isinstance(image_input, str):
+            if image_input.startswith('data:image'):
+                base64_data = image_input.split(',')[1]
+            else:
+                base64_data = image_input
+            image_data = base64.b64decode(base64_data)
+            image = Image.open(io.BytesIO(image_data))
+        else:
+            image = image_input
+        
+        if image.mode != 'RGBA':
+            image = image.convert('RGBA')
         
         overlay = Image.new('RGBA', image.size, (255, 255, 255, 0))
         draw = ImageDraw.Draw(overlay)
@@ -27,7 +41,10 @@ def add_watermark(image_path: str, watermark_text: str) -> Dict[str, Any]:
             font_size = min(image.width, image.height) // 20
             font = ImageFont.truetype("arial.ttf", font_size)
         except:
-            font = ImageFont.load_default()
+            try:
+                font = ImageFont.load_default()
+            except:
+                font = ImageFont.load_default()
         
         bbox = draw.textbbox((0, 0), watermark_text, font=font)
         text_width = bbox[2] - bbox[0]
@@ -36,84 +53,43 @@ def add_watermark(image_path: str, watermark_text: str) -> Dict[str, Any]:
         x = (image.width - text_width) // 2
         y = (image.height - text_height) // 2
         
-        text_color = (255, 255, 255, 128)
+        alpha_value = int(255 * opacity)
+        text_color = (255, 255, 255, alpha_value)
+        shadow_color = (0, 0, 0, int(alpha_value * 0.5))
         
-        draw.text((x-2, y-2), watermark_text, fill=(0, 0, 0, 64), font=font)
+        draw.text((x-2, y-2), watermark_text, fill=shadow_color, font=font)
         draw.text((x, y), watermark_text, fill=text_color, font=font)
         
-        watermarked = Image.alpha_composite(image.convert('RGBA'), overlay)
+        watermarked = Image.alpha_composite(image, overlay)
         final_image = watermarked.convert('RGB')
         
-        base_dir = os.path.dirname(image_path)
-        base_name, ext = os.path.splitext(os.path.basename(image_path))
-        new_filename = f"{base_name}_watermarked{ext}"
-        new_path = os.path.join(base_dir, new_filename)
-        
-        final_image.save(new_path, quality=95)
-        output_size = os.path.getsize(new_path)
-        
-        return {
-            "success": True,
-            "message": "Watermark added successfully",
-            "input_path": image_path,
-            "output_path": new_path,
-            "output_size_bytes": output_size,
-            "watermark_text": watermark_text
-        }
+        return final_image
         
     except Exception as e:
-        return {
-            "success": False,
-            "error": str(e),
-            "input_path": image_path,
-            "output_path": None
-        }
+        print(f"Error adding watermark: {e}")
+        return image_input if isinstance(image_input, Image.Image) else None
 
-def remove_watermark(image_path: str, alpha: float = 2.0, beta: float = -160) -> Dict[str, Any]:
-    """
-    Attempt to remove watermarks from an image using contrast and brightness adjustment.
-    
-    Args:
-        image_path: The path to the input image file.
-        alpha: Contrast control (1.0-3.0, default 2.0). Higher values increase contrast.
-        beta: Brightness control (-255 to 255, default -160). Negative values decrease brightness.
-        
-    Returns:
-        A dictionary containing success status, file paths, and operation details.
-        On success: success=True, input_path, output_path, output_size_bytes, alpha, beta, message.
-        On failure: success=False, error message, input_path, output_path=None.
-    """
+def remove_watermark(image_input: Union[str, Image.Image], alpha: float = 2.0, beta: float = -160) -> Union[str, Image.Image]:
     try:
-        img = cv2.imread(image_path)
-        
-        if img is None:
-            raise ValueError("Could not load image")
-        
-        new = alpha * img + beta
+        if isinstance(image_input, str):
+            if image_input.startswith('data:image'):
+                base64_data = image_input.split(',')[1]
+            else:
+                base64_data = image_input
+            image_data = base64.b64decode(base64_data)
+            image = Image.open(io.BytesIO(image_data))
+        else:
+            image = image_input
+        
+        img_array = np.array(image)
+        
+        new = alpha * img_array + beta
         new = np.clip(new, 0, 255).astype(np.uint8)
         
-        base_dir = os.path.dirname(image_path)
-        base_name, ext = os.path.splitext(os.path.basename(image_path))
-        new_filename = f"{base_name}_cleaned{ext}"
-        new_path = os.path.join(base_dir, new_filename)
-        
-        cv2.imwrite(new_path, new)
-        output_size = os.path.getsize(new_path)
+        result_image = Image.fromarray(new)
         
-        return {
-            "success": True,
-            "message": "Watermark removal attempted successfully",
-            "input_path": image_path,
-            "output_path": new_path,
-            "output_size_bytes": output_size,
-            "alpha": alpha,
-            "beta": beta
-        }
+        return result_image
         
     except Exception as e:
-        return {
-            "success": False,
-            "error": str(e),
-            "input_path": image_path,
-            "output_path": None
-        }
+        print(f"Error removing watermark: {e}")
+        return image_input if isinstance(image_input, Image.Image) else None