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

Automatically generated by Colab.

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

import gradio as gr
from PIL import Image
import numpy as np

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.
    """
    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)
    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="Hex Color to Remove", placeholder="#FFFFFF"),
        gr.Slider(minimum=0, maximum=255, value=40, step=1, label="Color Tolerance"),
        gr.ColorPicker(label="Pick a Color")  # Color picker remains
    ],
    outputs=[
        gr.Image(type="numpy", label="Processed Image")  # Removed output textbox
    ],
    title="Image Colour Remover",
    description="""
    Remove a specific colour from your image with precision!

    ### How to Use:
    1. Upload an image
    2. Choose the colour to remove
    3. Adjust the colour tolerance (higher numbers will remove similair colours to hex code entered)
    4. See the magic happen!
    """,
)

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