|
import gradio as gr |
|
from transformers import pipeline |
|
from PIL import Image, ImageOps |
|
import time |
|
from rembg import remove |
|
|
|
|
|
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): |
|
|
|
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""" |
|
|
|
|
|
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() |
