Update app.py

app.py

@@ -25,13 +25,15 @@ def resize_image(image, model_input_size=(1024, 1024)):
     image = image.resize(model_input_size, Image.BILINEAR)
     return image
 
+
+
 def process(image, background_image=None):
     # 이미지 준비
-    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)
-    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0) / 255.0
+    orig_image = Image.fromarray(image).convert("RGBA")
+    w, h = orig_image.size
+    resized_image = resize_image(orig_image)
+    im_np = np.array(resized_image).astype(np.float32) / 255.0
+    im_tensor = torch.tensor(im_np).permute(2, 0, 1).unsqueeze(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()
@@ -39,34 +41,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()
-
-    # mask 배열이 예상대로 2차원인지 확인하고, 아니라면 조정
+    mask = result.cpu().numpy()
+    
+    # 마스크 후처리 및 알파 채널로 변환
+    mask = (mask * 255).astype(np.uint8)
     if mask.ndim > 2:
         mask = mask.squeeze()
+    
+    # 최종 이미지 생성: 마스크를 알파 채널로 사용
+    orig_array = np.array(orig_image)
+    final_array = np.concatenate((orig_array[:, :, :3], mask[:, :, None]), axis=-1)
 
-    # mask 배열을 명확히 uint8로 변환
-    mask = mask.astype(np.uint8)
-
-    # 마스크를 알파 채널로 사용하여 최종 이미지 생성
-    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:
-        final_image = merge_images(background_image, foreground_image)
+    # 선택적 배경 이미지와 합성
+    if background_image is not None:
+        background = background_image.convert("RGBA")
+        foreground = Image.fromarray(final_array, 'RGBA')
+        final_image = merge_images(background, foreground)
     else:
-        final_image = foreground_image
+        final_image = Image.fromarray(final_array, 'RGBA')
 
     return final_image
 
 
+
 def merge_images(background_image, foreground_image):
     """
     배경 이미지에 배경이 제거된 이미지를 투명하게 삽입합니다.