# -*- coding: utf-8 -*-
"""app code for remove.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1GoN403JoCKCLavITeCYoMzFXKctQ83en
"""

from IPython import get_ipython
from IPython.display import display
# %%
import gradio as gr
from PIL import Image
import numpy as np
import cv2  # Import OpenCV

def sanitize_hex_color(color):
    """
    Sanitize and validate hex color input.
    """
    if color.startswith('rgb'):
        rgb_vals = color.strip('rgb()').split(',')
        if len(rgb_vals) == 3:
            try:
                r, g, b = map(int, rgb_vals)
                return f'#{r:02x}{g:02x}{b:02x}'.upper()
            except ValueError:
                raise ValueError("Invalid RGB color format")
    color = color.strip().upper()
    color = color.lstrip('#')
    if len(color) != 6:
        raise ValueError("Invalid hex color. Use '#RRGGBB' format")
    try:
        int(color, 16)
    except ValueError:
        raise ValueError("Invalid hexadecimal color")
    return f'#{color}'

def remove_background_color(image, hex_color, tolerance):
    """
    Remove pixels of a specified hex color from an image,
    with edge detection for refinement.
    """
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    image = image.convert("RGBA")
    try:
        hex_color = sanitize_hex_color(hex_color)
        target_color = tuple(int(hex_color.lstrip('#')[i:i+2], 16) for i in (0, 2, 4))
    except ValueError as e:
        raise ValueError(str(e))
    data = np.array(image)
    mask = np.all(np.abs(data[:,:,:3] - target_color) <= tolerance, axis=2)

    # Edge detection using Canny
    edges = cv2.Canny(cv2.cvtColor(data, cv2.COLOR_RGBA2GRAY), 100, 200)
    # Adjust thresholds (100, 200) as needed

    # Refine mask based on edges
    kernel = np.ones((3, 3), np.uint8)
    dilated_edges = cv2.dilate(edges, kernel, iterations=1)
    mask = np.logical_or(mask, dilated_edges)

    alpha_channel = np.where(mask, 0, 255)
    data[:, :, 3] = alpha_channel
    return data
def color_remover_interface(image, hex_color, tolerance, _=None):
    """
    Gradio interface function for color removal.
    """
    try:
        if image is None:
            raise ValueError("Please upload an image")
        hex_color = sanitize_hex_color(hex_color)
        processed_image = remove_background_color(image, hex_color, tolerance)
        return processed_image  # Return only the processed image
    except Exception as e:
        gr.Warning(str(e))
        return image  # Return original image in case of error

# Deploy using Gradio Client
iface = gr.Interface(
    fn=color_remover_interface,
    inputs=[
        gr.Image(type="numpy", label="Upload Image"),
        gr.Textbox(label="Enter the hex code of the colour you want to remove", placeholder="#FFFFFF"),
        gr.Slider(minimum=0, maximum=255, value=40, step=1, label="Colour Tolerance, adjust this to be higher if background doesn't remove"),
        gr.ColorPicker(label="Pick a Colour to find its hex value")  # Color picker remains
    ],
    outputs=[
        gr.Image(type="pil", image_mode="RGBA", label="Processed Image", format='png')
    ],
    title="Background remover for any image",
    description="""
    Remove a specific colour from your image with precision!

    ### How to Use:
    1. Upload an image
    2. Choose the colour to remove (Enter the hex code, the colour picker at the bottom of the page will help find this hex code)
    3. Adjust the colour tolerance (higher numbers will remove more)
    4. See the magic happen!
    """,
)

iface.launch(share=True)  # Launch the interface with a shareable link