import gradio as gr
from PIL import Image
from rembg import remove
import numpy as np
import random

# 배경 변경 함수
def change_background(image, background_color=None):
    """
    배경을 제거한 후 새로운 배경색을 삽입합니다.
    """
    # 배경 제거
    image_array = np.array(image)
    transparent_image = remove(image_array)

    # 새로운 배경 색 설정
    if background_color is None:
        background_color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))

    # 투명한 부분에 배경 색상 추가
    background = Image.new("RGB", image.size, background_color)
    background.paste(Image.fromarray(transparent_image).convert("RGB"), mask=Image.fromarray(transparent_image).split()[3])

    return background

# Gradio 인터페이스 함수
def process_image(image):
    # 4개의 배경 변경 결과 생성
    modified_images = [
        change_background(image, background_color=(255, 0, 0)),  # 빨간 배경
        change_background(image, background_color=(0, 255, 0)),  # 초록 배경
        change_background(image, background_color=(0, 0, 255)),  # 파란 배경
        change_background(image),  # 랜덤 배경
    ]
    return modified_images

# Gradio UI 구성
with gr.Blocks() as demo:
    gr.Markdown("# 이미지 배경 변경")
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="이미지 업로드")
        with gr.Column():
            output_gallery = gr.Gallery(label="결과")
    submit_button = gr.Button("배경 변경 실행")
    submit_button.click(process_image, inputs=[input_image], outputs=[output_gallery])

# 앱 실행
if __name__ == "__main__":
    demo.launch()