app.py

import numpy as np
import torch
import torch.nn.functional as F
from torchvision.transforms.functional import normalize
from huggingface_hub import hf_hub_download
import gradio as gr
from briarmbg import BriaRMBG
from PIL import Image

# 모델 초기화 및 로드
net = BriaRMBG()
model_path = hf_hub_download("briaai/RMBG-1.4", 'model.pth')
if torch.cuda.is_available():
    net.load_state_dict(torch.load(model_path))
    net = net.cuda()
else:
    net.load_state_dict(torch.load(model_path, map_location="cpu"))
net.eval()

def resize_image(image, model_input_size=(1024, 1024)):
    image = image.convert('RGB')
    image = image.resize(model_input_size, Image.BILINEAR)
    return image

def process(image, background_image=None):
    # prepare input
    orig_image = Image.fromarray(image).convert("RGBA")
    w, h = orig_im_size = orig_image.size
    image = resize_image(orig_image)
    im_np = np.array(image)
    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0) / 255.0
    im_tensor = normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
    if torch.cuda.is_available():
        im_tensor = im_tensor.cuda()

    # inference
    with torch.no_grad():
        result = net(im_tensor)
    
    # post process
    result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode='bilinear', align_corners=False), 0)
    result = torch.sigmoid(result)  # 마스크의 투명도를 결정하는 확률 값으로 변환
    mask = (result * 255).byte().cpu().numpy()  # 마스크를 0~255 사이의 값으로 변환

    # 마스크를 RGBA 이미지로 변환하고, 마스크의 알파 채널을 사용
    mask_image = Image.fromarray(mask).convert("L")
    final_image = Image.new("RGBA", orig_image.size)
    final_image.paste(orig_image, mask=mask_image)

    return new_im  # or any other variable you wish to return


def merge_images(background_image, foreground_image):
    """
    배경 이미지에 배경이 제거된 이미지를 투명하게 삽입합니다.
    배경이 제거된 이미지는 배경 이미지 중앙에 30% 크기로 축소되어 삽입됩니다.
    """
    # 이미 RGBA 모드로 변환된 이미지를 직접 사용합니다.
    background = background_image.convert("RGBA")
    foreground = foreground_image.convert("RGBA")
    
    # 전경 이미지를 배경 이미지의 30% 크기로 조정
    scale_factor = 0.3
    new_size = (int(background.width * scale_factor), int(foreground.height * background.width / foreground.width * scale_factor))
    foreground_resized = foreground.resize(new_size, Image.Resampling.LANCZOS)
    
    # 전경 이미지를 배경 이미지의 가운데에 위치시키기 위한 좌표 계산
    x = (background.width - new_size[0]) // 2
    y = (background.height - new_size[1]) // 2
    
    # 배경 이미지 위에 전경 이미지를 붙임
    background.paste(foreground_resized, (x, y), foreground_resized)
    
    return background


demo = gr.Interface(
    fn=process,
    inputs=[
        gr.Image(type="numpy", label="Image to remove background"),  # 첫 번째 인자
        gr.Image(type="pil", label="Optional: Background Image", optional=True)  # 두 번째 인자, 선택적
    ],
    outputs="image",
    title=title,
    description=description
)