Update app.py

app.py

@@ -10,109 +10,111 @@ import PIL
 from PIL import Image
 from typing import Tuple
 
-# 모델 초기화 및 로드
-net = BriaRMBG()
+net=BriaRMBG()
+# model_path = "./model1.pth"
 model_path = hf_hub_download("briaai/RMBG-1.4", 'model.pth')
 if torch.cuda.is_available():
     net.load_state_dict(torch.load(model_path))
-    net = net.cuda()
+    net=net.cuda()
 else:
-    net.load_state_dict(torch.load(model_path, map_location="cpu"))
-net.eval()
+    net.load_state_dict(torch.load(model_path,map_location="cpu"))
+net.eval() 
 
-def resize_image(image, model_input_size=(1024, 1024)):
+    
+def resize_image(image):
     image = image.convert('RGB')
+    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_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()
-
-    # 추론
-    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()
-
-    if mask.ndim > 2:
-        mask = mask.squeeze()
-
-    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 = orig_image
-
-    return final_image
-
-
+def process(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
+    # prepare input
+    orig_image = Image.fromarray(image)
+    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)
+    im_tensor = torch.unsqueeze(im_tensor,0)
+    im_tensor = torch.divide(im_tensor,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
+    result=net(im_tensor)
+    # post process
+    result = torch.squeeze(F.interpolate(result[0][0], size=(h,w), mode='bilinear') ,0)
+    ma = torch.max(result)
+    mi = torch.min(result)
+    result = (result-mi)/(ma-mi)    
+    # image to pil
+    im_array = (result*255).cpu().data.numpy().astype(np.uint8)
+    pil_im = Image.fromarray(np.squeeze(im_array))
+    # paste the mask on the original image
+    new_im = Image.new("RGBA", pil_im.size, (0,0,0,0))
+    new_im.paste(orig_image, mask=pil_im)
+    # new_orig_image = orig_image.convert('RGBA')
+
+    return new_im
+    # return [new_orig_image, new_im]
+
+def calculate_position(org_size, add_size, position):
+    if position == "상단 좌측":
+        return (0, 0)
+    elif position == "상단 가운데":
+        return ((org_size[0] - add_size[0]) // 2, 0)
+    elif position == "상단 우측":
+        return (org_size[0] - add_size[0], 0)
+    elif position == "중앙 좌측":
+        return (0, (org_size[1] - add_size[1]) // 2)
+    elif position == "중앙 가운데":
+        return ((org_size[0] - add_size[0]) // 2, (org_size[1] - add_size[1]) // 2)
+    elif position == "중앙 우측":
+        return (org_size[0] - add_size[0], (org_size[1] - add_size[1]) // 2)
+    elif position == "하단 좌측":
+        return (0, org_size[1] - add_size[1])
+    elif position == "하단 가운데":
+        return ((org_size[0] - add_size[0]) // 2, org_size[1] - add_size[1])
+    elif position == "하단 우측":
+        return (org_size[0] - add_size[0], org_size[1] - add_size[1])
+
+def merge(org_image, add_image, scale, position):
+    scale_percentage = scale / 100.0
+    new_size = (int(add_image.width * scale_percentage), int(add_image.height * scale_percentage))
+    add_image = add_image.resize(new_size, Image.Resampling.LANCZOS)
     
-    # 배경 이미지 위에 전경 이미지를 붙임
-    background.paste(foreground_resized, (x, y), foreground_resized)
+    position = calculate_position(org_image.size, add_image.size, position)
+    org_image.paste(add_image, position, add_image)
     
-    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."
-
-
-def add_blue_background(image):
-    # 배경 제거된 이미지에 푸른색 배경을 추가하는 함수
-    blue_background = Image.new("RGBA", image.size, "blue")
-    final_image = Image.alpha_composite(blue_background, image.convert("RGBA"))
-    return final_image
-
-
-inputs = "image"      # 이전: gr.inputs.Image(type="pil") -> 변경: "image"
-outputs = ["image"]   # 이전 방식에서 변경되었을 수 있는 출력 부분
-
-demo = gr.Interface(fn=process,
-                    inputs=inputs,
-                    outputs=outputs,
-                    title="Your Demo Title",
-                    description="A brief description of your app.")
-
-if __name__ == "__main__":
-    demo.launch()
-    
+    return org_image
+
+
+
+with gr.Blocks() as demo:
+    with gr.Tab("Background Removal"):
+        with gr.Column():
+            gr.Markdown("## BRIA RMBG 1.4")
+            gr.HTML('''
+              <p style="margin-bottom: 10px; font-size: 94%">
+                This is a demo for BRIA RMBG 1.4 that using
+                <a href="https://huggingface.co/briaai/RMBG-1.4" target="_blank">BRIA RMBG-1.4 image matting model</a> as backbone. 
+              </p>
+            ''')
+            input_image = gr.Image(type="pil")
+            output_image = gr.Image()
+            process_button = gr.Button("Remove Background")
+            process_button.click(fn=process, inputs=input_image, outputs=output_image)
+
+    with gr.Tab("Merge"):
+        with gr.Column():
+            org_image = gr.Image(label="Background", type='pil', image_mode='RGBA', height="80vh")
+            add_image = gr.Image(label="Foreground", type='pil', image_mode='RGBA', height="80vh")
+            scale = gr.Slider(minimum=10, maximum=200, step=1, value=100, label="Scale of Foreground Image (%)")
+            position = gr.Radio(choices=["중앙 가운데", "상단 좌측", "상단 가운데", "상단 우측", "중앙 좌측", "중앙 우측", "하단 좌측", "하단 가운데", "하단 우측"], value="중앙 가운데", label="Position of Foreground Image")
+            merge_button = gr.Button("Merge Images")
+            result_merge = gr.Image(height="80vh")
+            merge_button.click(fn=merge, inputs=[org_image, add_image, scale, position], outputs=result_merge)
+
+demo.launch()