File size: 720 Bytes
97c5c48
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import gradio as gr
from transformers import pipeline
from PIL import Image, ImageOps

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

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

def remove_background(img):
    segformer_b2_clothes_result = segformer_b2_clothes(img)

	return segformer_b2_clothes_result

iface = gr.Interface(fn=remove_background, 
                     inputs=gr.Image(type='pil'), 
                     outputs=gr.Image(label='segformer_b2_clothes', type='pil'))
iface.launch()