Spaces:
Running
Running
File size: 2,720 Bytes
11aa79a |
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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
import base64
from rembg import remove, new_session
from typing import Annotated
from fastapi import FastAPI, UploadFile, Request, File, Form
from fastapi.templating import Jinja2Templates
from fastapi.responses import RedirectResponse, HTMLResponse, Response
import uvicorn
import os
os.environ['U2NET_HOME'] = './models/'
model_names = [
'u2net',
'u2net_human_seg',
'u2net_cloth_seg',
'isnet-general-use',
]
sessions = {
'u2net': new_session(model_name=model_names[0]),
}
app = FastAPI()
templates = Jinja2Templates(directory="templates")
@app.get("/")
def health(request: Request):
return templates.TemplateResponse('dynamic.html', { "request": request })
@app.post("/remove")
async def remove_bg(
request: Request,
file: Annotated[UploadFile, File()],
mask_only: Annotated[str, Form()] = 'off',
name_of_model: Annotated[str, Form()] = 'u2net'
) -> HTMLResponse:
try:
if name_of_model not in sessions.keys():
sessions[name_of_model] = new_session(model_name=name_of_model)
current_session = sessions[name_of_model]
only_mask = mask_only == 'on'
data = file.file.read()
output_array = remove(data, only_mask=only_mask, session=current_session)
output_img = base64.b64encode(output_array).decode('utf-8')
file.file.close()
encoded_image = base64.b64encode(data).decode('utf-8')
return templates.TemplateResponse(
'dynamic.html',
{
"request" : request,
"image" : encoded_image,
"output_img" : output_img
}
)
except Exception as error_msg:
return templates.TemplateResponse(
"error.html",
{
"request" : request,
"error_msg" : str(error_msg),
}
)
@app.post("/remove_bg")
async def remove_bg(
request: Request,
file: UploadFile,
mask_only: str = 'off',
name_of_model: str = 'u2net'
):
try:
if name_of_model not in sessions.keys():
sessions[name_of_model] = new_session(model_name=name_of_model)
current_session = sessions[name_of_model]
only_mask = mask_only == 'on'
data = file.file.read()
output_array = remove(data, only_mask=only_mask, session=current_session)
file.file.close()
return Response(content=output_array, media_type="image/png")
except Exception as error:
return f"Oopss!!!! {error}"
@app.get("/remove")
def remove_bg_redirect():
return RedirectResponse('/')
if __name__ == '__main__':
uvicorn.run(app, host='0.0.0.0', port=7860)
|