# 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()



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):
    """
    입력된 이미지에서 배경을 제거하는 함수
    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

# 예제 이미지 경로 설정
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: 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=[[img] for img in example_images],  # 리스트 형식으로 경로 전달
        inputs=input_image,
        label="예제 이미지"
    )

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