from wordcloud import WordCloud, STOPWORDS
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import jieba

# 打开文件
text = open("zhihu.txt", encoding="utf-8").read()

# 中文分词
text = ' '.join(jieba.cut(text))
print(text[:100])

# 停用词
stopwords = STOPWORDS
STOPWORDS.add("的")
STOPWORDS.add("是")
STOPWORDS.add("和")
STOPWORDS.add("对")
STOPWORDS.add("上")
STOPWORDS.add("人")
STOPWORDS.add("有")
STOPWORDS.add("就")
STOPWORDS.add("一")
STOPWORDS.add("不")
STOPWORDS.add("要")
STOPWORDS.add("等等")
STOPWORDS.add("在")
STOPWORDS.add("比如")
STOPWORDS.add("自己")
STOPWORDS.add("我们")
STOPWORDS.add("不能")

# 生成对象
# 使用蒙版图片
mask = np.array(Image.open("bird.jpg"))  # 使用蒙版图片
wc = WordCloud(
               mask=mask,
               font_path="STKAITI.TTF",
               mode="RGBA",
               background_color=None,
               stopwords=STOPWORDS
).generate(text)

# 显示词云
plt.imshow(wc, interpolation="bilinear")
plt.axis("off")
plt.show()

# 保存文件
# wc.to_file("zhihu_wordcloud.png")