Update app.py

app.py

@@ -1,232 +1,245 @@
-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)
+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:
+        gr.Markdown("Manually Select the area.")
+        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 Single 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)
+    
+import click
+@click.command()
+@click.option("--debug", is_flag=True, default=False, help="Enable debug mode.")
+@click.option("--share", is_flag=True, default=False, help="Enable sharing of the interface.")
+
+def main(debug, share):
+    demo1 = ui_1()
+    demo2 = ui_2()
+    demo3=ui_3()
+    demo=gr.TabbedInterface([demo1,demo2,demo3], title="Meta Watermark Remover",tab_names=["Meta Single Image","Meta Bulk Images","Manual Remove"])
+    demo.queue().launch(debug=debug, share=share)#,server_port=9000)
+    #Run on local network
+    # laptop_ip="192.168.0.30"
+    # port=8080
+    # demo.queue().launch(debug=debug, share=share,server_name=laptop_ip,server_port=port)
+
+if __name__ == "__main__":
+    main()    
+