Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 81 additions & 43 deletions stylecloud/stylecloud.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import os
from PIL import Image
from matplotlib.colors import makeMappingArray, to_rgb
import matplotlib.pyplot as plt
import numpy as np
import fire
from shutil import rmtree
Expand All @@ -19,30 +20,33 @@ def file_to_text(file_path: str):
read as a dict of word/weights.
"""

if not file_path.endswith(".csv"):
with open(file_path, "r", encoding="utf-8") as f:
text = f.read()
return text
else: # parse as a CSV

if file_path.endswith(".csv"):
with open(file_path, "r", encoding="utf-8") as f:
r = csv.reader(f)
header = next(r)

assert len(header) <= 2, "The input CSV has too many columns."

# If a single-column CSV, read as a bulk text
if len(header) == 1:
texts = ""
for row in r:
texts += row[0] + "\n"
# If a two-column CSV, read as words/weights

# If a two-column CSV, read as words/frequencies.
elif len(header) == 2:
texts = {}
for row in r:
texts[row[0]] = float(row[1])
texts[row[0]] = float(row[1])
return texts



else:
# Read file as a bulk text
with open(file_path, "r", encoding="utf-8") as f:
text = f.read()
return text


def gen_fa_mask(
icon_name: str = "fas fa-grin",
size: int = 512,
Expand Down Expand Up @@ -82,7 +86,10 @@ def gen_fa_mask(


def gen_palette(palette: str):
"""Generates the corresponding palette function from `palettable`."""
"""
Generates the corresponding palette function from `palettable`.
"""

palette_split = palette.split(".")
palette_name = palette_split[-1]

Expand All @@ -98,7 +105,10 @@ def gen_palette(palette: str):


def gen_mask_array(icon_dir: str, invert_mask: bool, size: int):
"""Generates a numpy array of an icon mask."""
"""
Generates a numpy array of an icon mask.
"""

icon = Image.open(os.path.join(icon_dir, "icon.png"))

if isinstance(size, int):
Expand Down Expand Up @@ -127,7 +137,10 @@ def gen_gradient_mask(
gradient_dir: str = "horizontal",
invert_mask: bool = False,
):
"""Generates a gradient color mask from a specified palette."""
"""
Generates a gradient color mask from a specified palette.
"""

mask_array = gen_mask_array(icon_dir, invert_mask, size)
mask_array = np.float32(mask_array)

Expand All @@ -153,7 +166,10 @@ def gen_gradient_mask(


def color_to_rgb(color):
"""Converts a color to a RGB tuple from (0-255)."""
"""
Converts a color to a RGB tuple from (0-255).
"""

if isinstance(color, tuple):
# if a RGB tuple already
return color
Expand All @@ -168,6 +184,7 @@ def gen_stylecloud(
file_path: str = None,
size: int = 512,
icon_name: str = "fas fa-flag",
mask_file: str = None,
palette: str = "cartocolors.qualitative.Bold_5",
colors: Union[str, List[str]] = None,
background_color: str = "white",
Expand All @@ -178,6 +195,8 @@ def gen_stylecloud(
add_stopwords: bool = False,
icon_dir: str = ".temp",
output_name: str = "stylecloud.png",
plot_cloud: bool = False,
plot_size: int = (12,12),
gradient: str = None,
font_path: str = os.path.join(STATIC_PATH, "Staatliches-Regular.ttf"),
random_state: int = None,
Expand All @@ -186,41 +205,49 @@ def gen_stylecloud(
pro_icon_path: str = None,
pro_css_path: str = None,
):
"""Generates a stylecloud!
:param text: Input text. Best used if calling the function directly.
:param file_path: File path of the input text/CSV. Best used on the CLI.
:param size: Size (length and width in pixels) of the stylecloud.
:param icon_name: Icon Name for the stylecloud shape. (e.g. 'fas fa-grin')
:param palette: Color palette (via palettable)
:param colors: Custom color(s) for text (name or hex). Overrides palette.
:param background_color: Background color (name or hex).
:param max_font_size: Maximum font size in the stylecloud.
:param max_words: Maximum number of words to include in the stylecloud.
:param stopwords: Boolean to filter out common stopwords.
:param custom_stopwords: list of custom stopwords.
:param add_stopwords: Whether to use custom_stopwords to add to default
:param icon_dir: Temp directory to store the icon mask image.
:param output_name: Output file name of the stylecloud.
:param gradient: Direction of gradient. (if not None, will use gradient)
:param font_path: Path to .ttf file for font to use in stylecloud.
:param random_state: Controls random state of words and colors.
:param collocations: Whether to include collocations (bigrams) of two words.
:param invert_mask: Whether to invert the icon mask.
:param pro_icon_path: Path to Font Awesome Pro .ttf file if using FA Pro.
:param pro_css_path: Path to Font Awesome Pro .css file if using FA Pro.
"""
Generates a stylecloud!

Parameters
------------
text : Input text or Word Freq Dict. Best used if calling the function directly.
file_path : File path of the input text/CSV. Best used on the CLI.
size : Size (length and width in pixels) of the stylecloud.
icon_name : Icon Name for the stylecloud shape. (e.g. 'fas fa-grin')
mask_file : File path of Word Cloud Mask to use. Overrides icon_name, size.
palette : Color palette (via palettable)
colors : Custom color(s) for text (name or hex). Overrides palette.
background_color : Background color (name or hex).
max_font_size : Maximum font size in the stylecloud.
max_words : Maximum number of words to include in the stylecloud.
stopwords : Boolean to filter out common stopwords.
custom_stopwords : list of custom stopwords.
add_stopwords : Whether to use custom_stopwords to add to default
icon_dir : Temp directory to store the icon mask image.
output_name : Output file name of the stylecloud.
plot_cloud : Whether to plot the word cloud or not.
plot_size : Size (length and width in inches) of the plot.
gradient : Direction of gradient. (if not None, will use gradient)
font_path : Path to .ttf file for font to use in stylecloud.
random_state : Controls random state of words and colors.
collocations : Whether to include collocations (bigrams) of two words.
invert_mask : Whether to invert the icon mask.
pro_icon_path : Path to Font Awesome Pro .ttf file if using FA Pro.
pro_css_path : Path to Font Awesome Pro .css file if using FA Pro.
"""

