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"]
|