File size: 1,940 Bytes
97c5c48
 
 
dec6da5
83b67a1
97c5c48
 
 
 
 
8993d92
97c5c48
 
 
 
 
83b67a1
 
 
97c5c48
83b67a1
dec6da5
97c5c48
dec6da5
405b28a
83b67a1
 
 
 
 
 
71c7d66
83b67a1
97c5c48
f1b81c6
 
 
 
 
 
97c5c48
dec6da5
83b67a1
 
 
97c5c48
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
import gradio as gr
from transformers import pipeline
from PIL import Image, ImageOps
import time
from rembg import remove

# Initialize Segmentation Pipeline
segformer_b2_clothes_pipe = pipeline("image-segmentation", model="mattmdjaga/segformer_b2_clothes")

def segformer_b2_clothes(img):
    result = segformer_b2_clothes_pipe(img)
    mask = result[0]['mask'].convert('L')
    mask = ImageOps.invert(mask)
    img.putalpha(mask)
    return img

def rembg_remove(img):
    return remove(img)

def remove_background(img):
    # segformer_b2_clothes
    start = time.time()
    segformer_b2_clothes_result = segformer_b2_clothes(img)
    end = time.time()
    segformer_b2_clothes_text = """[mattmdjaga/segformer_b2_clothes](https://huggingface.co/mattmdjaga/segformer_b2_clothes) \n""" + str(end-start) + """ seconds"""

    #rembg
    start = time.time()
    rembg_result = rembg_remove(img)
    end = time.time()
    rembg_text = "[rembg](https://huggingface.co/spaces/openskyml/remove-background-on-image) \n" + str(end-start) + " seconds"
    
    return segformer_b2_clothes_text, segformer_b2_clothes_result, rembg_text, rembg_result 

iface = gr.Interface(fn=remove_background,
                     title='Remove Background Comparison',
                     description="""
                     Compares [mattmdjaga/segformer_b2_clothes](https://huggingface.co/mattmdjaga/segformer_b2_clothes) and [rembg](https://huggingface.co/spaces/openskyml/remove-background-on-image) background removal.
                     """,
                     theme = gr.themes.Base(primary_hue="teal",secondary_hue="teal",neutral_hue="slate"),
                     inputs=gr.Image(type='pil'), 
                     outputs=[gr.Markdown(),
                              gr.Image(label='segformer_b2_clothes', type='pil'),
                              gr.Markdown(),
                              gr.Image(label='rembg', type='pil')])
iface.launch()