42 lines
1.7 KiB
Python
42 lines
1.7 KiB
Python
|
import aiohttp
|
||
|
|
import base64
|
||
|
|
from nonebot import on_command
|
||
|
|
from nonebot.adapters.onebot.v11 import GroupMessageEvent, MessageSegment
|
||
|
|
from nonebot.log import logger
|
||
|
|
from .translation import translate
|
||
|
|
|
||
|
|
deepdanbooru = on_command(".gettag", aliases={"鉴赏", "查书"})
|
||
|
|
|
||
|
|
|
||
|
|
@deepdanbooru.handle()
|
||
|
|
async def deepdanbooru_handle(event: GroupMessageEvent):
|
||
|
|
url = ""
|
||
|
|
for seg in event.message['image']:
|
||
|
|
url = seg.data["url"]
|
||
|
|
if url:
|
||
|
|
async with aiohttp.ClientSession() as session:
|
||
|
|
logger.info(f"正在获取图片")
|
||
|
|
async with session.get(url) as resp:
|
||
|
|
bytes = await resp.read()
|
||
|
|
str_img = str(base64.b64encode(bytes), "utf-8")
|
||
|
|
message = MessageSegment.at(event.user_id)
|
||
|
|
start = "data:image/jpeg;base64,"
|
||
|
|
str0 = start+str_img
|
||
|
|
async with aiohttp.ClientSession() as session:
|
||
|
|
async with session.post('https://mayhug-rainchan-anime-image-label.hf.space/api/predict/', json={"data": [str0, 0.6,"ResNet101"]}) as resp:
|
||
|
|
if resp.status != 200:
|
||
|
|
await deepdanbooru.finish(f"识别失败，错误代码为{resp.status}")
|
||
|
|
jsonresult = await resp.json()
|
||
|
|
data = jsonresult['data'][0]
|
||
|
|
logger.info(f"TAG查询完毕")
|
||
|
|
tags = ""
|
||
|
|
for label in data['confidences']:
|
||
|
|
tags = tags+label["label"]+","
|
||
|
|
tags_ch = await translate(tags.replace("_", " "), "zh")
|
||
|
|
if tags_ch == tags.replace("_", " "):
|
||
|
|
message = message+tags
|
||
|
|
message = message+tags+f"\n机翻结果:"+tags_ch
|
||
|
|
await deepdanbooru.finish(message)
|
||
|
|
else:
|
||
|
|
await deepdanbooru.finish(f"未找到图片")
|