app.py

import gradio as gr
import numpy as np
import imageio
import cv2
import numpy as np
from inpaint import InpaintingTester
import os 
import shutil
import re 
import uuid

def create_mask(watermark, mask_type="white"):
    """
    Create a mask for the watermark region.
    mask_type: 'white' for white mask and 'black' for black mask
    """
    h, w, _ = watermark.shape
    if mask_type == "white":
        return np.ones((h, w), dtype=np.uint8) * 255  # White mask
    elif mask_type == "black":
        return np.zeros((h, w), dtype=np.uint8)  # Black mask
    return None


def inpaint_watermark(watermark, mask):
    """Inpaint the watermark region using the mask."""
    return cv2.inpaint(watermark, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)


def place_inpainted_back(image, inpainted_region, location):
    """Place the inpainted region back into the original image."""
    x_start, y_start, x_end, y_end = location
    image[y_start:y_end, x_start:x_end] = inpainted_region
    return image


def extract_watermark(image, height_ratio=0.15, width_ratio=0.15, margin=0):
    """Extract watermark from the image using given ratios and margin."""
    h, w, _ = image.shape
    crop_h, crop_w = int(h * height_ratio), int(w * width_ratio)
    x_start, y_start = w - crop_w, h - crop_h
    watermark = image[y_start:h-margin, x_start:w-margin]
    location = (x_start, y_start, w-margin, h-margin)
    return watermark, location


def load_inpainting_model():
    """Load the inpainting model."""
    save_path = "./output"
    # resize_to = None  # Default size from config
    resize_to = (480,480)
    return InpaintingTester(save_path, resize_to)


def process_image_with_model(image_path, mask_path, tester):
    """Process the image using the inpainting model and return the cleaned image path."""
    image_mask_pairs = [(image_path, mask_path)]
    return tester.process_multiple_images(image_mask_pairs)[0]





def img_file_name(image_path):
    global image_folder
    text=os.path.basename(image_path)
    text=text.split(".")[0]
    # Remove all non-alphabetic characters and convert to lowercase
    text = re.sub(r'[^a-zA-Z\s]', '', text)  # Retain only alphabets and spaces
    text = text.lower().strip()             # Convert to lowercase and strip leading/trailing spaces
    text = text.replace(" ", "_")           # Replace spaces with underscores
    
    # Truncate or handle empty text
    truncated_text = text[:25] if len(text) > 25 else text if len(text) > 0 else "empty"
    
    # Generate a random string for uniqueness
    random_string = uuid.uuid4().hex[:8].upper()
    
    # Construct the file name
    file_name = f"{image_folder}/{truncated_text}_{random_string}.png"
    return file_name
    
def logo_remover(image_path):
    image = cv2.imread(image_path)
    image = cv2.resize(image, (1280, 1280))  # Resize image if needed

    # Extract watermark and location
    first_crop, first_location = extract_watermark(image, 0.50, 0.50, 0)
    watermark, location = extract_watermark(first_crop, 0.12, 0.26, 27) #height, side, margin

    
    # Create black and white masks
    mask1 = create_mask(first_crop, "black")
    mask2 = create_mask(watermark, "white")
    combined_mask = place_inpainted_back(mask1, mask2, location)

    # Save temporary files
    input_image = "./input/temp.png"
    input_mask = "./input/temp_mask.png"
    # temp_image = cv2.resize(first_crop, (512, 512))
    temp_image=first_crop
    cv2.imwrite(input_image, temp_image)
    # temp_mask = cv2.resize(combined_mask, (512, 512))
    temp_mask=combined_mask
    cv2.imwrite(input_mask, temp_mask)

    
    clean_image_path = process_image_with_model(input_image, input_mask, tester)
    
    # Check if the image was loaded correctly
    if clean_image_path is None:
        print(f"Failed to load image: {clean_image_path}")
        return  # Or handle the error accordingly
    clean_image = cv2.imread(clean_image_path)
    clean_image = cv2.resize(clean_image, (combined_mask.shape[1], combined_mask.shape[0]))
    result_image = place_inpainted_back(image, clean_image, first_location)
    save_path=img_file_name(image_path)
    cv2.imwrite(save_path, result_image)
    return save_path





