import gradio as gr
import rembg
from rembg import remove, new_session
from PIL import Image
import numpy as np
import logging
import time

# Set up logging
logging.basicConfig(level=logging.INFO)
# Log the version of rembg
logging.info(f"rembg version: {rembg.__version__}")
# Define model options with separate names and descriptions
MODEL_OPTIONS = {
    "": "Select a model",
    "u2net": "A pre-trained model for general use cases (default)",
    "isnet-general-use": "A new pre-trained model for general use cases",
    "isnet-anime": "High-accuracy segmentation for anime characters",
    "silueta": "A reduced-size version of u2net (43MB)",
    # "sam_prompt": "A pre-trained model for any use case (Segment Anything Model)", # remove sam model
    "unet": "Lightweight version of u2net model",
    "u2netp": "A lightweight version of u2net model",
    "u2net_human_seg": "A pre-trained model for human segmentation",
    "u2net_cloth_seg": "A pre-trained model for cloth parsing in human portraits",
}

def hex_to_rgba(hex_color):
    hex_color = hex_color.lstrip('#')
    if len(hex_color) == 6:
        hex_color += 'FF'  # Add full opacity if no alpha is provided
    return tuple(int(hex_color[i:i + 2], 16) for i in (0, 2, 4, 6))

def remove_background(input_image, bg_color, model_choice, alpha_matting, post_process_mask, only_mask):
    try:
        # Extract the model name from the choice
        model_name = model_choice.split(' | ')[0] if model_choice else ""
        
        # Set up the session with the chosen model, or None if no model is selected
        session = new_session(model_name) if model_name else None
        
        # Convert hex color to RGBA tuple
        bg_color_rgba = hex_to_rgba(bg_color) if bg_color else None
        
        # Prepare additional options
        remove_kwargs = {
            "session": session,
            "bgcolor": bg_color_rgba if bg_color_rgba else None,
            "alpha_matting": alpha_matting,
            "post_process_mask": post_process_mask,
            "only_mask": only_mask
        }

        # Add alpha matting parameters if enabled
        if alpha_matting:
            remove_kwargs.update({
                "alpha_matting": True,
                "alpha_matting_foreground_threshold": 270,
                "alpha_matting_background_threshold": 20,
                "alpha_matting_erode_size": 11
            })
        
        logging.info(f'Model name={model_name}')
        logging.info(remove_kwargs)

        # Convert PIL Image to numpy array
        input_array = np.array(input_image)

        # Use the remove function
        if session or bg_color_rgba:
            output_array = remove(input_array, **{k: v for k, v in remove_kwargs.items() if v is not None})
        else:
            output_array = remove(input_array)  # Use the default remove function

        logging.info("Background removed")

        # Convert numpy array back to PIL Image
        output_image = Image.fromarray(output_array)

        # Convert to RGB mode if necessary
        if not only_mask and output_image.mode != 'RGB':
            output_image = output_image.convert('RGB')
            logging.info("Converted to RGB mode")

        # Save the output image to a temporary file
        # Generate a unique timestamp for the output file name
        timestamp = time.strftime("%Y%m%d-%H%M%S")
        output_path = f"output_remove_background_{timestamp}.png"
        output_image.save(output_path)
        logging.info(f"Saved output image {output_path}")

        return output_image, output_path

    except Exception as e:
        logging.error(f"An error occurred: {e}")
        return None, None

examples = [

    ['scifi_man1.jpg']
]
# Gradio interface
iface = gr.Interface(
    fn=remove_background,
    inputs=[
        gr.Image(type="pil"),
        gr.ColorPicker(label="Background Color", value=None),
        gr.Dropdown(
            choices=[""] + [f"{k} | {v}" for k, v in MODEL_OPTIONS.items() if k != ""],
            label="Model Selection",
            value="",
            type="value"
        ),
        gr.Checkbox(label="Enable Alpha Matting", value=False),
        gr.Checkbox(label="Post-Process Mask (post process the mask to get better results)", value=False),
        gr.Checkbox(label="Only Return Mask ", value=False)
    ],
    outputs=[
        gr.Image(type="pil", label="Output Image"), 
        gr.File(label="Download the output image")
    ],
    examples=examples,
    title="Advanced Background Remover v2.7",
    description="Upload an image to remove the background. Customize the result with different options, including background color, model selection, alpha matting, and more.",
    allow_flagging="never",
)

if __name__ == "__main__":
    
    iface.launch()