import gradio as gr
from PIL import Image
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

portrait_matting = pipeline(Tasks.portrait_matting, model='damo/cv_unet_image-matting')

def matting(image, use_color, color):
    result = portrait_matting(image)
    image_pil = Image.fromarray(result['output_img'])
    alpha_channel = image_pil.split()[3]

    # 마스크로 사용할 이미지 생성
    mask_image = Image.new("L", image_pil.size, 0)
    mask_image.paste(alpha_channel, alpha_channel)

    # new image with background
    if use_color:
        new_image = Image.new("RGBA", image_pil.size, color)
        new_image.paste(image_pil, (0,0), image_pil)
        new_image = new_image.convert("RGB")
    else:
        new_image = image_pil

    return [mask_image, new_image] 

def merge(org_image, add_image):
    add_image = add_image.resize(org_image.size)
    org_image.paste(add_image, (0,0), add_image)
    return [org_image]

with gr.Blocks() as demo:
    with gr.Tab(label="Portrait Matting"):
        with gr.Row():
            with gr.Column():
                image = gr.Image(height="40vh")
                with gr.Row():
                    use_color = gr.Checkbox(label="use background color", value=True, container=False)
                    color = gr.ColorPicker(info="background color",value="#89cff0", container=False)
                btn_matt= gr.Button()
            result_matt = gr.Gallery(height="40vh")
    with gr.Tab(label="Merge"):
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    org_image = gr.Image(label="background", type='pil', image_mode='RGBA', height="40vh") 
                    add_image = gr.Image(label="foreground", type='pil', image_mode='RGBA', height="40vh")
                btn_merge = gr.Button()
            result_merge = gr.Gallery(height="40vh")
    
    btn_matt.click(
        fn=matting,
        inputs=[image, use_color, color],
        outputs=[result_matt],
    )

    btn_merge.click(
        fn=merge,
        inputs=[org_image, add_image],
        outputs=[result_merge],
    )

demo.launch()