Update app.py

app.py

@@ -4,12 +4,9 @@ import torch.nn.functional as F
 from torchvision.transforms.functional import normalize
 from huggingface_hub import hf_hub_download
 import gradio as gr
-from gradio_imageslider import ImageSlider
 from briarmbg import BriaRMBG
 import PIL
 from PIL import Image
-from typing import Tuple
-
 import os
 import requests
 from moviepy.editor import VideoFileClip
@@ -22,13 +19,10 @@ def search_pexels_images(query):
     response = requests.get(url, headers=headers)
     data = response.json()
     
-    # 고해상도 이미지 URL만 선택하여 리스트 생성
     images_urls = []
     for photo in data.get('photos', []):
-        # 'large2x' 해상도의 이미지가 제공되는 경우, 해당 URL 사용
         if 'src' in photo and 'large2x' in photo['src']:
             images_urls.append(photo['src']['large2x'])
-        # 'large2x' 해상도의 이미지가 없는 경우, 'large' 또는 'original'을 대체로 사용
         elif 'large' in photo['src']:
             images_urls.append(photo['src']['large'])
         elif 'original' in photo['src']:
@@ -36,36 +30,29 @@ def search_pexels_images(query):
 
     return images_urls
 
-
 def show_search_results(query):
     images_urls = search_pexels_images(query)
     return images_urls
-    
 
-net=BriaRMBG()
-# model_path = "./model1.pth"
+net = BriaRMBG()
 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):
     image = image.convert('RGB')
     model_input_size = (1024, 1024)
     image = image.resize(model_input_size, Image.BILINEAR)
     return image
 
-
 def process(image):
-    # 이미지가 numpy 배열인 경우에만 PIL.Image 객체로 변환
     if isinstance(image, np.ndarray):
         orig_image = Image.fromarray(image)
     else:
-        # 이미 PIL.Image.Image 객체인 경우, 변환 없이 사용
         orig_image = image
 
     w, h = orig_im_size = orig_image.size
@@ -78,48 +65,34 @@ def process(image):
     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)
 
     return 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])
-
+    positions = {
+        "Center": ((org_size[0] - add_size[0]) // 2, (org_size[1] - add_size[1]) // 2),
+        "Top Left": (0, 0),
+        "Top Center": ((org_size[0] - add_size[0]) // 2, 0),
+        "Top Right": (org_size[0] - add_size[0], 0),
+        "Center Left": (0, (org_size[1] - add_size[1]) // 2),
+        "Center Right": (org_size[0] - add_size[0], (org_size[1] - add_size[1]) // 2),
+        "Bottom Left": (0, org_size[1] - add_size[1]),
+        "Bottom Center": ((org_size[0] - add_size[0]) // 2, org_size[1] - add_size[1]),
+        "Bottom Right": (org_size[0] - add_size[0], org_size[1] - add_size[1])
+    }
+    return positions.get(position, positions['Center'])
 
 def merge(org_image, add_image, scale, position, display_size):
-    # 사용자가 선택한 디스플레이 크기에 따라 결과 이미지 크기 조절
     display_width, display_height = map(int, display_size.split('x'))
-    
-    # 이미지 병합 로직
     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)
@@ -129,54 +102,30 @@ def merge(org_image, add_image, scale, position, display_size):
     merged_image.paste(org_image, (0, 0))
     merged_image.paste(add_image, position, add_image)
     
-    # 결과 이미지 디스플레이 크기 조절
     final_image = merged_image.resize((display_width, display_height), Image.Resampling.LANCZOS)
     
     return final_image
 
-
-with gr.Blocks() as demo:
-    with gr.Tab("Background Removal"):
-        with gr.Column():
-            gr.Markdown("누끼따기의 왕 '누킹'(Nuking)")
-            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=400)  # 예시로 높이 조절
-            add_image = gr.Image(label="Foreground", type='pil', image_mode='RGBA', height=400)  # 예시로 높이 조절
-            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")
-            display_size = gr.Textbox(value="1024x768", label="Display Size (Width x Height)")
-            btn_merge = gr.Button("Merge Images")
-            result_merge = gr.Image()
-            
-            btn_merge.click(
-                fn=merge,
-                inputs=[org_image, add_image, scale, position, display_size],
-                outputs=result_merge,
-            )
-
-            
-    with gr.TabItem("Image Search"):
-        with gr.Column():
-            gr.Markdown("### FREE Image Search")
-            search_query = gr.Textbox(label="사진 검색")
-            search_btn = gr.Button("검색")
-            images_output = gr.Gallery(label="검색 결과 이미지")
-            search_btn.click(
-                fn=show_search_results,
-                inputs=search_query,
-                outputs=images_output
-            )
-
-demo.launch()
+css = """
+footer { 
+    visibility: hidden; 
+}
+"""
+
+with gr.Blocks(css=css) as demo:
+    with gr.Column():
+        org_image = gr.Image(label="Background", type='pil', image_mode='RGBA', height=400)
+        add_image = gr.Image(label="Foreground", type='pil', image_mode='RGBA', height=400)
+        scale = gr.Slider(minimum=10, maximum=200, step=1, value=100, label="Scale of Foreground Image (%)")
+        position = gr.Radio(choices=["Center", "Top Left", "Top Center", "Top Right", "Center Left", "Center Right", "Bottom Left", "Bottom Center", "Bottom Right"], value="Center", label="Position of Foreground Image")
+        display_size = gr.Textbox(value="1024x768", label="Display Size (Width x Height)")
+        btn_merge = gr.Button("Merge Images")
+        result_merge = gr.Image()
+        
+        btn_merge.click(
+            fn=merge,
+            inputs=[org_image, add_image, scale, position, display_size],
+            outputs=result_merge,
+        )
+
+demo.launch()