Spaces:
Sleeping
Sleeping
File size: 2,491 Bytes
d39fc00 b2502f7 d39fc00 b4a01b7 |
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 |
from typing import Union, List
import requests
import uvicorn
from fastapi import BackgroundTasks, FastAPI
import img_label
from img_nsfw import init_nsfw_pipe, check_nsfw
import model
app = FastAPI()
def write_scan_img_result(image_id: int, scans: List[int], img: str, callback: str):
score_general_threshold = 0.35
score_character_threshold = 0.85
nsfw_tags = []
img_tags = []
if 0 in scans:
nsfw_tags = check_nsfw(img, pipe)
if 1 in scans:
img_tags = img_label.label_img(
image=img,
model="SwinV2",
l_score_general_threshold=score_general_threshold,
l_score_character_threshold=score_character_threshold,
)['general_res']
print(nsfw_tags)
print(img_tags)
img_tags = list(map(lambda x: model.ImageTag(tag=x['tag'], confidence=x['confidence']), img_tags))
callBackReq = model.ImageScanCallbackRequest(id=image_id, isValid=True, tags=img_tags)
try:
requests.post(callback, json=callBackReq.dict())
except Exception as ex:
print(ex)
nsfw_tags = list(map(lambda x: model.ImageScanTag(type="Moderation",
confidence=x['confidence']), nsfw_tags))
ret = model.ImageScanResponse(ok=True, error="", deleted=False, blockedFor=[], tags=nsfw_tags)
return ret
def write_scan_model_result(model_name: str, callback: str):
pass
# @app.post("/model-scan")
# async def send_notification(email: str, background_tasks: BackgroundTasks):
# background_tasks.add_task(write_scan_model_result, email, callback="")
# return {"message": "Notification sent in the background"}
@app.post("/image-scan")
async def image_scan_handler(req: model.ImageScanRequest, background_tasks: BackgroundTasks):
if not req.wait:
background_tasks.add_task(write_scan_img_result,
image_id=req.imageId,
scans=req.scans,
img=req.url, callback=req.callbackUrl)
return model.ImageScanResponse(ok=True, error="", deleted=False, blockedFor=[], tags=[])
else:
ret = write_scan_img_result(image_id=req.imageId,
scans=req.scans,
img=req.url, callback=req.callbackUrl)
return ret
if __name__ == "__main__":
global pipe
pipe = init_nsfw_pipe()
uvicorn.run(app, host="0.0.0.0", port=7860)
|