wordcloud_cli.py 3.51 KB
Newer Older
Raphael Boidol's avatar
Raphael Boidol committed
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
74
75
76
77
78
79
80
81
82
83
#!/usr/bin/env python
# -*- coding: utf-8 -*-
r"""Command-line tool to generate word clouds
Usage::
    $ cat word.txt | wordcloud_cli.py

    $ wordcloud_cli.py --text=words.txt --stopwords=stopwords.txt
"""
import argparse
import wordcloud as wc
import numpy as np
import sys
import io
from PIL import Image

def main(args):
    wordcloud = wc.WordCloud(stopwords=args.stopwords, mask=args.mask,
        width=args.width, height=args.height, font_path=args.font_path,
        margin=args.margin, relative_scaling=args.relative_scaling,
        color_func=args.color_func, background_color=args.background_color).generate(args.text)
    image = wordcloud.to_image()

    b = io.BytesIO()
    image.save(b, format='png')
    with args.imagefile:
        args.imagefile.write(b.getvalue())

def parse_args(arguments):
    prog = 'python wordcloud_cli.py'
    description = ('A simple command line interface for wordcloud module.')
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument('--text', metavar='file', type=argparse.FileType(), default='-',
        help='specify file of words to build the word cloud (default: stdin)')
    parser.add_argument('--stopwords', metavar='file', type=argparse.FileType(),
        help='specify file of stopwords (containing one word per line) to remove from the given text after parsing')
    parser.add_argument('--imagefile', metavar='file', type=argparse.FileType('w'), default='-',
        help='file the completed PNG image should be written to (default: stdout)')
    parser.add_argument('--fontfile', metavar='path', dest='font_path',
        help='path to font file you wish to use (default: DroidSansMono)')
    parser.add_argument('--mask', metavar='file', type=argparse.FileType(),
        help='mask to use for the image form')
    parser.add_argument('--colormask', metavar='file', type=argparse.FileType(),
        help='color mask to use for image coloring')
    parser.add_argument('--relative_scaling', type=float, default=0,
        metavar='rs', help=' scaling of words by frequency (0 - 1)')
    parser.add_argument('--margin', type=int, default=2,
        metavar='width', help='spacing to leave around words')
    parser.add_argument('--width', type=int, default=400,
        metavar='width', help='define output image width')
    parser.add_argument('--height', type=int, default=200,
        metavar='height', help='define output image height')
    parser.add_argument('--color', metavar='color',
        help='use given color as coloring for the image - accepts any value from PIL.ImageColor.getcolor')
    parser.add_argument('--background', metavar='color', default='black', type=str, dest='background_color',
        help='use given color as background color for the image - accepts any value from PIL.ImageColor.getcolor')
    args = parser.parse_args(arguments)

    if args.colormask and args.color:
        raise ValueError('specify either a color mask or a color function')

    with args.text:
        args.text = args.text.read()

    if args.stopwords:
        with args.stopwords:
            args.stopwords = set(map(str.strip, args.stopwords.readlines()))

    if args.mask:
        args.mask = np.array(Image.open(args.mask))

    color_func = wc.random_color_func
    if args.colormask:
        image = np.array(Image.open(args.colormask))
        color_func = wc.ImageColorGenerator(image)

    if args.color:
        color_func = wc.get_single_color_func(args.color)

    args.color_func = color_func
    return args

if __name__ == '__main__':
    main(parse_args(sys.argv[1:]))