Update app.py

app.py

@@ -26,10 +26,9 @@ def resize_image(image, model_input_size=(1024, 1024)):
     return image
 
 
-
 def process(image, background_image=None):
     # 이미지 준비
-    orig_image = Image.fromarray(image).convert("RGBA")
+    orig_image = Image.fromarray(image).convert("RGB")
     w, h = orig_image.size
     resized_image = resize_image(orig_image)
     im_np = np.array(resized_image).astype(np.float32) / 255.0
@@ -45,24 +44,25 @@ def process(image, background_image=None):
     # 후처리
     result = torch.squeeze(F.interpolate(result[0][0], size=(h, w), mode='bilinear', align_corners=False), 0)
     result = torch.sigmoid(result)
-    mask = result.cpu().numpy()
-    
-    # 마스크 후처리 및 알파 채널로 변환
-    mask = (mask * 255).astype(np.uint8)
+    mask = (result * 255).byte().cpu().numpy()
+
     if mask.ndim > 2:
         mask = mask.squeeze()
-    
-    # 최종 이미지 생성: 마스크를 알파 채널로 사용
-    orig_array = np.array(orig_image)
-    final_array = np.concatenate((orig_array[:, :, :3], mask[:, :, None]), axis=-1)
-
-    # 선택적 배경 이미지와 합성
-    if background_image is not None:
-        background = background_image.convert("RGBA")
-        foreground = Image.fromarray(final_array, 'RGBA')
-        final_image = merge_images(background, foreground)
+
+    mask = mask.astype(np.uint8)
+
+    # 마스크를 알파 채널로 사용하여 최종 이미지 생성
+    final_image = Image.new("RGBA", orig_image.size)
+    orig_image.putalpha(Image.fromarray(mask, 'L'))
+
+    if background_image:
+        # 배경 이미지가 제공된 경우, 배경 이미지 크기 조정
+        background = background_image.convert("RGBA").resize(orig_image.size)
+        # 배경과 전경(알파 적용된 원본 이미지) 합성
+        final_image = Image.alpha_composite(background, orig_image)
     else:
-        final_image = Image.fromarray(final_array, 'RGBA')
+        # 배경 이미지가 없는 경우, 투명도가 적용된 원본 이미지를 최종 이미지로 사용
+        final_image = orig_image
 
     return final_image