Update app.py

app.py

@@ -3,27 +3,16 @@ import cv2
 import numpy as np
 import torch
 import gradio as gr
-import spaces
-
-from glob import glob
-from typing import Tuple
 
 from PIL import Image
-from gradio_imageslider import ImageSlider
 from transformers import AutoModelForImageSegmentation
 from torchvision import transforms
 
-import requests
-from io import BytesIO
-import zipfile
-
-
 torch.set_float32_matmul_precision('high')
 torch.jit.script = lambda f: f
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-### image_proc.py
 def refine_foreground(image, mask, r=90):
     if mask.size != image.size:
         mask = mask.resize(image.size)
@@ -33,15 +22,12 @@ def refine_foreground(image, mask, r=90):
     image_masked = Image.fromarray((estimated_foreground * 255.0).astype(np.uint8))
     return image_masked
 
-
 def FB_blur_fusion_foreground_estimator_2(image, alpha, r=90):
-    # Thanks to the source: https://github.com/Photoroom/fast-foreground-estimation
     alpha = alpha[:, :, None]
     F, blur_B = FB_blur_fusion_foreground_estimator(
         image, image, image, alpha, r)
     return FB_blur_fusion_foreground_estimator(image, F, blur_B, alpha, r=6)[0]
 
-
 def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
     if isinstance(image, Image.Image):
         image = np.array(image) / 255.0
@@ -57,9 +43,8 @@ def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
     F = np.clip(F, 0, 1)
     return F, blurred_B
 
-
 class ImagePreprocessor():
-    def __init__(self, resolution: Tuple[int, int] = (1024, 1024)) -> None:
+    def __init__(self, resolution=(1024, 1024)) -> None:
         self.transform_image = transforms.Compose([
             transforms.Resize(resolution),
             transforms.ToTensor(),
@@ -70,159 +55,47 @@ class ImagePreprocessor():
         image = self.transform_image(image)
         return image
 
-
-usage_to_weights_file = {
-    'General': 'BiRefNet',
-    'General-Lite': 'BiRefNet_lite',
-    'General-Lite-2K': 'BiRefNet_lite-2K',
-    'Matting': 'BiRefNet-matting',
-    'Portrait': 'BiRefNet-portrait',
-    'DIS': 'BiRefNet-DIS5K',
-    'HRSOD': 'BiRefNet-HRSOD',
-    'COD': 'BiRefNet-COD',
-    'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs',
-    'General-legacy': 'BiRefNet-legacy'
-}
-
-birefnet = AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file['General'])), trust_remote_code=True)
+birefnet = AutoModelForImageSegmentation.from_pretrained('zhengpeng7/BiRefNet-matting', trust_remote_code=True)
 birefnet.to(device)
 birefnet.eval()
 
+def predict(image):
+    if image is None:
+        raise gr.Error("Please upload an image.")
 
