pip install numpy pillow

GWMRTool.py

@@ -0,0 +1,881 @@
+"""
+Gemini Watermark Removal Tool
+Removes visible Gemini watermarks using reverse alpha blending.
+
+Based on the mathematical formula:
+    watermarked = α × logo + (1 - α) × original
+    
+Solving for original:
+    original = (watermarked - α × logo) / (1 - α)
+Python port with embedded alpha maps.
+
+Usage:
+    python GWMRTool.py image.jpg
+    python GWMRTool.py -i input.jpg -o output.jpg
+    python GWMRTool.py -i ./input_folder/ -o ./output_folder/
+"""
+
+import numpy as np
+from PIL import Image
+import argparse
+from pathlib import Path
+import sys
+import os
+from typing import Tuple, Optional
+
+
+# ============================================================================
+# EMBEDDED ALPHA MAPS
+# ============================================================================
+
+# Default paths for alpha maps (relative to script location)
+SCRIPT_DIR = Path(__file__).parent.resolve()
+DEFAULT_ALPHA_MAP_SMALL_PATH = SCRIPT_DIR / "bg_48.png"
+DEFAULT_ALPHA_MAP_LARGE_PATH = SCRIPT_DIR / "bg_96.png"
+
+
+class AlphaMapManager:
+    """
+    Manages loading and caching of alpha maps.
+    """
+    
+    _cache = {}
+    
+    @classmethod
+    def load_alpha_map(cls, path: Path, expected_size: int) -> np.ndarray:
+        """
+        Load alpha map from file with caching.
+        
+        Args:
+            path: Path to alpha map image
+            expected_size: Expected size (48 or 96)
+            
+        Returns:
+            Alpha map as float32 array [0, 1]
+        """
+        cache_key = str(path)
+        
+        if cache_key in cls._cache:
+            return cls._cache[cache_key]
+        
+        if not path.exists():
+            raise FileNotFoundError(f"Alpha map not found: {path}")
+        
+        # Load image
+        img = Image.open(path)
+        img_array = np.array(img)
+        
+        # Calculate alpha map using max of RGB channels
+        if len(img_array.shape) == 3:
+            if img_array.shape[2] >= 3:
+                # Use max of RGB channels for brightness
+                gray = np.max(img_array[:, :, :3], axis=2)
+            else:
+                gray = img_array[:, :, 0]
+        else:
+            gray = img_array
+        
+        # Normalize to [0, 1]
+        alpha_map = gray.astype(np.float32) / 255.0
+        
+        # Validate size
+        if alpha_map.shape[0] != expected_size or alpha_map.shape[1] != expected_size:
+            print(f"Warning: Alpha map size {alpha_map.shape} doesn't match expected {expected_size}x{expected_size}")
+            # Resize if needed
+            img_resized = Image.fromarray(gray).resize(
+                (expected_size, expected_size), 
+                Image.Resampling.LANCZOS
+            )
+            alpha_map = np.array(img_resized).astype(np.float32) / 255.0
+        
+        # Cache and return
+        cls._cache[cache_key] = alpha_map
+        return alpha_map
+    
+    @classmethod
+    def clear_cache(cls):
+        """Clear the alpha map cache."""
+        cls._cache.clear()
+
+
+# ============================================================================
+# BLEND MODES (Translated from C++ blend_modes.cpp)
+# ============================================================================
+
+class BlendModes:
+    """
+    Alpha blending operations for watermark removal/addition.
+    Direct translation from C++ blend_modes.cpp
+    """
+    
+    # Constants from C++ code
+    ALPHA_THRESHOLD = 0.002  # Ignore very small alpha (noise)
+    MAX_ALPHA = 0.99         # Avoid division by near-zero
+    
+    @staticmethod
+    def calculate_alpha_map(bg_capture: np.ndarray) -> np.ndarray:
+        """
+        Calculate alpha map from background capture.
+        Uses max of RGB channels for brightness (handles colored logos).
+        
+        Args:
+            bg_capture: Background capture image (grayscale or RGB/RGBA)
+            
+        Returns:
+            Alpha map as float32 array normalized to [0, 1]
+        """
+        if bg_capture is None or bg_capture.size == 0:
+            raise ValueError("Empty background capture")
+        
+        # Handle different channel configurations
+        if len(bg_capture.shape) == 3:
+            if bg_capture.shape[2] >= 3:
+                # RGB/RGBA: Use max of RGB channels for brightness
+                gray = np.max(bg_capture[:, :, :3], axis=2)
+            else:
+                gray = bg_capture[:, :, 0]
+        else:
+            # Already grayscale
+            gray = bg_capture.copy()
+        
+        # Convert to float and normalize to [0, 1]
+        alpha_map = gray.astype(np.float32) / 255.0
+        
+        return alpha_map
+    
+    @staticmethod
+    def remove_watermark_alpha_blend(
+        image: np.ndarray,
+        alpha_map: np.ndarray,
+        position: Tuple[int, int],
+        logo_value: float = 255.0
+    ) -> np.ndarray:
+        """
+        Remove watermark using reverse alpha blending (vectorized).
+        
+        Gemini applies: watermarked = alpha * logo + (1 - alpha) * original
+        We reverse:     original = (watermarked - alpha * logo) / (1 - alpha)
+        
+        Args:
+            image: Input image (RGB, uint8)
+            alpha_map: Alpha map (float32, normalized to [0, 1])
+            position: (x, y) position of watermark top-left corner
+            logo_value: Logo pixel value (default 255 for white logo)
+            
+        Returns:
+            Image with watermark removed
+        """
+        if image is None or image.size == 0:
+            raise ValueError("Empty image")
+        if alpha_map is None or alpha_map.size == 0:
+            raise ValueError("Empty alpha map")
+        if len(image.shape) != 3 or image.shape[2] != 3:
+            raise ValueError("Image must be RGB (3 channels)")
+        if alpha_map.dtype != np.float32:
+            alpha_map = alpha_map.astype(np.float32)
+        
+        # Make a copy to avoid modifying original
+        result = image.copy()
+        
+        # Calculate region of interest
+        x, y = position
+        h, w = alpha_map.shape[:2]
+        
+        # Clip to image bounds
+        x1 = max(0, x)
+        y1 = max(0, y)
+        x2 = min(image.shape[1], x + w)
+        y2 = min(image.shape[0], y + h)
+        
+        if x1 >= x2 or y1 >= y2:
+            return result
+        
+        # Calculate ROI offsets
+        alpha_x1 = x1 - x
+        alpha_y1 = y1 - y
+        alpha_x2 = alpha_x1 + (x2 - x1)
+        alpha_y2 = alpha_y1 + (y2 - y1)
+        
+        # Get ROIs
+        alpha_region = alpha_map[alpha_y1:alpha_y2, alpha_x1:alpha_x2]
+        image_region = result[y1:y2, x1:x2]
+        
+        # Convert image region to float
+        image_f = image_region.astype(np.float32)
+        
+        # Create mask for pixels to process (alpha >= threshold)
+        process_mask = alpha_region >= BlendModes.ALPHA_THRESHOLD
+        
+        # Clamp alpha values
+        alpha_clamped = np.minimum(alpha_region, BlendModes.MAX_ALPHA)
+        one_minus_alpha = 1.0 - alpha_clamped
+        
+        # Avoid division by zero
+        one_minus_alpha = np.maximum(one_minus_alpha, 1e-6)
+        
+        # Expand alpha to 3 channels for broadcasting
+        alpha_3c = alpha_clamped[:, :, np.newaxis]
+        one_minus_alpha_3c = one_minus_alpha[:, :, np.newaxis]
+        process_mask_3c = process_mask[:, :, np.newaxis]
+        
+        # Reverse alpha blending formula:
+        # original = (watermarked - alpha * logo) / (1 - alpha)
+        original = (image_f - alpha_3c * logo_value) / one_minus_alpha_3c
+        original = np.clip(original, 0.0, 255.0)
+        
+        # Only apply to pixels that need processing
+        image_f = np.where(process_mask_3c, original, image_f)
+        
+        # Convert back to 8-bit
+        result[y1:y2, x1:x2] = image_f.astype(np.uint8)
+        
+        return result
+    
+    @staticmethod
+    def add_watermark_alpha_blend(
+        image: np.ndarray,
+        alpha_map: np.ndarray,
+        position: Tuple[int, int],
+        logo_value: float = 255.0
+    ) -> np.ndarray:
+        """
+        Add watermark using alpha blending (same as Gemini).
+        
+        Formula: result = alpha * logo + (1 - alpha) * original
+        
+        Args:
+            image: Input image (RGB, uint8)
+            alpha_map: Alpha map (float32, normalized to [0, 1])
+            position: (x, y) position of watermark top-left corner
+            logo_value: Logo pixel value (default 255 for white logo)
+            
+        Returns:
+            Image with watermark added
+        """
+        if image is None or image.size == 0:
+            raise ValueError("Empty image")
+        if alpha_map is None or alpha_map.size == 0:
+            raise ValueError("Empty alpha map")
+        if len(image.shape) != 3 or image.shape[2] != 3:
+            raise ValueError("Image must be RGB (3 channels)")
+        if alpha_map.dtype != np.float32:
+            alpha_map = alpha_map.astype(np.float32)
+        
+        # Make a copy to avoid modifying original
+        result = image.copy()
+        
+        # Calculate region of interest
+        x, y = position
+        h, w = alpha_map.shape[:2]
+        
+        # Clip to image bounds
+        x1 = max(0, x)
+        y1 = max(0, y)
+        x2 = min(image.shape[1], x + w)
+        y2 = min(image.shape[0], y + h)
+        
+        if x1 >= x2 or y1 >= y2:
+            return result
+        
+        # Calculate ROI offsets
+        alpha_x1 = x1 - x
+        alpha_y1 = y1 - y
+        alpha_x2 = alpha_x1 + (x2 - x1)
+        alpha_y2 = alpha_y1 + (y2 - y1)
+        
+        # Get ROIs
+        alpha_region = alpha_map[alpha_y1:alpha_y2, alpha_x1:alpha_x2]
+        image_region = result[y1:y2, x1:x2]
+        
+        # Convert image region to float
+        image_f = image_region.astype(np.float32)
+        
+        # Create mask for pixels to process
+        process_mask = alpha_region >= BlendModes.ALPHA_THRESHOLD
+        
+        # Expand alpha to 3 channels for broadcasting
+        alpha_3c = alpha_region[:, :, np.newaxis]
+        one_minus_alpha_3c = (1.0 - alpha_region)[:, :, np.newaxis]
+        process_mask_3c = process_mask[:, :, np.newaxis]
+        
+        # Alpha blending formula:
+        # result = alpha * logo + (1 - alpha) * original
+        blended = alpha_3c * logo_value + one_minus_alpha_3c * image_f
+        blended = np.clip(blended, 0.0, 255.0)
+        
+        # Only apply to pixels that need processing
+        image_f = np.where(process_mask_3c, blended, image_f)
+        
+        # Convert back to 8-bit
+        result[y1:y2, x1:x2] = image_f.astype(np.uint8)
+        
+        return result
+
+
+# ============================================================================
+# WATERMARK REMOVER
+# ============================================================================
+
+class GeminiWatermarkRemover:
+    """
+    Removes Gemini watermarks using reverse alpha blending.
+    Uses bg_48.png and bg_96.png alpha maps by default.
+    """
+    
+    # Watermark size configurations
+    SMALL_SIZE = 48
+    LARGE_SIZE = 96
+    SMALL_MARGIN = 32
+    LARGE_MARGIN = 64
+    SIZE_THRESHOLD = 1024
+    
+    # Default logo value (white)
+    DEFAULT_LOGO_VALUE = 255.0
+    
+    def __init__(self, 
+                 alpha_map_small_path: Path = None,
+                 alpha_map_large_path: Path = None,
+                 logo_value: float = DEFAULT_LOGO_VALUE):
+        """
+        Initialize the watermark remover.
+        
+        Args:
+            alpha_map_small_path: Path to 48x48 alpha map (default: bg_48.png)
+            alpha_map_large_path: Path to 96x96 alpha map (default: bg_96.png)
+            logo_value: Logo pixel value (default 255 for white)
+        """
+        # Use default paths if not specified
+        self.alpha_map_small_path = alpha_map_small_path or DEFAULT_ALPHA_MAP_SMALL_PATH
+        self.alpha_map_large_path = alpha_map_large_path or DEFAULT_ALPHA_MAP_LARGE_PATH
+        self.logo_value = logo_value
+        
+        # Lazy-loaded alpha maps
+        self._alpha_map_small = None
+        self._alpha_map_large = None
+    
+    @property
+    def alpha_map_small(self) -> np.ndarray:
+        """Lazy-load small alpha map."""
+        if self._alpha_map_small is None:
+            self._alpha_map_small = AlphaMapManager.load_alpha_map(
+                Path(self.alpha_map_small_path), 
+                self.SMALL_SIZE
+            )
+        return self._alpha_map_small
+    
+    @property
+    def alpha_map_large(self) -> np.ndarray:
+        """Lazy-load large alpha map."""
+        if self._alpha_map_large is None:
+            self._alpha_map_large = AlphaMapManager.load_alpha_map(
+                Path(self.alpha_map_large_path), 
+                self.LARGE_SIZE
+            )
+        return self._alpha_map_large
+    
+    def detect_watermark_size(self, image: np.ndarray) -> Tuple[int, int]:
+        """
+        Auto-detect watermark size based on image dimensions.
+        
+        Args:
+            image: Input image as numpy array
+            
+        Returns:
+            Tuple of (size, margin) for watermark
+        """
+        height, width = image.shape[:2]
+        
+        if width <= self.SIZE_THRESHOLD or height <= self.SIZE_THRESHOLD:
+            return self.SMALL_SIZE, self.SMALL_MARGIN
+        else:
+            return self.LARGE_SIZE, self.LARGE_MARGIN
+    
+    def get_watermark_position(self, image: np.ndarray, 
+                                size: int, margin: int) -> Tuple[int, int]:
+        """
+        Get the (x, y) position of the watermark top-left corner.
+        Watermark is in bottom-right corner.
+        
+        Args:
+            image: Input image
+            size: Watermark size (48 or 96)
+            margin: Margin from edge
+            
+        Returns:
+            Tuple of (x, y) position
+        """
+        height, width = image.shape[:2]
+        
+        x = width - margin - size
+        y = height - margin - size
+        
+        return x, y
+    
+    def remove_watermark(self, image: np.ndarray, 
+                         force_size: str = None,
+                         verbose: bool = False) -> np.ndarray:
+        """
+        Remove watermark from image.
+        
+        Args:
+            image: Input image as numpy array (RGB)
+            force_size: Force 'small' (48x48) or 'large' (96x96)
+            verbose: Print debug information
+            
+        Returns:
+            Image with watermark removed
+        """
+        # Ensure RGB
+        if len(image.shape) == 2:
+            # Grayscale - convert to RGB
+            image = np.stack([image] * 3, axis=-1)
+        elif len(image.shape) == 3 and image.shape[2] == 4:
+            # RGBA - drop alpha
+            image = image[:, :, :3]
+        elif len(image.shape) != 3 or image.shape[2] != 3:
+            raise ValueError("Image must be RGB or RGBA")
+        
+        # Detect or force watermark size
+        if force_size == 'small':
+            size, margin = self.SMALL_SIZE, self.SMALL_MARGIN
+        elif force_size == 'large':
+            size, margin = self.LARGE_SIZE, self.LARGE_MARGIN
+        else:
+            size, margin = self.detect_watermark_size(image)
+        
+        if verbose:
+            print(f"  Watermark size: {size}x{size}, margin: {margin}px")
+        
+        # Get watermark position
+        x, y = self.get_watermark_position(image, size, margin)
+        
+        if verbose:
+            print(f"  Watermark position: ({x}, {y})")
+        
+        # Validate position
+        if x < 0 or y < 0:
+            if verbose:
+                print(f"  Warning: Image too small for {size}x{size} watermark")
+            return image.copy()
+        
+        # Get appropriate alpha map
+        if size == self.SMALL_SIZE:
+            alpha_map = self.alpha_map_small
+        else:
+            alpha_map = self.alpha_map_large
+        
+        # Remove watermark
+        result = BlendModes.remove_watermark_alpha_blend(
+            image=image,
+            alpha_map=alpha_map,
+            position=(x, y),
+            logo_value=self.logo_value
+        )
+        
+        return result
+    
+    def add_watermark(self, image: np.ndarray,
+                      force_size: str = None,
+                      verbose: bool = False) -> np.ndarray:
+        """
+        Add watermark to image (for testing purposes).
+        
+        Args:
+            image: Input image as numpy array (RGB)
+            force_size: Force 'small' (48x48) or 'large' (96x96)
+            verbose: Print debug information
+            
+        Returns:
+            Image with watermark added
+        """
+        # Ensure RGB
+        if len(image.shape) == 2:
+            image = np.stack([image] * 3, axis=-1)
+        elif len(image.shape) == 3 and image.shape[2] == 4:
+            image = image[:, :, :3]
+        elif len(image.shape) != 3 or image.shape[2] != 3:
+            raise ValueError("Image must be RGB or RGBA")
+        
+        # Detect or force watermark size
+        if force_size == 'small':
+            size, margin = self.SMALL_SIZE, self.SMALL_MARGIN
+        elif force_size == 'large':
+            size, margin = self.LARGE_SIZE, self.LARGE_MARGIN
+        else:
+            size, margin = self.detect_watermark_size(image)
+        
+        if verbose:
+            print(f"  Watermark size: {size}x{size}, margin: {margin}px")
+        
+        # Get watermark position
+        x, y = self.get_watermark_position(image, size, margin)
+        
+        # Validate position
+        if x < 0 or y < 0:
+            if verbose:
+                print(f"  Warning: Image too small for {size}x{size} watermark")
+            return image.copy()
+        
+        # Get appropriate alpha map
+        if size == self.SMALL_SIZE:
+            alpha_map = self.alpha_map_small
+        else:
+            alpha_map = self.alpha_map_large
+        
+        # Add watermark
+        result = BlendModes.add_watermark_alpha_blend(
+            image=image,
+            alpha_map=alpha_map,
+            position=(x, y),
+            logo_value=self.logo_value
+        )
+        
+        return result
+
+
+# ============================================================================
+# FILE OPERATIONS
+# ============================================================================
+
+def load_image(path: str) -> np.ndarray:
+    """Load image as RGB numpy array."""
+    img = Image.open(path).convert('RGB')
+    return np.array(img)
+
+
+def save_image(image: np.ndarray, path: str, quality: int = 95):
+    """Save numpy array as image."""
+    img = Image.fromarray(image)
+    
+    # Determine format from extension
+    ext = Path(path).suffix.lower()
+    
+    if ext in ['.jpg', '.jpeg']:
+        img.save(path, 'JPEG', quality=quality)
+    elif ext == '.png':
+        img.save(path, 'PNG')
+    elif ext == '.webp':
+        img.save(path, 'WEBP', quality=quality)
+    elif ext == '.bmp':
+        img.save(path, 'BMP')
+    else:
+        img.save(path)
+
+
+def process_single_file(input_path: str, 
+                        output_path: str = None,
+                        remover: GeminiWatermarkRemover = None,
+                        force_size: str = None,
+                        verbose: bool = False,
+                        add_mode: bool = False):
+    """
+    Process a single image file.
+    
+    Args:
+        input_path: Path to input image
+        output_path: Path for output (None = overwrite input)
+        remover: Watermark remover instance
+        force_size: Force 'small' or 'large'
+        verbose: Print verbose output
+        add_mode: Add watermark instead of removing
+    """
+    if output_path is None:
+        output_path = input_path
+    
+    if verbose:
+        print(f"Processing: {input_path}")
+    
+    # Load image
+    image = load_image(input_path)
+    
+    if verbose:
+        print(f"  Image size: {image.shape[1]}x{image.shape[0]}")
+    
+    # Process
+    if add_mode:
+        result = remover.add_watermark(image, force_size, verbose)
+        action = "Added watermark"
+    else:
+        result = remover.remove_watermark(image, force_size, verbose)
+        action = "Removed watermark"
+    
+    # Save result
+    save_image(result, output_path)
+    
+    if verbose:
+        print(f"  {action}, saved to: {output_path}")
+
+
+def process_directory(input_dir: str, 
+                      output_dir: str,
+                      remover: GeminiWatermarkRemover = None,
+                      force_size: str = None,
+                      verbose: bool = False,
+                      add_mode: bool = False):
+    """
+    Process all images in a directory.
+    
+    Args:
+        input_dir: Input directory path
+        output_dir: Output directory path
+        remover: Watermark remover instance
+        force_size: Force 'small' or 'large'
+        verbose: Print verbose output
+        add_mode: Add watermark instead of removing
+    """
+    input_path = Path(input_dir)
+    output_path = Path(output_dir)
+    
+    # Create output directory
+    output_path.mkdir(parents=True, exist_ok=True)
+    
+    # Supported formats
+    formats = {'.jpg', '.jpeg', '.png', '.webp', '.bmp'}
+    
+    # Find all images
+    images = [f for f in input_path.iterdir() 
+              if f.is_file() and f.suffix.lower() in formats]
+    
+    total = len(images)
+    if verbose:
+        print(f"Found {total} images to process")
+    
+    processed = 0
+    errors = 0
+    
+    for idx, img_path in enumerate(images, 1):
+        out_file = output_path / img_path.name
+        try:
+            if verbose:
+                print(f"\n[{idx}/{total}] ", end="")
+            process_single_file(
+                str(img_path), str(out_file), 
+                remover, force_size, verbose, add_mode
+            )
+            processed += 1
+        except Exception as e:
+            print(f"Error processing {img_path.name}: {e}")
+            errors += 1
+    
+    print(f"\nCompleted: {processed} processed, {errors} errors")
+
+
+# ============================================================================
+# CLI
+# ============================================================================
+
+def print_banner():
+    """Print ASCII banner."""
+    banner = r"""
+╔═══════════════════════════════════════════════════════════════════╗
+║                                                                   ║
+║         ██████╗ ███████╗███╗   ███╗██╗███╗   ██╗██╗               ║
+║        ██╔════╝ ██╔════╝████╗ ████║██║████╗  ██║██║               ║
+║        ██║  ███╗█████╗  ██╔████╔██║██║██╔██╗ ██║██║               ║
+║        ██║   ██║██╔══╝  ██║╚██╔╝██║██║██║╚██╗██║██║               ║
+║        ╚██████╔╝███████╗██║ ╚═╝ ██║██║██║ ╚████║██║               ║
+║         ╚═════╝ ╚══════╝╚═╝     ╚═╝╚═╝╚═╝  ╚═══╝╚═╝               ║
+║                                                                   ║
+║            Watermark Removal Tool (Python Edition)                ║
+║              Reverse Alpha Blending Technology                    ║
+║                                                                   ║
+╚═══════════════════════════════════════════════════════════════════╝
+    """
+    print(banner)
+
+
+def check_alpha_maps():
+    """Check if alpha map files exist and print status."""
+    small_exists = DEFAULT_ALPHA_MAP_SMALL_PATH.exists()
+    large_exists = DEFAULT_ALPHA_MAP_LARGE_PATH.exists()
+    
+    print("Alpha Map Status:")
+    print(f"  bg_48.png (48x48): {'✓ Found' if small_exists else '✗ Not found'} - {DEFAULT_ALPHA_MAP_SMALL_PATH}")
+    print(f"  bg_96.png (96x96): {'✓ Found' if large_exists else '✗ Not found'} - {DEFAULT_ALPHA_MAP_LARGE_PATH}")
+    print()
+    
+    if not small_exists or not large_exists:
+        print("WARNING: Missing alpha map files!")
+        print("Please ensure bg_48.png and bg_96.png are in the same directory as this script.")
+        print()
+        return False
+    
+    return True
+
+
+def main():
+    """Main entry point."""
+    parser = argparse.ArgumentParser(
+        description='Remove Gemini watermarks using reverse alpha blending',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog='''
+Examples:
+  # Simple mode - edit in place
+  python GWMRTool.py image.jpg
+  
+  # Specify output file
+  python GWMRTool.py -i input.jpg -o output.jpg
+  
+  # Batch processing
+  python GWMRTool.py -i ./input_folder/ -o ./output_folder/
+  
+  # Force specific watermark size
+  python GWMRTool.py -i image.jpg -o output.jpg --force-large
+  
+  # Add watermark (for testing)
+  python GWMRTool.py -i clean.jpg -o watermarked.jpg --add
+
+Required Files:
+  bg_48.png  - 48x48 alpha map (for images <= 1024px)
+  bg_96.png  - 96x96 alpha map (for images > 1024px)
+  
+  Place these files in the same directory as this script.
+
+Mathematical Formula:
+  Watermark Application:  watermarked = α × logo + (1 - α) × original
+  Reverse (Removal):      original = (watermarked - α × logo) / (1 - α)
+'''
+    )
+    
+    parser.add_argument('simple_input', nargs='?', 
+                        help='Image file (simple mode - edits in place)')
+    parser.add_argument('-i', '--input', 
+                        help='Input image file or directory')
+    parser.add_argument('-o', '--output', 
+                        help='Output image file or directory')
+    parser.add_argument('--force-small', action='store_true',
+                        help='Force 48x48 watermark size')
+    parser.add_argument('--force-large', action='store_true',
+                        help='Force 96x96 watermark size')
+    parser.add_argument('-v', '--verbose', action='store_true',
+                        help='Enable verbose output')
+    parser.add_argument('-q', '--quiet', action='store_true',
+                        help='Suppress all output except errors')
+    parser.add_argument('-b', '--banner', action='store_true',
+                        help='Show ASCII banner')
+    parser.add_argument('--add', action='store_true',
+                        help='Add watermark instead of removing (for testing)')
+    parser.add_argument('--alpha-small', 
+                        help='Path to custom 48x48 alpha map')
+    parser.add_argument('--alpha-large',
+                        help='Path to custom 96x96 alpha map')
+    parser.add_argument('--logo-value', type=float, default=255.0,
+                        help='Logo pixel value (default: 255 for white)')
+    parser.add_argument('--check', action='store_true',
+                        help='Check alpha map files and exit')
+    parser.add_argument('--version', action='version', 
+                        version='Gemini Watermark Remover v1.0.0 (Python)')
+    
+    args = parser.parse_args()
+    
+    # Show banner
+    if args.banner and not args.quiet:
+        print_banner()
+    
+    # Check mode
+    if args.check:
+        check_alpha_maps()
+        sys.exit(0)
+    
+    # Verbose mode (unless quiet)
+    verbose = args.verbose and not args.quiet
+    
+    # Determine force size
+    force_size = None
+    if args.force_small:
+        force_size = 'small'
+    elif args.force_large:
+        force_size = 'large'
+    
+    # Custom alpha map paths
+    alpha_small_path = Path(args.alpha_small) if args.alpha_small else None
+    alpha_large_path = Path(args.alpha_large) if args.alpha_large else None
+    
+    # Check default alpha maps exist (if not using custom)
+    if alpha_small_path is None and not DEFAULT_ALPHA_MAP_SMALL_PATH.exists():
+        print(f"Error: Alpha map not found: {DEFAULT_ALPHA_MAP_SMALL_PATH}")
+        print("Please ensure bg_48.png is in the same directory as this script.")
+        print("Or specify a custom path with --alpha-small")
+        sys.exit(1)
+    
+    if alpha_large_path is None and not DEFAULT_ALPHA_MAP_LARGE_PATH.exists():
+        print(f"Error: Alpha map not found: {DEFAULT_ALPHA_MAP_LARGE_PATH}")
+        print("Please ensure bg_96.png is in the same directory as this script.")
+        print("Or specify a custom path with --alpha-large")
+        sys.exit(1)
+    
+    # Initialize remover
+    try:
+        remover = GeminiWatermarkRemover(
+            alpha_map_small_path=alpha_small_path,
+            alpha_map_large_path=alpha_large_path,
+            logo_value=args.logo_value
+        )
+    except Exception as e:
+        print(f"Error initializing: {e}")
+        sys.exit(1)
+    
+    # Determine input/output
+    if args.simple_input:
+        # Simple mode - edit in place
+        if not Path(args.simple_input).exists():
+            print(f"Error: File not found: {args.simple_input}")
+            sys.exit(1)
+        
+        try:
+            process_single_file(
+                args.simple_input, None, remover, 
+                force_size, verbose, args.add
+            )
+        except Exception as e:
+            print(f"Error: {e}")
+            sys.exit(1)
+            
+    elif args.input:
+        input_path = Path(args.input)
+        
+        if not input_path.exists():
+            print(f"Error: Input not found: {args.input}")
+            sys.exit(1)
+        
+        if input_path.is_dir():
+            # Directory mode
+            if not args.output:
+                print("Error: Output directory required for batch processing (-o)")
+                sys.exit(1)
+            try:
+                process_directory(
+                    args.input, args.output, remover, 
+                    force_size, verbose, args.add
+                )
+            except Exception as e:
+                print(f"Error: {e}")
+                sys.exit(1)
+        else:
+            # Single file mode
+            output = args.output if args.output else args.input
+            try:
+                process_single_file(
+                    args.input, output, remover, 
+                    force_size, verbose, args.add
+                )
+            except Exception as e:
+                print(f"Error: {e}")
+                sys.exit(1)
+    else:
+        # No input provided - show help
+        parser.print_help()
+        print("\n" + "="*60)
+        check_alpha_maps()
+        sys.exit(0)
+    
+    if not args.quiet:
+        print("Done!")
+
+
+# ============================================================================
+# ENTRY POINT
+# ============================================================================
+
+if __name__ == '__main__':
+    main()