File size: 1,652 Bytes
b5ba7a5 |
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 47 48 49 50 51 52 53 |
from fastapi import FastAPI, Body
from modules.api.models import *
from modules.api import api
import gradio as gr
import rembg
# models = [
# "None",
# "u2net",
# "u2netp",
# "u2net_human_seg",
# "u2net_cloth_seg",
# "silueta",
# ]
def rembg_api(_: gr.Blocks, app: FastAPI):
@app.post("/rembg")
async def rembg_remove(
input_image: str = Body("", title='rembg input image'),
model: str = Body("u2net", title='rembg model'),
return_mask: bool = Body(False, title='return mask'),
alpha_matting: bool = Body(False, title='alpha matting'),
alpha_matting_foreground_threshold: int = Body(240, title='alpha matting foreground threshold'),
alpha_matting_background_threshold: int = Body(10, title='alpha matting background threshold'),
alpha_matting_erode_size: int = Body(10, title='alpha matting erode size')
):
if not model or model == "None":
return
input_image = api.decode_base64_to_image(input_image)
image = rembg.remove(
input_image,
session=rembg.new_session(model),
only_mask=return_mask,
alpha_matting=alpha_matting,
alpha_matting_foreground_threshold=alpha_matting_foreground_threshold,
alpha_matting_background_threshold=alpha_matting_background_threshold,
alpha_matting_erode_size=alpha_matting_erode_size,
)
return {"image": api.encode_pil_to_base64(image).decode("utf-8")}
try:
import modules.script_callbacks as script_callbacks
script_callbacks.on_app_started(rembg_api)
except:
pass
|