-@spaces.GPU
-def predict(images, resolution, weights_file):
-    assert (images is not None), 'AssertionError: images cannot be None.'
-
-    global birefnet
-    # Load BiRefNet with chosen weights
-    _weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else usage_to_weights_file['General']))
-    print('Using weights: {}.'.format(_weights_file))
-    birefnet = AutoModelForImageSegmentation.from_pretrained(_weights_file, trust_remote_code=True)
-    birefnet.to(device)
-    birefnet.eval()
-
-    try:
-        resolution = [int(int(reso)//32*32) for reso in resolution.strip().split('x')]
-    except:
-        resolution = (1024, 1024) if weights_file not in ['General-Lite-2K'] else (2560, 1440)
-        print('Invalid resolution input. Automatically changed to 1024x1024 or 2K.')
-
-    if isinstance(images, list):
-        # For tab_batch
-        save_paths = []
-        save_dir = 'preds-BiRefNet'
-        if not os.path.exists(save_dir):
-            os.makedirs(save_dir)
-        tab_is_batch = True
-    else:
-        images = [images]
-        tab_is_batch = False
-
-    for idx_image, image_src in enumerate(images):
-        if isinstance(image_src, str):
-            if os.path.isfile(image_src):
-                image_ori = Image.open(image_src)
-            else:
-                response = requests.get(image_src)
-                image_data = BytesIO(response.content)
-                image_ori = Image.open(image_data)
-        else:
-            image_ori = Image.fromarray(image_src)
-
-        image = image_ori.convert('RGB')
-        # Preprocess the image
-        image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
-        image_proc = image_preprocessor.proc(image)
-        image_proc = image_proc.unsqueeze(0)
-
-        # Prediction
-        with torch.no_grad():
-            preds = birefnet(image_proc.to(device))[-1].sigmoid().cpu()
-        pred = preds[0].squeeze()
-
-        # Show Results
-        pred_pil = transforms.ToPILImage()(pred)
-        image_masked = refine_foreground(image, pred_pil)
-        image_masked.putalpha(pred_pil.resize(image.size))
-
-        torch.cuda.empty_cache()
-
-        if tab_is_batch:
-            save_file_path = os.path.join(save_dir, "{}.png".format(os.path.splitext(os.path.basename(image_src))[0]))
-            image_masked.save(save_file_path)
-            save_paths.append(save_file_path)
-
-    if tab_is_batch:
-        zip_file_path = os.path.join(save_dir, "{}.zip".format(save_dir))
-        with zipfile.ZipFile(zip_file_path, 'w') as zipf:
-            for file in save_paths:
-                zipf.write(file, os.path.basename(file))
-        return save_paths, zip_file_path
-    else:
-        return (image_masked, image_ori)
-
-
-examples = [[_] for _ in glob('examples/*')][:]
-# Add the option of resolution in a text box.
-for idx_example, example in enumerate(examples):
-    examples[idx_example].append('1024x1024')
-examples.append(examples[-1].copy())
-examples[-1][1] = '512x512'
-
-examples_url = [
-    ['https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg'],
-]
-for idx_example_url, example_url in enumerate(examples_url):
-    examples_url[idx_example_url].append('1024x1024')
-
-descriptions = ('Upload a picture, our model will extract a highly accurate segmentation of the subject in it.\n)'
-                 ' The resolution used in our training was `1024x1024`, thus the suggested resolution to obtain good results!\n'
-                 ' Our codes can be found at https://github.com/ZhengPeng7/BiRefNet.\n'
-                 ' We also maintain the HF model of BiRefNet at https://huggingface.co/ZhengPeng7/BiRefNet for easier access.')
-
-tab_image = gr.Interface(
-    fn=predict,
-    inputs=[
-        gr.Image(label='Upload an image'),
-        gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`.", label="Resolution"),
-        gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want.")
-    ],
-    outputs=ImageSlider(label="BiRefNet's prediction", type="pil"),
-    examples=examples,
-    api_name="image",
-    description=descriptions,
-)
+    image_ori = Image.fromarray(image)
+    image = image_ori.convert('RGB')
+    
+    # Preprocess the image
+    image_preprocessor = ImagePreprocessor(resolution=(1024, 1024))
+    image_proc = image_preprocessor.proc(image)
+    image_proc = image_proc.unsqueeze(0)
 
-tab_text = gr.Interface(
-    fn=predict,
-    inputs=[
-        gr.Textbox(label="Paste an image URL"),
-        gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`.", label="Resolution"),
-        gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want.")
-    ],
-    outputs=ImageSlider(label="BiRefNet's prediction", type="pil"),
-    examples=examples_url,
-    api_name="text",
-    description=descriptions+'\nTab-URL is partially modified from https://huggingface.co/spaces/not-lain/background-removal, thanks to this great work!',
-)
+    # Prediction
+    with torch.no_grad():
+        preds = birefnet(image_proc.to(device))[-1].sigmoid().cpu()
+    pred = preds[0].squeeze()
 
-tab_batch = gr.Interface(
-    fn=predict,
-    inputs=[
-        gr.File(label="Upload multiple images", type="filepath", file_count="multiple"),
-        gr.Textbox(lines=1, placeholder="Type the resolution (`WxH`) you want, e.g., `1024x1024`.", label="Resolution"),
-        gr.Radio(list(usage_to_weights_file.keys()), value='General', label="Weights", info="Choose the weights you want.")
-    ],
-    outputs=[gr.Gallery(label="BiRefNet's predictions"), gr.File(label="Download masked images.")],
-    api_name="batch",
-    description=descriptions+'\nTab-batch is partially modified from https://huggingface.co/spaces/NegiTurkey/Multi_Birefnetfor_Background_Removal, thanks to this great work!',
-)
+    # Show Results
+    pred_pil = transforms.ToPILImage()(pred)
+    image_masked = refine_foreground(image, pred_pil)
+    image_masked.putalpha(pred_pil.resize(image.size))
 
-demo = gr.TabbedInterface(
-    [tab_image, tab_text, tab_batch],
-    ['image', 'text', 'batch'],
-    title="BiRefNet demo for subject extraction (general / matting / salient / camouflaged / portrait).",
+    torch.cuda.empty_cache()
+
+    # Save as PNG
+    output_path = "output.png"
+    image_masked.save(output_path)
+
+    return output_path
+
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="numpy"),
+    outputs=gr.Image(type="filepath"),
+    title="BiRefNet Matting",
+    description="Upload an image to perform matting using BiRefNet."
 )
 
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    iface.launch(debug=True)