app.py

import gradio as gr
import numpy as np
import cv2
from PIL import Image
import tempfile
import os
from pathlib import Path
from typing import Optional, Tuple, List
import subprocess
import shutil
import time

# Try to import ML libraries with graceful fallbacks
try:
    import torch
    HAS_TORCH = True
except ImportError:
    HAS_TORCH = False
    print("Warning: torch not installed. Using CPU-only methods.")

try:
    from transformers import pipeline, AutoImageProcessor, AutoModelForImageSegmentation
    from diffusers import StableDiffusionInpaintPipeline, DPMSolverMultistepScheduler
    HAS_TRANSFORMERS = True
except ImportError:
    HAS_TRANSFORMERS = False
    print("Warning: transformers/diffusers not installed. Using OpenCV fallback.")

try:
    from rembg import remove as rembg_remove
    HAS_REMBG = True
except ImportError:
    HAS_REMBG = False

try:
    from skimage import morphology, measure
    from skimage.filters import threshold_otsu
    HAS_SKIMAGE = True
except ImportError:
    HAS_SKIMAGE = False

print(f"Available: torch={HAS_TORCH}, transformers={HAS_TRANSFORMERS}, rembg={HAS_REMBG}, skimage={HAS_SKIMAGE}")


class WatermarkRemover:
    """Main class for watermark removal using various techniques."""
    
    def __init__(self):
        self.device = "cuda" if HAS_TORCH and torch.cuda.is_available() else "cpu"
        self.sd_pipeline = None
        self.sam_processor = None
        self.sam_model = None
        self._load_models()
    
    def _load_models(self):
        """Load all required models with error handling."""
        print(f"Initializing on {self.device}...")
        
        if HAS_TRANSFORMERS and HAS_TORCH:
            try:
                # Try to load a lightweight inpainting model
                print("Attempting to load inpainting model...")
                # Use a smaller model for better compatibility
                model_id = "runwayml/stable-diffusion-inpainting"
                
                # Check if model exists locally first
                from huggingface_hub import snapshot_download
                try:
                    model_path = snapshot_download(model_id, local_files_only=True)
                    print(f"Found model at {model_path}")
                except:
                    print("Model not cached, skipping download for faster startup")
                    self.sd_pipeline = None
                    return
                
                self.sd_pipeline = StableDiffusionInpaintPipeline.from_pretrained(
                    model_id,
                    torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
                    local_files_only=True,
                )
                self.sd_pipeline.scheduler = DPMSolverMultistepScheduler.from_config(
                    self.sd_pipeline.scheduler.config
                )
                self.sd_pipeline = self.sd_pipeline.to(self.device)
                print("✓ SD inpainting model loaded")
            except Exception as e:
                print(f"⚠ Could not load SD model: {e}")
                self.sd_pipeline = None
            
            try:
                print("Attempting to load SAM model...")
                self.sam_processor = AutoImageProcessor.from_pretrained(
                    "facebook/sam-vit-base",
                    local_files_only=True
                )
                self.sam_model = AutoModelForImageSegmentation.from_pretrained(
                    "facebook/sam-vit-base",
                    local_files_only=True
                )
                self.sam_model = self.sam_model.to(self.device)
                print("✓ SAM model loaded")
            except Exception as e:
                print(f"⚠ Could not load SAM model: {e}")
                self.sam_processor = None
                self.sam_model = None
        else:
            print("⚠ ML libraries not available, using OpenCV methods only")
    
    def detect_watermark_region(self, image: np.ndarray, method: str = "automatic") -> np.ndarray:
        """Detect watermark region in image. Returns binary mask."""
        h, w = image.shape[:2]
        mask = np.zeros((h, w), dtype=np.uint8)
        
        if method == "automatic":
            # Convert to grayscale
            if len(image.shape) == 3:
                gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
            else:
                gray = image
            
            corner_size = min(h, w) // 4
            corners = [
                ("tr", 0, w-corner_size),
                ("br", h-corner_size, w-corner_size),
                ("bl", h-corner_size, 0),
                ("tl", 0, 0)
            ]
            
            for name, y0, x0 in corners:
                region = gray[y0:y0+corner_size, x0:x0+corner_size]
                
                # Edge detection
                edges = cv2.Canny(region, 50, 150)
                edge_density = np.sum(edges > 0) / edges.size
                
                # Threshold analysis
                _, thresh = cv2.threshold(region, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
                white_ratio = np.sum(thresh > 128) / thresh.size
                
                if edge_density > 0.05 or (0.1 < white_ratio < 0.9):
                    mask[y0:y0+corner_size, x0:x0+corner_size] = 255
            
            # Laplacian variance for high-contrast regions
            laplacian = cv2.Laplacian(gray, cv2.CV_64F)
            laplacian = np.abs(laplacian)
            _, high_var = cv2.threshold(laplacian.astype(np.uint8), 30, 255, cv2.THRESH_BINARY)
            
            kernel = np.ones((5, 5), np.uint8)
            high_var = cv2.dilate(high_var, kernel, iterations=2)
            
            contours, _ = cv2.findContours(high_var, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            
            for cnt in contours:
                x, y, cw, ch = cv2.boundingRect(cnt)
                in_corner = (x < corner_size and y < corner_size) or \
                           (x > w - corner_size - cw and y < corner_size) or \
                           (x < corner_size and y > h - corner_size - ch) or \
                           (x > w - corner_size - cw and y > h - corner_size - ch)
                
                if in_corner and cw * ch > 100:
                    cv2.rectangle(mask, (x, y), (x+cw, y+ch), 255, -1)
        
        elif method == "bottom_right":
            corner_size = min(h, w) // 5
            mask[h-corner_size:h, w-corner_size:w] = 255
            
        elif method == "bottom_left":
            corner_size = min(h, w) // 5
            mask[h-corner_size:h, 0:corner_size] = 255
            
        elif method == "top_right":
            corner_size = min(h, w) // 5
            mask[0:corner_size, w-corner_size:w] = 255
            
        elif method == "top_left":
            corner_size = min(h, w) // 5
            mask[0:corner_size, 0:corner_size] = 255
            
        elif method == "center":
            ch, cw = h // 4, w // 4
            mask[h//2-ch:h//2+ch, w//2-cw:w//2+cw] = 255
        
        # Morphological operations
        kernel = np.ones((15, 15), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        mask = cv2.dilate(mask, kernel, iterations=2)
        
        return mask
    
    def inpaint_image(self, image: np.ndarray, mask: np.ndarray, 
                     method: str = "opencv", prompt: str = "") -> np.ndarray:
        """Inpaint image using specified method."""
        # Ensure RGB
        if len(image.shape) == 2:
            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
        elif image.shape[2] == 4:
            image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
        
        mask_binary = (mask > 128).astype(np.uint8) * 255
        
        # Always use OpenCV for reliability
        print("Using OpenCV inpainting (most reliable)")
        result = cv2.inpaint(image, mask_binary, 3, cv2.INPAINT_TELEA)
        return result
    
    def remove_watermark(self, image_input, detection_method: str = "automatic",
                        inpaint_method: str = "opencv", custom_prompt: str = "",
                        manual_mask: Optional[np.ndarray] = None) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Main watermark removal function."""
        # Load image
        if isinstance(image_input, str):
            image = cv2.imread(image_input)
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        elif isinstance(image_input, np.ndarray):
            image = image_input.copy()
        elif isinstance(image_input, Image.Image):
            image = np.array(image_input)
            if len(image.shape) == 3 and image.shape[2] == 4:
                image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
        elif isinstance(image_input, dict):
            image = np.array(image_input.get("image", image_input))
            if len(image.shape) == 3 and image.shape[2] == 4:
                image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
        else:
            raise ValueError(f"Unsupported image type: {type(image_input)}")
        
        # Ensure RGB
        if len(image.shape) == 2:
            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
        
        # Get mask
        if manual_mask is not None:
            mask = manual_mask
        else:
            mask = self.detect_watermark_region(image, detection_method)
        
        # Create visualization
        mask_viz = image.copy()
        red_overlay = np.zeros_like(image)
        red_overlay[mask > 128] = [255, 0, 0]
        mask_viz = cv2.addWeighted(mask_viz, 0.7, red_overlay, 0.3, 0)
        
        # Inpaint
        result = self.inpaint_image(image, mask, inpaint_method, custom_prompt)
        
        return mask_viz, result, mask
    
    def enhance_image(self, image: np.ndarray, enhancement_type: str) -> np.ndarray:
        """Apply image enhancements."""
        if enhancement_type == "sharpen":
            kernel = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]])
            return cv2.filter2D(image, -1, kernel)
        
        elif enhancement_type == "denoise":
            return cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21)
        
        elif enhancement_type == "contrast":
            lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB)
            l, a, b = cv2.split(lab)
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
            l = clahe.apply(l)
            return cv2.cvtColor(cv2.merge([l,a,b]), cv2.COLOR_LAB2RGB)
        
        elif enhancement_type == "brightness_up":
            return cv2.convertScaleAbs(image, alpha=1.0, beta=30)
        
        elif enhancement_type == "brightness_down":
            return cv2.convertScaleAbs(image, alpha=1.0, beta=-30)
        
        elif enhancement_type == "saturation_up":
            hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV).astype(np.float32)
            hsv[:,:,1] = np.clip(hsv[:,:,1] * 1.3, 0, 255)
            return cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2RGB)
        
        elif enhancement_type == "upscale_2x":
            return cv2.resize(image, None, fx=2, fy=2, interpolation=cv2.INTER_LANCZOS4)
        
        return image
    
    def remove_background(self, image: np.ndarray) -> np.ndarray:
        """Remove image background."""
        if HAS_REMBG:
            try:
                pil_img = Image.fromarray(image)
                result = rembg_remove(pil_img)
                return np.array(result)
            except Exception as e:
                print(f"rembg failed: {e}")
        
        # Fallback: simple threshold-based segmentation
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
            _, mask = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
            
            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            if contours:
                largest = max(contours, key=cv2.contourArea)
                mask = np.zeros_like(mask)
                cv2.drawContours(mask, [largest], -1, 255, -1)
                
                rgba = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
                rgba[:,:,3] = mask
                return rgba
        
        return image
    
    def process_video(self, video_path: str, detection_method: str = "automatic",
                     inpaint_method: str = "opencv", progress=None) -> str:
        """Process video frame by frame."""
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise ValueError(f"Could not open video: {video_path}")
        
        fps = cap.get(cv2.CAP_PROP_FPS)
        if fps <= 0:
            fps = 30.0
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        if total_frames <= 0:
            cap.release()
            raise ValueError("Could not determine frame count")
        
        # Create temp output
        temp_output = tempfile.mktemp(suffix='.mp4')
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(temp_output, fourcc, fps, (width, height))
        
        # Detect watermark on first frame
        ret, first_frame = cap.read()
        if not ret:
            cap.release()
            raise ValueError("Could not read first frame")
        
        first_frame_rgb = cv2.cvtColor(first_frame, cv2.COLOR_BGR2RGB)
        mask = self.detect_watermark_region(first_frame_rgb, detection_method)
        
        # Dilate mask for video
        kernel = np.ones((20, 20), np.uint8)
        mask = cv2.dilate(mask, kernel, iterations=3)
        mask_binary = (mask > 128).astype(np.uint8) * 255
        
        # Process frames
        frame_count = 0
        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            result = cv2.inpaint(frame_rgb, mask_binary, 3, cv2.INPAINT_TELEA)
            result_bgr = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
            out.write(result_bgr)
            
            frame_count += 1
            if progress is not None and total_frames > 0:
                progress(frame_count / total_frames, f"Processing {frame_count}/{total_frames}")
        
        cap.release()
        out.release()
        
        # Try to re-encode with ffmpeg
        final_output = self._reencode_video(temp_output)
        return final_output
    
    def _reencode_video(self, input_path: str) -> str:
        """Re-encode video with better codec."""
        output_path = tempfile.mktemp(suffix='.mp4')
        
        ffmpeg_cmd = [
            'ffmpeg', '-y', '-i', input_path,
            '-c:v', 'libx264', '-preset', 'fast', '-crf', '23',
            '-c:a', 'aac', '-b:a', '128k',
            '-movflags', '+faststart',
            output_path
        ]
        
        try:
            result = subprocess.run(ffmpeg_cmd, check=True, capture_output=True, timeout=300)
            if os.path.exists(output_path) and os.path.getsize(output_path) > 0:
                if os.path.exists(input_path):
                    os.remove(input_path)
                return output_path
        except Exception as e:
            print(f"ffmpeg failed: {e}")
        
        if os.path.exists(output_path):
            os.remove(output_path)
        return input_path


# Global instance
_watermark_remover = None

def get_remover():
    """Get or create watermark remover instance."""
    global _watermark_remover
    if _watermark_remover is None:
        _watermark_remover = WatermarkRemover()
    return _watermark_remover


def process_image(image, detection_method, inpaint_method, custom_prompt, 
                  apply_enhancement, enhancement_type):
    """Process single image."""
    if image is None:
        return None, None, None, "⚠ Please upload an image first."
    
    try:
        remover = get_remover()
        
        # Handle different image input types
        if isinstance(image, dict):
            img_array = image.get("image", image)
            if isinstance(img_array, np.ndarray):
                img_array = img_array
            else:
                img_array = np.array(img_array)
        elif isinstance(image, np.ndarray):
            img_array = image.copy()
        elif isinstance(image, Image.Image):
            img_array = np.array(image)
        elif isinstance(image, str):
            img_array = cv2.imread(image)
            img_array = cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB)
        else:
            img_array = np.array(image)
        
        # Remove watermark
        mask_viz, result, mask = remover.remove_watermark(
            img_array, 
            detection_method=detection_method,
            inpaint_method=inpaint_method,
            custom_prompt=custom_prompt
        )
        
        # Apply enhancement
        if apply_enhancement and enhancement_type != "none":
            result = remover.enhance_image(result, enhancement_type)
        
        return mask_viz, result, mask, "✅ Processing complete!"
        
    except Exception as e:
        import traceback
        error_msg = f"❌ Error: {str(e)}"
        print(error_msg)
        print(traceback.format_exc())
        return None, None, None, error_msg


def process_video_file(video, detection_method, inpaint_method, progress=gr.Progress()):
    """Process video file."""
    if video is None:
        return None, "⚠ Please upload a video first."
    
    try:
        remover = get_remover()
        
        # Get video path
        if isinstance(video, str):
            video_path = video
        elif hasattr(video, 'name'):
            video_path = video.name
        else:
            video_path = str(video)
        
        if not os.path.exists(video_path):
            return None, f"❌ Video file not found: {video_path}"
        
        progress(0, "Starting video processing...")
        
        output_path = remover.process_video(
            video_path, 
            detection_method=detection_method,
            inpaint_method=inpaint_method,
            progress=lambda p, msg: progress(p, msg)
        )
        
        if os.path.exists(output_path):
            return output_path, "✅ Video processing complete!"
        else:
            return None, "❌ Failed to create output video"
        
    except Exception as e:
        import traceback
        error_msg = f"❌ Error: {str(e)}"
        print(error_msg)
        print(traceback.format_exc())
        return None, error_msg


def batch_process_images(files, detection_method, inpaint_method, custom_prompt,
                        apply_enhancement, enhancement_type, progress=gr.Progress()):
    """Process multiple images in batch."""
    if not files:
        return [], "⚠ No files uploaded."
    
    results = []
    remover = get_remover()
    
    # Handle files list
    if isinstance(files, str):
        files = [files]
    elif hasattr(files, 'name'):
        files = [files.name]
    elif isinstance(files, (list, tuple)):
        files = [f.name if hasattr(f, 'name') else str(f) for f in files]
    
    for i, file_path in enumerate(files):
        progress((i) / len(files), f"Processing {i+1}/{len(files)}")
        
        try:
            if not os.path.exists(file_path):
                print(f"File not found: {file_path}")
                continue
            
            img = cv2.imread(file_path)
            if img is None:
                print(f"Could not read: {file_path}")
                continue
            
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            
            _, result, _, _ = remover.remove_watermark(
                img, detection_method, inpaint_method, custom_prompt
            )
            
            if apply_enhancement and enhancement_type != "none":
                result = remover.enhance_image(result, enhancement_type)
            
            results.append(result)
            
        except Exception as e:
            print(f"Error processing {file_path}: {e}")
    
    progress(1.0, "Complete!")
    
    if results:
        return results, f"✅ Processed {len(results)} images."
    else:
        return [], "❌ No images were successfully processed."


def apply_enhancement_only(image, enh_type, remove_bg):
    """Apply enhancement without watermark removal."""
    if image is None:
        return None, "⚠ Please upload an image."
    
    try:
        remover = get_remover()
        
        # Handle image input
        if isinstance(image, dict):
            img_array = image.get("image", image)
            if isinstance(img_array, np.ndarray):
                result = img_array.copy()
            else:
                result = np.array(img_array)
        elif isinstance(image, np.ndarray):
            result = image.copy()
        elif isinstance(image, Image.Image):
            result = np.array(image)
        elif isinstance(image, str):
            result = cv2.imread(image)
            result = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
        else:
            result = np.array(image)
        
        status_msg = ""
        
        if remove_bg:
            result = remover.remove_background(result)
            status_msg = "Background removed. "
        
        result = remover.enhance_image(result, enh_type)
        
        return result, status_msg + f"Applied {enh_type} enhancement."
        
    except Exception as e:
        return None, f"❌ Error: {str(e)}"


# Custom CSS
custom_css = """
.watermark-tool {
    max-width: 1400px;
    margin: 0 auto;
}
.header-title {
    text-align: center;
    margin-bottom: 1rem;
}
.header-title h1 {
    font-size: 2.5rem;
    font-weight: 700;
    background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
}
.built-with {
    text-align: center;
    margin-top: 0.5rem;
    font-size: 0.9rem;
}
.built-with a {
    color: #667eea;
    text-decoration: none;
    font-weight: 500;
}
.built-with a:hover {
    text-decoration: underline;
}
.tool-description {
    text-align: center;
    color: #666;
    margin-bottom: 2rem;
}
"""


# Create Gradio application
with gr.Blocks() as demo:
    
    # Header
    with gr.Row():
        with gr.Column():
            gr.Markdown(
                """
                <div class="header-title">
                    <h1>🎨 AI Watermark Remover</h1>
                </div>
                <div class="built-with">
                    Built with <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank">anycoder</a>
                </div>
                <div class="tool-description">
                    Remove watermarks from images and videos using AI inpainting. 
                    Supports automatic detection and manual masking.
                </div>
                """,
                elem_classes="watermark-tool"
            )
    
    # Main tabs
    with gr.Tabs():
        
        # Image Watermark Removal Tab
        with gr.Tab("🖼️ Image Removal"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Input Settings")
                    
                    image_input = gr.Image(
                        label="Upload Image",
                        type="numpy",
                        height=400
                    )
                    
                    with gr.Accordion("Detection Settings", open=True):
                        detection_method = gr.Radio(
                            choices=[
                                ("Automatic (AI)", "automatic"),
                                ("Bottom Right", "bottom_right"),
                                ("Bottom Left", "bottom_left"),
                                ("Top Right", "top_right"),
                                ("Top Left", "top_left"),
                                ("Center", "center")
                            ],
                            value="automatic",
                            label="Watermark Location"
                        )
                    
                    with gr.Accordion("Inpainting Settings", open=True):
                        inpaint_method = gr.Radio(
                            choices=[
                                ("OpenCV (Fast)", "opencv"),
                            ],
                            value="opencv",
                            label="Processing Method"
                        )
                        
                        custom_prompt = gr.Textbox(
                            label="Custom Prompt (optional)",
                            placeholder="Describe what should replace watermark",
                            value="",
                            info="For future AI model support"
                        )
                    
                    with gr.Accordion("Enhancement", open=False):
                        apply_enhancement = gr.Checkbox(
                            label="Apply post-processing",
                            value=False
                        )
                        enhancement_type = gr.Dropdown(
                            choices=[
                                ("None", "none"),
                                ("Sharpen", "sharpen"),
                                ("Denoise", "denoise"),
                                ("Enhance Contrast", "contrast"),
                                ("Increase Brightness", "brightness_up"),
                                ("Decrease Brightness", "brightness_down"),
                                ("Increase Saturation", "saturation_up"),
                                ("2x Upscale", "upscale_2x")
                            ],
                            value="none",
                            label="Enhancement Type"
                        )
                    
                    process_btn = gr.Button(
                        "🚀 Remove Watermark",
                        variant="primary",
                        size="lg"
                    )
                
                with gr.Column(scale=1):
                    gr.Markdown("### Results")
                    
                    with gr.Tabs():
                        with gr.Tab("Result"):
                            result_output = gr.Image(
                                label="Clean Image",
                                height=400
                            )
                        with gr.Tab("Detection Mask"):
                            mask_viz_output = gr.Image(
                                label="Detected Region",
                                height=400
                            )
                        with gr.Tab("Raw Mask"):
                            mask_output = gr.Image(
                                label="Binary Mask",
                                height=400
                            )
                    
                    status_text = gr.Textbox(
                        label="Status",
                        interactive=False
                    )
            
            process_btn.click(
                fn=process_image,
                inputs=[image_input, detection_method, inpaint_method, custom_prompt,
                       apply_enhancement, enhancement_type],
                outputs=[mask_viz_output, result_output, mask_output, status_text],
                api_visibility="public"
            )
        
        # Video Watermark Removal Tab
        with gr.Tab("🎬 Video Removal"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Video Input")
                    
                    video_input = gr.Video(
                        label="Upload Video",
                        format="mp4"
                    )
                    
                    with gr.Accordion("Settings", open=True):
                        video_detection = gr.Radio(
                            choices=[
                                ("Automatic", "automatic"),
                                ("Bottom Right", "bottom_right"),
                                ("Bottom Left", "bottom_left")
                            ],
                            value="automatic",
                            label="Watermark Location"
                        )
                        
                        video_inpaint = gr.Radio(
                            choices=[("OpenCV (Fast)", "opencv")],
                            value="opencv",
                            label="Processing Method"
                        )
                    
                    video_process_btn = gr.Button(
                        "🎬 Process Video",
                        variant="primary",
                        size="lg"
                    )
                
                with gr.Column(scale=1):
                    gr.Markdown("### Result")
                    
                    video_output = gr.Video(
                        label="Clean Video",
                        format="mp4"
                    )
                    
                    video_status = gr.Textbox(
                        label="Status",
                        interactive=False
                    )
            
            video_process_btn.click(
                fn=process_video_file,
                inputs=[video_input, video_detection, video_inpaint],
                outputs=[video_output, video_status],
                api_visibility="public"
            )
        
        # Batch Processing Tab
        with gr.Tab("📁 Batch Processing"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Batch Input")
                    
                    batch_files = gr.File(
                        label="Upload Multiple Images",
                        file_count="multiple",
                        file_types=["image"]
                    )
                    
                    batch_detection = gr.Radio(
                        choices=[
                            ("Automatic", "automatic"),
                            ("Bottom Right", "bottom_right"),
                            ("Bottom Left", "bottom_left")
                        ],
                        value="automatic",
                        label="Detection Method"
                    )
                    
                    batch_inpaint = gr.Radio(
                        choices=[
                            ("OpenCV Fast", "opencv"),
                        ],
                        value="opencv",
                        label="Processing Method"
                    )
                    
                    batch_enhance = gr.Checkbox(
                        label="Apply Enhancement",
                        value=False
                    )
                    
                    batch_process_btn = gr.Button(
                        "📦 Process Batch",
                        variant="primary",
                        size="lg"
                    )
                
                with gr.Column(scale=2):
                    gr.Markdown("### Results")
                    
                    batch_output = gr.Gallery(
                        label="Processed Images",
                        columns=3,
                        rows=2,
                        height=600
                    )
                    
                    batch_status = gr.Textbox(
                        label="Status",
                        interactive=False
                    )
            
            batch_process_btn.click(
                fn=batch_process_images,
                inputs=[batch_files, batch_detection, batch_inpaint, 
                       gr.State(""), batch_enhance, gr.State("none")],
                outputs=[batch_output, batch_status],
                api_visibility="public"
            )
        
        # Image Enhancement Tab
        with gr.Tab("✨ Enhancement"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Enhance Images")
                    
                    enhance_input = gr.Image(
                        label="Upload Image",
                        type="numpy"
                    )
                    
                    enhance_type = gr.Dropdown(
                        choices=[
                            ("Sharpen", "sharpen"),
                            ("Denoise", "denoise"),
                            ("Enhance Contrast", "contrast"),
                            ("Increase Brightness", "brightness_up"),
                            ("Decrease Brightness", "brightness_down"),
                            ("Increase Saturation", "saturation_up"),
                            ("2x Upscale", "upscale_2x")
                        ],
                        value="sharpen",
                        label="Enhancement Type"
                    )
                    
                    bg_remove_check = gr.Checkbox(
                        label="Remove Background",
                        value=False,
                        info="AI background removal"
                    )
                    
                    enhance_btn = gr.Button(
                        "✨ Apply Enhancement",
                        variant="primary",
                        size="lg"
                    )
                
                with gr.Column(scale=1):
                    gr.Markdown("### Result")
                    
                    enhance_output = gr.Image(
                        label="Enhanced Image",
                        height=400
                    )
                    
                    enhance_status = gr.Textbox(
                        label="Status",
                        interactive=False
                    )
            
            enhance_btn.click(
                fn=apply_enhancement_only,
                inputs=[enhance_input, enhance_type, bg_remove_check],
                outputs=[enhance_output, enhance_status],
                api_visibility="public"
            )
        
        # About Tab
        with gr.Tab("ℹ️ About"):
            gr.Markdown(
                """
                ## About AI Watermark Remover
                
                This tool removes watermarks from images and videos using computer vision techniques.
                
                ### Features
                
                - 🎯 **Automatic Detection**: AI-powered watermark location detection
                - 🖼️ **Image Processing**: Remove watermarks from single images
                - 🎬 **Video Support**: Process videos frame by frame
                - 📁 **Batch Processing**: Handle multiple images at once
                - ✨ **Enhancement Tools**: Sharpen, denoise, adjust colors, upscale
                
                ### How It Works
                
                1. **Detection**: Analyzes corners and edges where watermarks typically appear
                2. **Masking**: Creates a binary mask covering the watermark region
                3. **Inpainting**: Uses OpenCV to fill the masked area seamlessly
                4. **Enhancement**: Optional post-processing to improve quality
                
                ### Tips for Best Results
                
                - Use "Automatic" detection for unknown watermark positions
                - Specify corner locations if you know where the watermark is
                - For videos, static watermarks work best
                - Try different enhancement options for best quality
                
                ### Limitations
                
                - Large watermarks (>30% of image) may be challenging
                - Semi-transparent watermarks over complex textures are harder
                - Video processing optimized for static watermarks
                """
            )


# Launch application
if __name__ == "__main__":
    print("🚀 Initializing Watermark Remover...")
    try:
        remover = get_remover()
        print("✓ Models initialized")
    except Exception as e:
        print(f"⚠ Model initialization warning: {e}")
    
    print("🌐 Starting Gradio server...")
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        css=custom_css,
        theme=gr.themes.Soft(
            primary_hue="indigo",
            secondary_hue="purple",
            neutral_hue="slate"
        ).set(
            button_primary_background_fill="*primary_600",
            button_primary_background_fill_hover="*primary_700",
            block_title_text_weight="600",
            block_label_text_weight="500",
        ),
        footer_links=[{"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}]
    )