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