Update app.py

app.py

@@ -23,7 +23,7 @@ def resize_image(image, model_input_size=(1024, 1024)):
     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)
@@ -33,54 +33,38 @@ def process(image, background_image=None):
     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 사이의 값으로 변환
+    result = torch.sigmoid(result)
+    mask = (result * 255).byte().cpu().numpy()
 
-    # 마스크를 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
+    # 선택적 배경 이미지 처리
+    if background_image is not None:
+        final_image = merge_images(background_image, final_image)
 
+    return final_image
 
-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
-
+title = "Background Removal"
+description = "This is a demo for BRIA RMBG 1.4 using the BRIA RMBG-1.4 image matting model as backbone."
 
 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)  # 두 번째 인자, 선택적
+        gr.Image(type="numpy", label="Image to remove background"),
+        gr.Image(type="pil", label="Optional: Background Image")  # 'optional=True'는 Gradio 버전에 따라 필요 없을 수 있음
     ],
     outputs="image",
     title=title,
     description=description
 )
+
+if __name__ == "__main__":
+    demo.launch()