68 lines
1.8 KiB
Python
68 lines
1.8 KiB
Python
## Modified from Akhaliq Hugging Face Demo
|
|
## https://huggingface.co/akhaliq
|
|
|
|
import gradio as gr
|
|
import os
|
|
import cv2
|
|
|
|
def inference(file, af, mask):
|
|
im = cv2.imread(file, cv2.IMREAD_COLOR)
|
|
cv2.imwrite(os.path.join("input.png"), im)
|
|
|
|
model = "u2net"
|
|
from rembg import remove
|
|
from rembg.session_base import BaseSession
|
|
from rembg.session_factory import new_session
|
|
|
|
input_path = 'input.png'
|
|
output_path = 'output.png'
|
|
|
|
with open(input_path, 'rb') as i:
|
|
with open(output_path, 'wb') as o:
|
|
input = i.read()
|
|
sessions: dict[str, BaseSession] = {}
|
|
output = remove(
|
|
input,
|
|
session=sessions.setdefault(
|
|
model, new_session(model)
|
|
),
|
|
alpha_matting_erode_size = af,
|
|
only_mask = (True if mask == "Mask only" else False)
|
|
)
|
|
o.write(output)
|
|
return os.path.join("output.png")
|
|
|
|
|
|
|
|
gr.Interface(
|
|
inference,
|
|
[
|
|
gr.inputs.Image(type="filepath", label="Input"),
|
|
gr.inputs.Slider(10, 25, default=10, label="Alpha matting"),
|
|
gr.inputs.Radio(
|
|
[
|
|
"Default",
|
|
"Mask only"
|
|
],
|
|
type="value",
|
|
default="Default",
|
|
label="Choices"
|
|
),
|
|
#gr.inputs.Dropdown([
|
|
# "u2net",
|
|
# "u2netp",
|
|
# "u2net_human_seg",
|
|
# "u2net_cloth_seg",
|
|
# "silueta"
|
|
# ],
|
|
# type="value",
|
|
# default="u2net",
|
|
# label="Models"
|
|
#),
|
|
],
|
|
# gr.outputs.Image(type="filepath", label="Output"),
|
|
gr.Image(),
|
|
examples=[["lion.png", 10, "Default", "u2net"], ["girl.jpg", 10, "Default", "u2net"]],
|
|
enable_queue=True
|
|
).launch(server_name="0.0.0.0" , share=True)
|