remove prompt injection, change watermark file and added remove watermark, both added to mcp server
efbd483
| from PIL import Image, ImageDraw, ImageFont | |
| import os | |
| from typing import Dict, Any | |
| import cv2 | |
| import numpy as np | |
| def add_watermark(image_path: str, watermark_text: str) -> Dict[str, Any]: | |
| """ | |
| 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. | |
| 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. | |
| """ | |
| try: | |
| image = Image.open(image_path) | |
| 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: | |
| 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 | |
| text_color = (255, 255, 255, 128) | |
| draw.text((x-2, y-2), watermark_text, fill=(0, 0, 0, 64), font=font) | |
| draw.text((x, y), watermark_text, fill=text_color, font=font) | |
| watermarked = Image.alpha_composite(image.convert('RGBA'), 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 | |
| } | |
| except Exception as e: | |
| return { | |
| "success": False, | |
| "error": str(e), | |
| "input_path": image_path, | |
| "output_path": 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. | |
| """ | |
| try: | |
| img = cv2.imread(image_path) | |
| if img is None: | |
| raise ValueError("Could not load image") | |
| new = alpha * img + 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) | |
| 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 | |
| } | |
| except Exception as e: | |
| return { | |
| "success": False, | |
| "error": str(e), | |
| "input_path": image_path, | |
| "output_path": None | |
| } |