import gradio as gr
from rembg import remove
from PIL import Image
import io

def merge_images(foreground, background):
    # foreground와 background를 바이트 배열로 변환
    foreground_bytes = io.BytesIO()
    foreground.save(foreground_bytes, format="PNG")
    foreground_bytes.seek(0)  # 포인터를 처음으로 이동
    background_bytes = io.BytesIO()
    background.save(background_bytes, format="PNG")
    background_bytes.seek(0)  # 포인터를 처음으로 이동
    
    # 배경 제거
    foreground_no_bg_bytes = remove(foreground_bytes.read())
    
    # 바이트 데이터를 BytesIO 스트림으로 변환하여 PIL 이미지로 로드
    foreground_no_bg = Image.open(io.BytesIO(foreground_no_bg_bytes)).convert("RGBA")
    background_image = Image.open(background_bytes).convert("RGBA")
    
    # 배경 이미지 크기 조정
    background_image = background_image.resize(foreground_no_bg.size)
    
    # 메인 이미지와 배경 이미지 합치기
    merged_image = Image.alpha_composite(background_image, foreground_no_bg)
    
    # 결과 이미지를 PIL 이미지 객체로 반환
    return merged_image

with gr.Blocks() as demo:
    with gr.Row():
        foreground_input = gr.Image(type="pil", label="Upload your main image")
        background_input = gr.Image(type="pil", label="Upload your background image")
    image_output = gr.Image(type="pil", label="Merged Image")
    
    btn = gr.Button("Merge Images")
    btn.click(fn=merge_images, inputs=[foreground_input, background_input], outputs=image_output)

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