1#!/usr/local/bin/python3.8 2""" 3Using custom colors 4=================== 5 6Using the recolor method and custom coloring functions. 7""" 8 9import numpy as np 10from PIL import Image 11from os import path 12import matplotlib.pyplot as plt 13import os 14import random 15 16from wordcloud import WordCloud, STOPWORDS 17 18 19def grey_color_func(word, font_size, position, orientation, random_state=None, 20 **kwargs): 21 return "hsl(0, 0%%, %d%%)" % random.randint(60, 100) 22 23 24# get data directory (using getcwd() is needed to support running example in generated IPython notebook) 25d = path.dirname(__file__) if "__file__" in locals() else os.getcwd() 26 27# read the mask image taken from 28# http://www.stencilry.org/stencils/movies/star%20wars/storm-trooper.gif 29mask = np.array(Image.open(path.join(d, "stormtrooper_mask.png"))) 30 31# movie script of "a new hope" 32# http://www.imsdb.com/scripts/Star-Wars-A-New-Hope.html 33# May the lawyers deem this fair use. 34text = open(path.join(d, 'a_new_hope.txt')).read() 35 36# pre-processing the text a little bit 37text = text.replace("HAN", "Han") 38text = text.replace("LUKE'S", "Luke") 39 40# adding movie script specific stopwords 41stopwords = set(STOPWORDS) 42stopwords.add("int") 43stopwords.add("ext") 44 45wc = WordCloud(max_words=1000, mask=mask, stopwords=stopwords, margin=10, 46 random_state=1).generate(text) 47# store default colored image 48default_colors = wc.to_array() 49plt.title("Custom colors") 50plt.imshow(wc.recolor(color_func=grey_color_func, random_state=3), 51 interpolation="bilinear") 52wc.to_file("a_new_hope.png") 53plt.axis("off") 54plt.figure() 55plt.title("Default colors") 56plt.imshow(default_colors, interpolation="bilinear") 57plt.axis("off") 58plt.show() 59