assert any([text, file_path]), "Either text or file_path must be specified."

if file_path:
text = file_to_text(file_path)

gen_fa_mask(icon_name, size, icon_dir, pro_icon_path, pro_css_path)

if gradient and colors is None:
pal_colors, mask_array = gen_gradient_mask(
size, palette, icon_dir, gradient, invert_mask
)

else: # Color each word randomly from the palette
mask_array = gen_mask_array(icon_dir, invert_mask, size)
if colors:
Expand All @@ -243,6 +270,9 @@ def pal_colors(word, font_size, position, orientation, random_state, **kwargs):

if add_stopwords:
custom_stopwords.extend(STOPWORDS)

if mask_file!=None:
mask_array = np.array(Image.open(mask_file))

# cleanup icon folder
rmtree(icon_dir)
Expand All @@ -252,7 +282,7 @@ def pal_colors(word, font_size, position, orientation, random_state, **kwargs):
font_path=font_path,
max_words=max_words,
mask=mask_array,
stopwords=custom_stopwords if stopwords else None,
stopwords = custom_stopwords if stopwords else None,
max_font_size=max_font_size,
random_state=random_state,
collocations=collocations,
Expand All @@ -261,14 +291,22 @@ def pal_colors(word, font_size, position, orientation, random_state, **kwargs):
# generate word cloud
if isinstance(text, str):
wc.generate_from_text(text)
else: # i.e. a dict of word:value from a CSV
elif isinstance(text, dict):
if stopwords: # manually remove stopwords since otherwise ignored
text = {k: v for k, v in text.items() if k not in custom_stopwords}
wc.generate_from_frequencies(text)

wc.recolor(color_func=pal_colors, random_state=random_state)
wc.to_file(output_name)



# Plots the word cloud
if plot_cloud:
plt.figure(figsize=plot_size)
plt.imshow(wc, interpolation='bilinear')
plt.axis("off")

def stylecloud_cli(**kwargs):
"""Entrypoint for the stylecloud CLI."""
fire.Fire(gen_stylecloud)
"""
Entrypoint for the stylecloud CLI.
"""
fire.Fire(gen_stylecloud)