diff --git a/GroupTagCloud.py b/GroupTagCloud.py index 9126958..5bf2335 100644 --- a/GroupTagCloud.py +++ b/GroupTagCloud.py @@ -12,17 +12,23 @@ import time import logging import jieba +import numpy as np +from PIL import Image class GroupTagCloud(ProcessInterface): recordMaxNum = 500 maxFrequency = 40 imgDir = 'TagCloud' - def __init__(self, fontPath): + def __init__(self, fontPath, maskPath = None): self.client = MongoClient() self.coll = self.client[dbName][collName] self.fontPath = fontPath - self.wordCloud = WordCloud(font_path=self.fontPath, width=400, height=400, max_words=100) + if maskPath is None: + self.wordCloud = WordCloud(font_path=self.fontPath, width=400, height=400, max_words=100) + else: + self.mask = np.array(Image.open(maskPath)) + self.wordCloud = WordCloud(font_path=self.fontPath, mask=self.mask, max_words=100) if not os.path.exists(self.imgDir): os.mkdir(self.imgDir) logging.info('GroupTagCloud connected to MongoDB.') diff --git a/girl.png b/girl.png new file mode 100644 index 0000000..18312a9 Binary files /dev/null and b/girl.png differ diff --git a/main.py b/main.py index 844a9b4..df8dd4a 100644 --- a/main.py +++ b/main.py @@ -22,7 +22,7 @@ GlobalTextHook({ '^ding$': 'dong', '鸭哥': '嘎?' }), PaiDuiHook(), HistoryRecorder(), - GroupTagCloud('/usr/share/fonts/truetype/wqy/wqy-microhei.ttc'), + GroupTagCloud('/usr/share/fonts/truetype/wqy/wqy-microhei.ttc','girl.png'), ActivityInfo('/usr/share/fonts/truetype/wqy/wqy-microhei.ttc'), GroupMessageForwarder([ '二群', '三群' ], [ 'AI二群测试中', 'AI三群测试' ]) ]