import torch
from PIL import Image
import numpy as np
from transformers import U2NetForImageMatting, U2NetProcessor
import gradio as gr

# U-2-Net 모델과 프로세서 초기화
model = U2NetForImageMatting.from_pretrained("u2net")
processor = U2NetProcessor.from_pretrained("u2net")

def remove_background(image):
    # 이미지를 모델에 맞게 처리
    inputs = processor(image, return_tensors="pt")
    # 모델을 사용하여 이미지에서 배경 제거
    outputs = model(**inputs)
    # 마스크 추출 및 후처리
    mask = outputs.matting_mask[0].detach().cpu().numpy()
    mask = (mask * 255).astype(np.uint8)
    mask = Image.fromarray(mask).resize(image.size)
    # 원본 이미지와 마스크를 결합하여 배경이 제거된 이미지 생성
    empty = Image.new("RGBA", image.size)
    final_image = Image.composite(image.convert("RGBA"), empty, mask)
    return final_image

# Gradio 인터페이스 설정
iface = gr.Interface(fn=remove_background, 
                     inputs=gr.inputs.Image(type="pil", label="Upload Image"), 
                     outputs="image",
                     title="Background Removal with U-2-Net",
                     description="Upload an image to remove its background.")

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