Update app.py

app.py

@@ -4,8 +4,11 @@ import torch.nn.functional as F
 from torchvision.transforms.functional import normalize
 from huggingface_hub import hf_hub_download
 import gradio as gr
+from gradio_imageslider import ImageSlider
 from briarmbg import BriaRMBG
+import PIL
 from PIL import Image
+from typing import Tuple
 
 # 모델 초기화 및 로드
 net = BriaRMBG()
@@ -24,7 +27,7 @@ def resize_image(image, model_input_size=(1024, 1024)):
 
 def process(image, background_image=None):
     # 이미지 준비
-    orig_image = Image.fromarray(image).convert("RGBA")
+    orig_image = Image.fromarray(image).convert("RGB")
     w, h = orig_im_size = orig_image.size
     image = resize_image(orig_image)
     im_np = np.array(image)
@@ -37,30 +40,33 @@ def process(image, background_image=None):
     with torch.no_grad():
         result = net(im_tensor)
     
-# 후처리
+    # 후처리
     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 사이의 값으로 변환
+    mask = (result * 255).byte().cpu().numpy()
 
-# mask 배열이 예상대로 2차원인지 확인하고, 아니라면 조정합니다.
+    # mask 배열이 예상대로 2차원인지 확인하고, 아니라면 조정
     if mask.ndim > 2:
-        mask = mask.squeeze()  # 차원 축소
+        mask = mask.squeeze()
 
-# mask 배열을 명확히 uint8로 변환합니다.
+    # mask 배열을 명확히 uint8로 변환
     mask = mask.astype(np.uint8)
 
-# mask를 PIL 이미지로 변환
-    mask_image = Image.fromarray(mask, 'L')  # 'L' 모드는 그레이스케일 이미지를 나타냅니다.
-    
-    final_image = Image.new("RGBA", orig_image.size)
-    final_image.paste(orig_image, mask=mask_image)
+    # 마스크를 알파 채널로 사용하여 최종 이미지 생성
+    orig_image = orig_image.convert("RGBA")
+    final_image = Image.new("RGBA", orig_image.size, (0, 0, 0, 0))
+    mask_image = Image.fromarray(mask, mode='L')
+    foreground_image = Image.composite(orig_image, final_image, mask_image)
 
     # 선택적 배경 이미지 처리
-    if background_image is not None:
-        final_image = merge_images(background_image, final_image)
+    if background_image:
+        final_image = merge_images(background_image, foreground_image)
+    else:
+        final_image = foreground_image
 
     return final_image
 
+
 def merge_images(background_image, foreground_image):
     """
     배경 이미지에 배경이 제거된 이미지를 투명하게 삽입합니다.