termcloud/termcloud.py at master · brook-stevens/termcloud · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
import nltk

from nltk.stem import WordNetLemmatizer
from nltk.corpus import stopwords

from wordcloud import WordCloud, STOPWORDS
import io

from os import path

def generate(text, font_path, no_lemmatization, width, height):
    documents = []
    buf = io.StringIO(text)
    for line in buf:
        line = line.rstrip('\n')
        if(line != ''):
            documents.append(line)

    engstopwords = stopwords.words('english')
    engstopwords.append("n't")
    engstopwords.append("ve")

    all_words = []

    for curr_doc_index, val in enumerate(documents):
        tokens = nltk.word_tokenize(documents[curr_doc_index])

        lemmatizer = WordNetLemmatizer()
        tokens_cleaned = []
        for tok in tokens:
            if tok.lower() in no_lemmatization:
                tokens_cleaned.append(lemmatizer.lemmatize(tok))
            else:
                tokens_cleaned.append(tok)

        tokens = tokens_cleaned

        tokens_cleaned = []
        for tok in tokens:
            if tok.lower() not in engstopwords:
                tokens_cleaned.append(tok)
        tokens = tokens_cleaned

        all_words.extend(tokens)


    # Convert all the required text into a single string here
    # and store them in word_string

    # you can specify fonts, stopwords, background color and other options
    all_words_string = ' '.join(str(e) for e in all_words)

    d = path.dirname(__file__)

    def grey_color_func(word, font_size, position, orientation, random_state=None,
                        **kwargs):
        #print "hsl(0, 0%%, %d%%)" % random.randint(60, 100)
        return "hsl(0, 0%, 40%)"

    wordcloud = WordCloud(font_path=font_path,
                            stopwords=STOPWORDS,
                            background_color='white',
                            width=width,
                            height=height,
                            #mask=mask,
                            max_font_size=48,
                          normalize_plurals=False,
                          color_func=grey_color_func
                          ).generate(all_words_string)

    image = wordcloud.to_image()

    return image