55 lines
1.8 KiB
Python
55 lines
1.8 KiB
Python
from ..config import config, nickname
|
||
from ..utils.save import save_img
|
||
from io import BytesIO
|
||
import base64
|
||
import aiohttp
|
||
import asyncio
|
||
from nonebot.adapters.onebot.v11 import MessageSegment
|
||
from nonebot.log import logger
|
||
|
||
|
||
async def check_safe_method(fifo, img_bytes, message):
|
||
if config.novelai_h:
|
||
for i in img_bytes:
|
||
await save_img(fifo, i)
|
||
message += MessageSegment.image(i)
|
||
else:
|
||
nsfw_count = 0
|
||
for i in img_bytes:
|
||
try:
|
||
label = await check_safe(i)
|
||
except RuntimeError as e:
|
||
logger.error(f"NSFWAPI调用失败，错误代码为{e.args}")
|
||
label = "unknown"
|
||
if label != "explicit":
|
||
message += MessageSegment.image(i)
|
||
else:
|
||
nsfw_count += 1
|
||
await save_img(fifo, i, label)
|
||
if nsfw_count > 0:
|
||
message += f"\n有{nsfw_count}张图片太涩了，{nickname}已经帮你吃掉了哦"
|
||
return message
|
||
|
||
|
||
async def check_safe(img_bytes: BytesIO):
|
||
# 检查图片是否安全
|
||
start = "data:image/jpeg;base64,"
|
||
image = img_bytes.getvalue()
|
||
image = str(base64.b64encode(image), "utf-8")
|
||
str0 = start + image
|
||
# 重试三次
|
||
for i in range(3):
|
||
async with aiohttp.ClientSession() as session:
|
||
# 调用API
|
||
async with session.post('https://hf.space/embed/mayhug/rainchan-image-porn-detection/api/predict/',
|
||
json={"data": [str0]}) as resp:
|
||
if resp.status == 200:
|
||
jsonresult = await resp.json()
|
||
break
|
||
else:
|
||
await asyncio.sleep(2)
|
||
error = resp.status
|
||
else:
|
||
raise RuntimeError(error)
|
||
return jsonresult["data"][0]["label"]
|