diff --git a/groupword.py b/groupword.py index 556c6a3..69b8fa9 100644 --- a/groupword.py +++ b/groupword.py @@ -1,83 +1,56 @@ -from asyncio import sleep from wordcloud import WordCloud from io import BytesIO from os.path import exists -from os import makedirs -from sys import executable +from os import makedirs, sep, remove +from PIL import Image from collections import defaultdict from requests import get from pagermaid import version -from pagermaid.utils import execute, alias_command +from pagermaid.utils import execute, alias_command, pip_install from pagermaid.listener import listener -imported = True -imported_ = True +pip_install("jieba") + +import jieba + punctuation = {33: ' ', 34: ' ', 35: ' ', 36: ' ', 37: ' ', 38: ' ', 39: ' ', 40: ' ', 41: ' ', 42: ' ', 43: ' ', 44: ' ', 45: ' ', 46: ' ', 47: ' ', 58: ' ', 59: ' ', 60: ' ', 61: ' ', 62: ' ', 63: ' ', 64: ' ', 91: ' ', 92: ' ', 93: ' ', 94: ' ', 95: ' ', 96: ' ', 123: ' ', 124: ' ', 125: ' ', 126: ' ', 65311: ' ', 65292: ' ', 65281: ' ', 12304: ' ', 12305: ' ', 65288: ' ', 65289: ' ', 12289: ' ', 12290: ' ', 65306: ' ', 65307: ' ', 8217: ' ', 8216: ' ', 8230: ' ', 65509: ' ', 183: ' '} -try: - import jieba -except ImportError: - imported = False -try: - import paddle -except ImportError: - imported_ = False @listener(is_plugin=True, outgoing=True, command=alias_command("groupword"), - description="拉取最新 300 条消息生成词云。", + description="拉取最新 500 条消息生成词云，回复图片可自定义背景图。（图片白色区域为不显示部分）", parameters="[任意内容启用AI分词]") async def group_word(context): + if not exists(f"plugins{sep}groupword"): + makedirs(f"plugins{sep}groupword") + # 自定义背景图 + reply = await context.get_reply_message() + if reply and reply.photo: + if exists(f"plugins{sep}groupword{sep}circle.jpg"): + remove(f"plugins{sep}groupword{sep}circle.jpg") + await context.client.download_media(reply, file=f"plugins{sep}groupword{sep}circle.jpg") + return await context.edit("自定义背景图设置完成。") + imported_1 = False if len(context.parameter) >= 1: + pip_install("paddlepaddle-tiny", alias="paddle") imported_1 = True - if not imported: - try: - await context.edit("支持库 `jieba` 未安装...\n正在尝试自动安装...") - await execute(f'{executable} -m pip install jieba') - await sleep(10) - result = await execute(f'{executable} -m pip show jieba') - if len(result) > 0: - await context.edit('支持库 `jieba` 安装成功...\n正在尝试自动重启...') - await context.client.disconnect() - else: - await context.edit(f"自动安装失败..请尝试手动安装 `{executable} -m pip install jieba` 随后，请重启 PagerMaid-Modify 。") - return - except: - return - if not imported_ and imported_1: - try: - await context.edit("支持库 `paddlepaddle-tiny` 未安装...\n正在尝试自动安装...") - await execute(f'{executable} -m pip install paddlepaddle-tiny') - await sleep(10) - result = await execute(f'{executable} -m pip show paddlepaddle-tiny') - if len(result) > 0 and not 'WARNING' in result: - await context.edit('支持库 `paddlepaddle-tiny` 安装成功...\n正在尝试自动重启...') - await context.client.disconnect() - else: - await context.edit(f"自动安装失败，可能是系统不支持..\nAI 分词不可用，切换到基础分词。\n" - f"您可以尝试手动安装 `{executable} -m pip install paddlepaddle-tiny` 。") - await sleep(4) - except: - return try: await context.edit('正在生成中。。。') except: return - if not exists("plugins/groupword"): - makedirs("plugins/groupword") - if not exists("plugins/groupword/wqy-microhei.ttc"): + if not exists(f"plugins{sep}groupword{sep}wqy-microhei.ttc"): await context.edit('正在拉取中文字体文件。。。（等待时间请评估你的服务器）') r = get('https://cdn.jsdelivr.net/gh/anthonyfok/fonts-wqy-microhei/wqy-microhei.ttc') - with open("plugins/groupword/wqy-microhei.ttc", "wb") as code: + with open(f"plugins{sep}groupword{sep}wqy-microhei.ttc", "wb") as code: code.write(r.content) words = defaultdict(int) count = 0 try: - if imported_ and imported_1: + if imported_1: try: jieba.enable_paddle() except: @@ -87,7 +60,7 @@ async def group_word(context): continue if msg.text and not msg.text.startswith('/') and not msg.text.startswith('-') and not '//' in msg.text: try: - if imported_ and imported_1: + if imported_1: for word in jieba.cut(msg.text.translate(punctuation), use_paddle=True): word = word.lower() words[word] += 1 @@ -106,9 +79,21 @@ async def group_word(context): except: return try: - image = WordCloud(font_path="plugins/groupword/wqy-microhei.ttc", width=800, - height=400).generate_from_frequencies( - words).to_image() + if not exists(f"plugins{sep}groupword{sep}circle.jpg"): + image = WordCloud(font_path=f"plugins{sep}groupword{sep}wqy-microhei.ttc", width=800, + height=400).generate_from_frequencies( + words).to_image() + else: + import numpy as np + background = Image.open(f"plugins{sep}groupword{sep}circle.jpg") + graph = np.array(background) # noqa + + image = WordCloud(font_path=f"plugins{sep}groupword{sep}wqy-microhei.ttc", width=800, + height=800, + mask=graph, + scale=5, + background_color='white').generate_from_frequencies( + words).to_image() stream = BytesIO() image.save(stream, 'PNG') except: diff --git a/list.json b/list.json index 8211e54..460a364 100644 --- a/list.json +++ b/list.json @@ -492,13 +492,13 @@ }, { "name": "groupword", - "version": "1.061", + "version": "1.1", "section": "chat", - "maintainer": "xtaodada", - "size": "3.32 kb", + "maintainer": "xtaodada, matsuri_nya", + "size": "4.5 kb", "supported": true, "des-short": "将聊天内容生成为词云", - "des": "拉取最新 300 条消息生成一张词云图片。\n\n命令：-groupword" + "des": "拉取最新 500 条消息生成词云，回复图片可自定义背景图。（图片白色区域为不显示部分）\n\n命令：-groupword" }, { "name": "kongge",