import gradio as gr
import cv2
import numpy as np
import torch
import torch.nn as nn
from PIL import Image
import torchvision.transforms as transforms
from skimage import restoration
import warnings
warnings.filterwarnings('ignore')

class WatermarkRemovalNet(nn.Module):
    def __init__(self):
        super(WatermarkRemovalNet, self).__init__()
        # Encoder
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),
            
            nn.Conv2d(64, 128, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),
            
            nn.Conv2d(128, 256, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, 3, padding=1),
            nn.ReLU(inplace=True),
        )
        
        # Decoder
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(256, 128, 2, stride=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, 3, padding=1),
            nn.ReLU(inplace=True),
            
            nn.ConvTranspose2d(128, 64, 2, stride=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.ReLU(inplace=True),
            
            nn.Conv2d(64, 3, 3, padding=1),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

class WatermarkRemover:
    def __init__(self):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model = WatermarkRemovalNet().to(self.device)
        self.transform = transforms.Compose([
            transforms.ToTensor(),
        ])
        
    def preprocess_image(self, image):
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        return self.transform(image).unsqueeze(0)
    
    def detect_watermark(self, img):
        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        denoised = cv2.fastNlMeansDenoising(gray)
        
        # Multi-scale watermark detection
        scales = [1.0, 0.5, 2.0]
        masks = []
        
        for scale in scales:
            scaled = cv2.resize(denoised, None, fx=scale, fy=scale)
            thresh = cv2.adaptiveThreshold(
                scaled, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                cv2.THRESH_BINARY_INV, 11, 2
            )
            
            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
            mask = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
            
            if scale != 1.0:
                mask = cv2.resize(mask, (denoised.shape[1], denoised.shape[0]))
            
            masks.append(mask)
        
        # Combine masks
        final_mask = np.zeros_like(denoised)
        for mask in masks:
            final_mask = cv2.bitwise_or(final_mask, mask)
            
        return final_mask

    def remove_watermark(self, img, mask, strength=0.8):
        # Apply inpainting
        inpainted = cv2.inpaint(img, mask, 3, cv2.INPAINT_TELEA)
        
        # Apply deep learning model
        tensor_img = self.preprocess_image(inpainted).to(self.device)
        with torch.no_grad():
            output = self.model(tensor_img)
        dl_result = output.squeeze(0).cpu().numpy().transpose(1, 2, 0)
        
        # Apply image restoration
        denoise_img = restoration.denoise_tv_chambolle(dl_result, weight=0.1)
        enhanced = cv2.detailEnhance(denoise_img.astype(np.float32), sigma_s=10, sigma_r=0.15)
        
        # Blend results
        result = cv2.addWeighted(enhanced, strength, img.astype(np.float32)/255, 1-strength, 0)
        return (result * 255).astype(np.uint8)

def process_image(input_image, strength, enhance_details):
    remover = WatermarkRemover()
    
    # Detect watermark
    mask = remover.detect_watermark(input_image)
    
    # Remove watermark
    result = remover.remove_watermark(input_image, mask, strength)
    
    # Optional detail enhancement
    if enhance_details:
        result = cv2.detailEnhance(result, sigma_s=10, sigma_r=0.15)
    
    return result, mask

# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # Removedor Avançado de Marca D'água
    Este aplicativo utiliza uma combinação de deep learning e processamento de imagem para remover marcas d'água.
    """)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Imagem Original", type="numpy")
            strength_slider = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.8,
                step=0.1,
                label="Intensidade da Remoção"
            )
            enhance_checkbox = gr.Checkbox(
                label="Melhorar Detalhes",
                value=True
            )
            process_btn = gr.Button("Processar Imagem", variant="primary")
            
        with gr.Column():
            output_image = gr.Image(label="Imagem Processada")
            mask_image = gr.Image(label="Máscara de Marca D'água")

    
    # Eventos
    process_btn.click(
        fn=process_image,
        inputs=[input_image, strength_slider, enhance_checkbox],
        outputs=[output_image, mask_image]
    )
    
    gr.Markdown("""
    ## Recursos Avançados:
    - Detecção multi-escala de marca d'água
    - Modelo de deep learning para reconstrução de imagem
    - Restauração avançada de imagem
    - Preservação de detalhes
    - Controle de intensidade ajustável
    
    ## Dicas de Uso:
    1. Ajuste a 'Intensidade da Remoção' para controlar o balanço entre remoção da marca e preservação da imagem
    2. Ative 'Melhorar Detalhes' para realçar a qualidade da imagem final
    3. Visualize a máscara para entender quais áreas estão sendo processadas
    """)

if __name__ == "__main__":
    demo.launch(share=True)