Upload 4 files

app.py

@@ -0,0 +1,232 @@
+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)

config.py

@@ -0,0 +1,22 @@
+import torch
+
+# GENERIC
+GPU_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+INIMAGE = "./input/image.jpg"
+MASKIMAGE = "./input/mask.jpg"
+OUTIMAGE = "./output/inpainted_img.png"
+RESIZE_TO = (512, 512)
+CUDA = True if torch.cuda.is_available() else False
+
+# DEEPFILLv2
+DEEPFILL_MODEL_PATH = "./model/deepfillv2_WGAN.pth"
+GPU_ID = -1
+INIT_TYPE = "xavier"
+INIT_GAIN = 0.02
+PAD_TYPE = "zero"
+IN_CHANNELS = 4
+OUT_CHANNELS = 3
+LATENT_CHANNELS = 48
+ACTIVATION = "elu"
+NORM = "in"
+NUM_WORKERS = 0

create_mask.py

@@ -0,0 +1,84 @@
+import numpy as np
+import cv2
+from config import RESIZE_TO, MASKIMAGE
+
+# image mask
+
+# free form mask
+# bbox mask
+
+
+def create_ff_mask():
+    config = {
+        "img_shape": list(RESIZE_TO),
+        "mv": 15,
+        "ma": 4.0,
+        "ml": 40,
+        "mbw": 5,
+    }
+
+    h, w = config["img_shape"]
+    mask = np.zeros((h, w))
+    num_v = np.random.randint(config["mv"])
+
+    for i in range(num_v):
+        start_x = np.random.randint(w)
+        start_y = np.random.randint(h)
+        for j in range(1 + np.random.randint(5)):
+            angle = 0.01 + np.random.randint(config["ma"])
+            if i % 2 == 0:
+                angle = 2 * 3.1415926 - angle
+            length = 10 + np.random.randint(config["ml"])
+            brush_w = 5 + np.random.randint(config["mbw"])
+            end_x = (start_x + length * np.sin(angle)).astype(np.int32)
+            end_y = (start_y + length * np.cos(angle)).astype(np.int32)
+
+            cv2.line(mask, (start_y, start_x), (end_y, end_x), 255.0, brush_w)
+            start_x, start_y = end_x, end_y
+
+    mask = mask.astype(np.uint8)
+    cv2.imwrite(MASKIMAGE, mask)
+
+
+def create_bbox_mask():
+    shape = list(RESIZE_TO)
+    margin = [10, 10]
+    bbox_shape = [30, 30]
+
+    def random_bbox(shape, margin, bbox_shape):
+        """Generate a random tlhw with configuration.
+        Args:
+            config: Config should have configuration including IMG_SHAPES,
+                    VERTICAL_MARGIN, HEIGHT, HORIZONTAL_MARGIN, WIDTH.
+        Returns:
+            tuple: (top, left, height, width)
+        """
+        img_height, img_width = shape
+        height, width = bbox_shape
+        ver_margin, hor_margin = margin
+        maxt = img_height - ver_margin - height
+        maxl = img_width - hor_margin - width
+        t = np.random.randint(low=ver_margin, high=maxt)
+        l = np.random.randint(low=hor_margin, high=maxl)
+        h = height
+        w = width
+        return (t, l, h, w)
+
+    bboxs = []
+    for i in range(20):
+        bbox = random_bbox(shape, margin, bbox_shape)
+        bboxs.append(bbox)
+
+    height, width = shape
+    mask = np.zeros((height, width), np.float32)
+    # print(mask.shape)
+    for bbox in bboxs:
+        h = int(bbox[2] * 0.1) + np.random.randint(int(bbox[2] * 0.2 + 1))
+        w = int(bbox[3] * 0.1) + np.random.randint(int(bbox[3] * 0.2) + 1)
+        mask[
+            (bbox[0] + h) : (bbox[0] + bbox[2] - h),
+            (bbox[1] + w) : (bbox[1] + bbox[3] - w),
+        ] = 255.0
+
+    mask = mask.astype(np.uint8)
+    cv2.imwrite(MASKIMAGE, mask)

inpaint.py

