Update app.py

app.py

@@ -37,12 +37,21 @@ 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 = (result * 255).byte().cpu().numpy()  # 마스크를 0~255 사이의 값으로 변환
 
-    mask_image = Image.fromarray(mask).convert("L")
+# mask 배열이 예상대로 2차원인지 확인하고, 아니라면 조정합니다.
+    if mask.ndim > 2:
+        mask = mask.squeeze()  # 차원 축소
+
+# 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)
 
@@ -52,6 +61,31 @@ def process(image, background_image=None):
 
     return final_image
 
+def merge_images(background_image, foreground_image):
+    """
+    배경 이미지에 배경이 제거된 이미지를 투명하게 삽입합니다.
+    배경이 제거된 이미지는 배경 이미지 중앙에 30% 크기로 축소되어 삽입됩니다.
+    """
+    background = background_image.convert("RGBA")
+    foreground = foreground_image.convert("RGBA")
+    
+    # 전경 이미지를 배경 이미지의 30% 크기로 조정
+    scale_factor = 0.3
+    foreground_width = int(background.width * scale_factor)
+    foreground_height = int(foreground.height * foreground_width / foreground.width)
+    new_size = (foreground_width, foreground_height)
+    foreground_resized = foreground.resize(new_size, Image.Resampling.LANCZOS)
+    
+    # 전경 이미지를 배경 이미지의 가운데에 위치시키기 위한 좌표 계산
+    x = (background.width - foreground_width) // 2
+    y = (background.height - foreground_height) // 2
+    
+    # 배경 이미지 위에 전경 이미지를 붙임
+    background.paste(foreground_resized, (x, y), foreground_resized)
+    
+    return background
+
+
 title = "Background Removal"
 description = "This is a demo for BRIA RMBG 1.4 using the BRIA RMBG-1.4 image matting model as backbone."