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

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

def remove_background(input_image):
    """
    입력된 이미지에서 배경을 제거하는 함수
    Parameters:
    input_image (PIL.Image 또는 str): 업로드된 이미지 파일의 객체 또는 파일 경로
    Returns:
    PIL.Image 또는 str: 배경이 제거된 이미지 또는 에러 메시지
    """
    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)
        logging.info("Background removed successfully.")
        
        # 투명도를 유지하기 위해 PNG 형식으로 변환
        return output_img.convert("RGBA")
    except Exception as e:
        logging.error(f"Error removing background: {e}")
        return f"에러가 발생했습니다: {str(e)}"

def clear_images():
    """입력 이미지와 출력 이미지를 모두 초기화하는 함수"""
    return None, None

# Gradio 인터페이스 구성
with gr.Blocks(css="""
    #btn-remove {
        background-color: #4CAF50; /* 초록색 */
        color: white;
    }
    #btn-reset {
        background-color: #f44336; /* 빨간색 */
        color: white;
    }
""") 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")  # elem_id로 버튼 스타일 지정
        btn_remove.click(fn=remove_background, inputs=input_image, outputs=output_image)
        
        # 초기화 버튼 추가
        btn_reset = gr.Button("초기화", elem_id="btn-reset")  # elem_id로 버튼 스타일 지정
        btn_reset.click(fn=clear_images, inputs=None, outputs=[input_image, output_image])
    
    # 예제 이미지
    gr.Examples(
        examples=[
            ["./example1.jpg"],
            ["./example2.jpg"]
        ],
        inputs=input_image,
        label="예제 이미지"
    )

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