import gradio as gr
from rembg import remove
from PIL import Image
import logging
import os

# 로깅 설정
logging.basicConfig(level=logging.INFO)

def remove_background(input_image):
    try:
        logging.info("Processing image.")
        
        if isinstance(input_image, str):
            input_img = Image.open(input_image)
        else:
            input_img = input_image

        # 배경 제거 (파란색 배경으로 설정)
        output_img = remove(input_img)
        background = Image.new("RGBA", output_img.size, (255, 255, 255, 1))  # 흰색색 배경
        output_img = Image.alpha_composite(background, output_img)
        
        logging.info("Background removed successfully.")
        return output_img
    except Exception as e:
        logging.error(f"Error removing background: {e}")
        return f"에러가 발생했습니다: {str(e)}"

def clear_images():
    return None, None

# 예제 이미지 경로 설정
example_images = [
    "./example1.jpg",
    "./example2.jpg"
]

# 예제 이미지 경로 확인
for img_path in example_images:
    if not os.path.exists(img_path):
        logging.warning(f"Example image {img_path} does not exist. Please provide valid paths.")

# Gradio 인터페이스 구성
with gr.Blocks(css="""
    #btn-remove {
        background-color: #4CAF50; /* 초록색 */
        color: black;
    }
    #btn-reset {
        background-color: #FF9800; /* 주황색으로 변경 */
        color: black;
    }
""") as demo:
    gr.Markdown("# 이미지 배경 제거기")
    gr.Markdown("이미지를 업로드하면 배경이 자동으로 제거되고 흰색 배경으로 대체됩니다.")

    with gr.Row():
        input_image = gr.Image(
            type="pil", 
            label="입력 이미지",
            interactive=True
        )
        output_image = gr.Image(
            type="pil",
            label="배경 제거된 이미지",
            interactive=False
        )
    
    with gr.Row():
        btn_remove = gr.Button("배경 제거", elem_id="btn-remove")
        btn_remove.click(fn=remove_background, inputs=input_image, outputs=output_image)
        
        btn_reset = gr.Button("초기화", elem_id="btn-reset")
        btn_reset.click(fn=clear_images, inputs=None, outputs=[input_image, output_image])
    
    gr.Examples(
        examples=[[img] for img in example_images],
        inputs=input_image,
        label="예제 이미지"
    )

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