@@ -0,0 +1,121 @@
+import os
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from types import SimpleNamespace
+from deepfillv2 import test_dataset, utils
+from config import *
+
+class InpaintingTester:
+    def __init__(self, save_path, resize_to=None):
+        if resize_to is None:
+            resize_to = RESIZE_TO
+        self.save_path = save_path
+        self.setsize = resize_to
+
+        # Build the generator network
+        opt = SimpleNamespace(
+            pad_type=PAD_TYPE,
+            in_channels=IN_CHANNELS,
+            out_channels=OUT_CHANNELS,
+            latent_channels=LATENT_CHANNELS,
+            activation=ACTIVATION,
+            norm=NORM,
+            init_type=INIT_TYPE,
+            init_gain=INIT_GAIN,
+            use_cuda=CUDA,
+            gpu_device=GPU_DEVICE,
+        )
+
+        # Initialize generator (only once)
+        self.generator = utils.create_generator(opt).eval()
+
+        # Load pretrained model weights
+        # print("-- INPAINT: Loading Pretrained Model --")
+        self.load_model_generator(self.generator)
+
+        # Move the generator to GPU
+        self.generator = self.generator.to(GPU_DEVICE)
+
+    def load_model_generator(self, generator):
+        pretrained_dict = torch.load(
+            DEEPFILL_MODEL_PATH, map_location=torch.device(GPU_DEVICE), weights_only=True
+        )
+        generator.load_state_dict(pretrained_dict)
+
+    def process_image(self, in_image, mask_image, save_image_path):
+        # Initialize dataset and dataloader
+        trainset = test_dataset.InpaintDataset(in_image, mask_image, self.setsize)
+        dataloader = DataLoader(
+            trainset,
+            batch_size=1,
+            shuffle=False,
+            num_workers=8,
+            pin_memory=True,
+        )
+
+        # Testing loop for a single image
+        for batch_idx, (img, mask) in enumerate(dataloader):
+            img = img.to(GPU_DEVICE)
+            mask = mask.to(GPU_DEVICE)
+
+            # Generator output
+            with torch.no_grad():
+                first_out, second_out = self.generator(img, mask)
+
+            # Combine outputs with input
+            first_out_wholeimg = img * (1 - mask) + first_out * mask
+            second_out_wholeimg = img * (1 - mask) + second_out * mask
+
+            masked_img = img * (1 - mask) + mask
+            mask = torch.cat((mask, mask, mask), 1)
+            img_list = [second_out_wholeimg]
+            name_list = ["second_out"]
+
+            # Save the sample image
+            results_path = os.path.dirname(save_image_path)
+            if not os.path.exists(results_path):
+                os.makedirs(results_path)
+
+            utils.save_sample_png(
+                sample_folder=results_path,
+                sample_name=os.path.basename(save_image_path),
+                img_list=img_list,
+                name_list=name_list,
+                pixel_max_cnt=255,
+            )
+
+    def process_multiple_images(self, image_mask_pairs):
+        # Iterate through a list of image/mask pairs and save results
+        png_images=[]
+        for img_path, mask_path in image_mask_pairs:
+            try:
+                save_image_path = os.path.join(self.save_path, os.path.basename(img_path))
+                print(f"Processing: {img_path} and {mask_path}")
+                self.process_image(img_path, mask_path, save_image_path)
+                extention = os.path.splitext(save_image_path)[1]
+                save_at=save_image_path.replace(extention, ".png")
+                png_images.append(save_at)
+            except Exception as e:
+                if self.save_path in png_images:
+                    png_images.pop()
+                    png_images.append(None)
+                print(f"Error: {e}")
+        # print("-- All Inpainting is finished --")
+        return png_images
+
+# Main execution
+# if __name__ == "__main__":
+#     save_path = "./output"
+#     resize_to = None  # Default size from config
+
+#     # List of image and mask pairs
+#     image_mask_pairs = [
+#         ( "./input/image.jpg", "./input/mask.jpg"),
+#     ]
+
+#     tester = InpaintingTester(save_path, resize_to)
+    
+#     # Process multiple images using a loop
+#     results=tester.process_multiple_images(image_mask_pairs)
+#     print(results)