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):
    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
    new_image = Image.new("RGBA", image_pil.size, (207,195,181))
    new_image.paste(image_pil, (0,0), image_pil)
    new_image = new_image.convert("RGB")

    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()
                btn_matt= gr.Button()
            result_matt = gr.Gallery()
    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')
                    add_image = gr.Image(label="foreground", type='pil', image_mode='RGBA')
                btn_merge = gr.Button()
            result_merge = gr.Image()

    btn_matt.click(
        fn=matting,
        inputs=[image],
        outputs=[result_matt],
    )

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

demo.launch()