Update app.py

app.py

@@ -22,32 +22,33 @@ 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.astype(np.uint8))
-    resized_image = resize_image(orig_image)
-    im_tensor = torch.tensor(np.array(resized_image), dtype=torch.float32).permute(2, 0, 1).unsqueeze(0) / 255.0
+def process(image):
+    # prepare input
+    orig_image = Image.fromarray(image).convert("RGBA")
+    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
     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()
-    
-    # 배경 제거 모델 추론
+
+    # inference
     with torch.no_grad():
-        output = net(im_tensor)
-    output = F.interpolate(output[0][0], size=orig_image.size[::-1], mode='bilinear', align_corners=False)
-    output = torch.sigmoid(output).cpu().numpy().squeeze()
-    mask = (output * 255).astype(np.uint8)
-    new_im = Image.fromarray(mask).convert("L")
+        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_image = Image.new("RGBA", orig_image.size)
-    orig_image_rgba = orig_image.convert("RGBA")
-    result_image.paste(orig_image_rgba, mask=new_im)
+    # 마스크를 RGBA 이미지로 변환하고, 마스크의 알파 채널을 사용
+    mask_image = Image.fromarray(mask).convert("L")
+    final_image = Image.new("RGBA", orig_image.size)
+    final_image.paste(orig_image, mask=mask_image)
 
-    # 배경 이미지가 제공된 경우, 합성
-    if background_image is not None:
-        result_image = merge_images(background_image, result_image)
+    return final_image
 
-    return result_image
 
 def merge_images(background_image, foreground_image):
     """