Update app.py

app.py

@@ -1,93 +1,181 @@
 import gradio as gr
 import cv2
 import numpy as np
+import torch
+import torch.nn as nn
 from PIL import Image
-import io
+import torchvision.transforms as transforms
+from skimage import restoration
+import warnings
+warnings.filterwarnings('ignore')
 
-def process_image(input_image, threshold_value, kernel_size):
-    # Converter para array numpy
-    img = np.array(input_image)
-    
-    # Converter para escala de cinza
-    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
-    
-    # Aplicar threshold adaptativo
-    thresh = cv2.adaptiveThreshold(
-        gray,
-        255,
-        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
-        cv2.THRESH_BINARY_INV,
-        threshold_value,
-        2
-    )
-    
-    # Criar kernel para operações morfológicas
-    kernel = np.ones((kernel_size, kernel_size), np.uint8)
+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)
     
-    # Aplicar operações morfológicas
-    opening = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
-    closing = cv2.morphologyEx(opening, cv2.MORPH_CLOSE, kernel)
+    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()
     
-    # Encontrar contornos da marca d'água
-    contours, _ = cv2.findContours(closing, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    # Detect watermark
+    mask = remover.detect_watermark(input_image)
     
-    # Criar máscara para inpainting
-    mask = np.zeros_like(gray)
-    for cnt in contours:
-        cv2.drawContours(mask, [cnt], -1, (255,255,255), -1)
+    # Remove watermark
+    result = remover.remove_watermark(input_image, mask, strength)
     
-    # Realizar inpainting
-    result = cv2.inpaint(img, mask, 3, cv2.INPAINT_TELEA)
+    # Optional detail enhancement
+    if enhance_details:
+        result = cv2.detailEnhance(result, sigma_s=10, sigma_r=0.15)
     
     return result, mask
 
-# Interface Gradio
+# Gradio Interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# Removedor de Marca D'água")
+    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")
-            threshold_slider = gr.Slider(
-                minimum=3, 
-                maximum=101,
-                step=2,
-                value=11,
-                label="Valor do Threshold"
+            strength_slider = gr.Slider(
+                minimum=0.1,
+                maximum=1.0,
+                value=0.8,
+                step=0.1,
+                label="Intensidade da Remoção"
             )
-            kernel_slider = gr.Slider(
-                minimum=1,
-                maximum=9,
-                step=2,
-                value=3,
-                label="Tamanho do Kernel"
+            enhance_checkbox = gr.Checkbox(
+                label="Melhorar Detalhes",
+                value=True
             )
-            process_btn = gr.Button("Processar Imagem")
+            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 Detectada")
+            mask_image = gr.Image(label="Máscara de Marca D'água")
+
     
     # Eventos
     process_btn.click(
         fn=process_image,
-        inputs=[input_image, threshold_slider, kernel_slider],
+        inputs=[input_image, strength_slider, enhance_checkbox],
         outputs=[output_image, mask_image]
     )
     
     gr.Markdown("""
-    ## Como usar:
-    1. Faça upload de uma imagem com marca d'água
-    2. Ajuste os parâmetros:
-        - Valor do Threshold: Controla a sensibilidade da detecção
-        - Tamanho do Kernel: Afeta o tamanho da área de processamento
-    3. Clique em "Processar Imagem"
-    4. Visualize o resultado e a máscara detectada
+    ## 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
     
-    ## Observações:
-    - O resultado pode variar dependendo do tipo e complexidade da marca d'água
-    - Ajuste os parâmetros para obter melhores resultados
-    - A máscara mostra as áreas detectadas como marca d'água
+    ## 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__":