# Define a function to handle the image editing and return the final result
def process_and_return(im):
    global tester
    # Save the composite image (base image) and mask to files
    base_image_path = "base_image.png"
    mask_image_path = "mask_image.png"
    
    # Save the composite image (base image)
    imageio.imwrite(base_image_path, im["composite"])
    
    # Extract the alpha channel (mask)
    alpha_channel = im["layers"][0][:, :, 3]
    
    # Create the mask: white (255) where drawn, black (0) elsewhere
    mask = np.zeros_like(alpha_channel, dtype=np.uint8)
    mask[alpha_channel > 0] = 255  # Set drawn areas to white (255)
    
    # Save the mask image
    imageio.imwrite(mask_image_path, mask)
    # Process the images using the inpainting model
    final_result = process_image_with_model(base_image_path, mask_image_path,tester)
    
    # Return the processed image
    return final_result

def ui_3():
    # Create a Gradio app
    with gr.Blocks() as demo:
        with gr.Row():
            # Create an ImageEditor component for uploading and editing the image
            im = gr.ImageEditor(
                type="numpy",
                canvas_size=(1, 1),  # Use canvas_size instead of crop_size
                layers=True,  # Allow layers in the editor
                transforms=["crop"],  # Allow cropping
                format="png",  # Save images in PNG format
                label="Base Image",
                show_label=True
            )
            # Create an Image component to display the processed result
            im2 = gr.Image(label="Processed Image", show_label=True)

        # Create a Button to trigger the image processing
        btn = gr.Button("Process Image")
        
        # Define an event listener to trigger the image processing when the button is clicked
        btn.click(process_and_return, inputs=im, outputs=im2)  # Output processed image
    return demo
# def handle_pil_image(image):
    
#     logo_remover(image)
    

def ui_1():
    test_examples=[["./input/image.jpg"]]
    gradio_input=[gr.Image(label='Upload an Image',type="filepath")]
    gradio_Output=[gr.Image(label='Display Image')]
    gradio_interface = gr.Interface(fn=logo_remover, inputs=gradio_input,outputs=gradio_Output ,
                                title="Meta Watermark Remover For Image",
                                examples=test_examples)
    return gradio_interface
from PIL import Image
import zipfile

def make_zip(image_list):
    zip_path = f"./temp/images/{uuid.uuid4().hex[:6]}.zip"
    with zipfile.ZipFile(zip_path, 'w') as zipf:
        for image in image_list:
            zipf.write(image, os.path.basename(image))
    return zip_path

def handle_multiple_files(image_files):
    image_list = []
    if len(image_files) == 1:
        saved_path=logo_remover(image_files[0])
        return saved_path
    else:
        for image_path in image_files:
            saved_path=logo_remover(image_path)
            image_list.append(saved_path)
        zip_path = make_zip(image_list)
        return zip_path
    

    
def ui_2():
    gradio_multiple_images = gr.Interface(
    handle_multiple_files,
    [gr.File(type='filepath', file_count='multiple',label='Upload Images')],
    [gr.File(label='Download File')],
    title='Meta Watermark Remover For Bulk Images',
    cache_examples=True
    )
    return gradio_multiple_images

# Load and process the inpainting model
tester = load_inpainting_model()
image_folder="./temp/images"
if not os.path.exists(image_folder):
    os.makedirs(image_folder)
    
# Launch the Gradio app
if __name__ == "__main__":
    demo2 = ui_1()
    demo3 = ui_2()
    demo1=ui_3()
    demo=gr.TabbedInterface([demo1,demo2,demo3], title="Meta Watermark Remover",tab_names=["Manual Remove","Meta Single Image","Meta Bulk Images"])
    demo.launch(show_error=True)