rm_bkg

app.py

@@ -8,6 +8,9 @@ import os
 import matplotlib.pyplot as plt
 import io
 from PIL import Image
+import rembg
+from typing import Any
+
 
 from huggingface_hub import hf_hub_download
 ckpt_path = hf_hub_download(repo_id="Viglong/OriNet", filename="celarge/dino_weight.pt", repo_type="model", cache_dir='./')
@@ -30,6 +33,68 @@ dino.load_state_dict(torch.load(ckpt_path, map_location='cpu'))
 print('weight loaded')
 val_preprocess   = AutoImageProcessor.from_pretrained(DINO_LARGE, cache_dir='./')
 
+def background_preprocess(input_image, do_remove_background):
+
+    rembg_session = rembg.new_session() if do_remove_background else None
+
+    if do_remove_background:
+        input_image = remove_background(input_image, rembg_session)
+        input_image = resize_foreground(input_image, 0.85)
+
+    return input_image
+
+def resize_foreground(
+    image: Image,
+    ratio: float,
+) -> Image:
+    image = np.array(image)
+    assert image.shape[-1] == 4
+    alpha = np.where(image[..., 3] > 0)
+    y1, y2, x1, x2 = (
+        alpha[0].min(),
+        alpha[0].max(),
+        alpha[1].min(),
+        alpha[1].max(),
+    )
+    # crop the foreground
+    fg = image[y1:y2, x1:x2]
+    # pad to square
+    size = max(fg.shape[0], fg.shape[1])
+    ph0, pw0 = (size - fg.shape[0]) // 2, (size - fg.shape[1]) // 2
+    ph1, pw1 = size - fg.shape[0] - ph0, size - fg.shape[1] - pw0
+    new_image = np.pad(
+        fg,
+        ((ph0, ph1), (pw0, pw1), (0, 0)),
+        mode="constant",
+        constant_values=((0, 0), (0, 0), (0, 0)),
+    )
+
+    # compute padding according to the ratio
+    new_size = int(new_image.shape[0] / ratio)
+    # pad to size, double side
+    ph0, pw0 = (new_size - size) // 2, (new_size - size) // 2
+    ph1, pw1 = new_size - size - ph0, new_size - size - pw0
+    new_image = np.pad(
+        new_image,
+        ((ph0, ph1), (pw0, pw1), (0, 0)),
+        mode="constant",
+        constant_values=((0, 0), (0, 0), (0, 0)),
+    )
+    new_image = Image.fromarray(new_image)
+    return new_image
+
+def remove_background(image: Image,
+    rembg_session: Any = None,
+    force: bool = False,
+    **rembg_kwargs,
+) -> Image:
+    do_remove = True
+    if image.mode == "RGBA" and image.getextrema()[3][0] < 255:
+        do_remove = False
+    do_remove = do_remove or force
+    if do_remove:
+        image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
+    return image
 
 def get_3angle(image):
     
@@ -80,10 +145,8 @@ def figure_to_img(fig):
         image = Image.open(buf).copy()
     return image
 
-# def generate_mutimodal(title, context, img):
-#     return f"Title:{title}\nContext:{context}\n...{img}"
-
-def generate_mutimodal(img):
+def infer_func(img, do_rm_bkg):
+    img = background_preprocess(img, do_rm_bkg)
     angles = get_3angle(img)
     
     fig, ax = plt.subplots(figsize=(8, 8))
@@ -123,9 +186,10 @@ def generate_mutimodal(img):
 
 server = gr.Interface(
     flagging_mode='never',
-    fn=generate_mutimodal, 
+    fn=infer_func, 
     inputs=[
-        gr.Image(height=512, width=512, label="upload your image")
+        gr.Image(height=512, width=512, label="upload your image"),
+        gr.Checkbox(label="Remove Background", value=True)
     ], 
     outputs=[
         gr.Image(height=512, width=512, label="result image"),

requirements.txt

@@ -5,4 +5,4 @@ pillow==10.2.0
 huggingface-hub==0.26.5
 gradio==5.9.0
 numpy==1.26.4
-
+rembg