import gradio as gr
import os
import cv2
from rembg import new_session, remove
from PIL import Image
from io import BytesIO

def inference(file, mask, model, alpha_influence, segmentation_strength, smoothing):
    im = cv2.imread(file, cv2.IMREAD_COLOR)
    cv2.imwrite(os.path.join("input.png"), im)

    input_path = 'input.png'
    output_path = 'output.png'

    with open(input_path, 'rb') as i:
        with open(output_path, 'wb') as o:
            input = i.read()
            output = remove(
                input, 
                only_mask=(True if mask == "Mask only" else False),
                alpha_matting=True,  # Habilitar el modo alpha matting
                alpha_matting_foreground_threshold=alpha_influence,  # Control de influencia del canal alfa
                alpha_matting_background_threshold=1 - alpha_influence,  # Control del canal alfa para el fondo
                alpha_matting_erode_size=int(segmentation_strength * 10),  # Control de fuerza de segmentación
                alpha_matting_smoothing=smoothing,  # Control de suavizado de bordes de la segmentación
                session=new_session(model)
            )

            o.write(output)

    return Image.open(BytesIO(output))

title = "RemBG"
description = "Gradio demo for RemBG. To use it, simply upload your image and adjust the alpha influence, segmentation strength, and smoothing."
article = "<p style='text-align: center;'><a href='https://github.com/danielgatis/rembg' target='_blank'>Github Repo</a></p>"

gr.Interface(
    inference, 
    [
        gr.inputs.Image(type="filepath", label="Input"),
        gr.inputs.Radio(
            [
                "Default", 
                "Mask only"
            ], 
            type="value",
            default="Default",
            label="Choices"
        ),
        gr.inputs.Dropdown([
            "u2net", 
            "u2netp", 
            "u2net_human_seg", 
            "u2net_cloth_seg", 
            "silueta",
            "isnet-general-use",
            "isnet-anime",
            "sam",
        ], 
        type="value",
        default="isnet-general-use",
        label="Models"
        ),
        gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.5, label="Alpha Influence"),
        gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.5, label="Segmentation Strength"),
        gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.5, label="Smoothing"),
    ], 
    gr.outputs.Image(type="PIL", label="Output"),
    title=title,
    description=description,
    article=article,
    examples=[["lion.png", "Default", "u2net", 0.5, 0.5, 0.5], ["girl.jpg", "Default", "u2net", 0.5, 0.5, 0.5], ["anime-girl.jpg", "Default", "isnet-anime", 0.5, 0.5, 0.5]],
    enable_queue=True
).launch()