Spaces:
Sleeping
Sleeping
from PIL import Image, ImageDraw, ImageFont | |
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_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_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: | |
PIL Image object with watermark applied | |
""" | |
try: | |
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) | |
try: | |
font_size = min(image.width, image.height) // 20 | |
font = ImageFont.truetype("arial.ttf", font_size) | |
except: | |
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] | |
text_height = bbox[3] - bbox[1] | |
x = (image.width - text_width) // 2 | |
y = (image.height - text_height) // 2 | |
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=shadow_color, font=font) | |
draw.text((x, y), watermark_text, fill=text_color, font=font) | |
watermarked = Image.alpha_composite(image, overlay) | |
final_image = watermarked.convert('RGB') | |
return final_image | |
except Exception as e: | |
print(f"Error adding watermark: {e}") | |
return image_input if isinstance(image_input, Image.Image) else None | |
def remove_watermark(image_input: Union[str, Image.Image], alpha: float = 2.0, beta: float = -160) -> Union[str, Image.Image]: | |
try: | |
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) | |
result_image = Image.fromarray(new) | |
return result_image | |
except Exception as e: | |
print(f"Error removing watermark: {e}") | |
return image_input if isinstance(image_input, Image.Image) else None |