import numpy as np import random as _random from random import Random import colorsys from PIL import Image, ImageDraw, ImageFont, ImageFilter import re import os import sys from .ewc_core import IntegralGrid from .tokenization import process_tokens, unigrams_and_bigrams _FILE = os.path.dirname(__file__) _STOPWORDS_PATH = os.path.join(_FILE, "stopwords") def _load_stopwords(): if os.path.exists(_STOPWORDS_PATH): with open(_STOPWORDS_PATH, encoding="utf-8") as f: return set(map(str.strip, f)) # Fallback: minimal English stopwords so the module still works without the file return {"the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with", "is", "are", "was", "were", "be", "been", "it", "its", "this", "that", "i", "you", "he", "she", "we", "they"} STOPWORDS = _load_stopwords() # --------------------------------------------------------------------------- # Color functions (mirrors ref/word_cloud/wordcloud/wordcloud.py) # --------------------------------------------------------------------------- def random_color_func(word=None, font_size=None, position=None, orientation=None, font_path=None, random_state=None): """Random hue color generation (HSL, saturation=80%, lightness=50%).""" if random_state is None: random_state = Random() return "hsl(%d, 80%%, 50%%)" % random_state.randint(0, 255) class colormap_color_func: """Color function backed by a matplotlib colormap.""" def __init__(self, colormap): import matplotlib.pyplot as plt self.colormap = plt.get_cmap(colormap) def __call__(self, word, font_size, position, orientation, random_state=None, **kwargs): if random_state is None: random_state = Random() r, g, b, _ = np.maximum(0, 255 * np.array( self.colormap(random_state.uniform(0, 1)))) return "rgb({:.0f}, {:.0f}, {:.0f})".format(r, g, b) def get_single_color_func(color): """Return a color func that varies only the HSV value for a given color. Accepted values are PIL/Pillow color strings, e.g. 'deepskyblue', '#00b4d2'. """ from PIL import ImageColor old_r, old_g, old_b = ImageColor.getrgb(color) h, s, v = colorsys.rgb_to_hsv(old_r / 255., old_g / 255., old_b / 255.) def single_color_func(word=None, font_size=None, position=None, orientation=None, font_path=None, random_state=None): if random_state is None: random_state = Random() r, g, b = colorsys.hsv_to_rgb(h, s, random_state.uniform(0.2, 1)) return "rgb({:.0f}, {:.0f}, {:.0f})".format(r * 255, g * 255, b * 255) return single_color_func import logging class EfficientWordCloud: """ EfficientWordCloud Generation Class. Uses C++ backend (ewc_core) for high-performance collision detection. Optimized for high-resolution generation (4k/8k+). Parameters ---------- width, height : int Canvas size (ignored when mask is provided). mask : ndarray or None Shape mask. White pixels (255) are treated as blocked, others as free. font_path : str or None Path to a TrueType font file. max_words : int Maximum number of words to place. min_font_size : int Smallest font size to use. max_font_size : int or None Largest font size. Derived automatically when None. background_color : color PIL-compatible background color. prefer_horizontal : float Probability a word is placed horizontally (0–1). mode : str PIL image mode ('RGB', 'RGBA', …). use_spiral_search : bool Use center-out sorted search (True) or reservoir sampling (False). scale : float Scaling factor between layout computation and final rendering. ``scale=2`` means the output image is 2× the canvas size in each dimension while layout is still computed at base resolution. Equivalent to ref's ``scale`` parameter. contour_width : float If > 0 and mask is set, draw the mask contour on the output image. contour_color : color PIL-compatible color for the mask contour (default 'black'). margin : int Pixel gap between words. stopwords : set of str or None Words to exclude when processing text. Defaults to built-in STOPWORDS. regexp : str or None Override the regex used to tokenize text (default ``r"\\w[\\w']+"``). collocations : bool Whether to detect bigrams (default True). collocation_threshold : int Dunning score threshold for bigrams (default 30). normalize_plurals : bool Strip trailing 's' to merge plurals (default True). include_numbers : bool Keep numeric tokens when processing text (default False). min_word_length : int Minimum character length for a token to be kept (default 0). color_func : callable or None ``color_func(word, font_size, position, orientation, font_path, random_state) -> color``. Overrides *colormap*. colormap : str or matplotlib colormap or None Matplotlib colormap used when *color_func* is None. random_state : int, Random, or None Seed for reproducibility. repeat : bool If True, repeat words (with decreasing weight) until *max_words* or *min_font_size* is reached (default False). relative_scaling : float (0–1) How much word frequency (vs rank) influences font size. 0 = rank only, 1 = fully frequency-driven. When *repeat* is True, defaults to 0. font_step : int Step size when reducing font size to find a fit. """ def __init__(self, width=400, height=200, mask=None, font_path=None, max_words=200, min_font_size=4, max_font_size=None, background_color="black", prefer_horizontal=0.9, mode="RGB", use_spiral_search=True, scale=1, contour_width=0, contour_color="black", margin=2, color_func=None, colormap=None, random_state=None, relative_scaling="auto", font_step=1, repeat=False, stopwords=None, regexp=None, collocations=True, collocation_threshold=30, normalize_plurals=True, include_numbers=False, min_word_length=0): self.width = width self.height = height self.mask = mask self.font_path = font_path self.max_words = max_words self.min_font_size = min_font_size self.max_font_size = max_font_size self.background_color = background_color self.prefer_horizontal = prefer_horizontal self.mode = mode self.use_spiral_search = use_spiral_search self.scale = scale self.contour_width = contour_width self.contour_color = contour_color self.repeat = repeat # relative_scaling default mirrors ref: 0 when repeat, else 0.5 if relative_scaling == "auto": self.relative_scaling = 0 if repeat else 0.5 else: self.relative_scaling = relative_scaling self.margin = margin self.font_step = font_step self.stopwords = stopwords if stopwords is not None else STOPWORDS self.regexp = regexp self.collocations = collocations self.collocation_threshold = collocation_threshold self.normalize_plurals = normalize_plurals self.include_numbers = include_numbers self.min_word_length = min_word_length # Random state if isinstance(random_state, int): self.random_state = Random(random_state) elif random_state is None: self.random_state = Random() else: self.random_state = random_state # Color function if color_func is not None: self.color_func = color_func elif colormap is not None: self.color_func = colormap_color_func(colormap) else: self.color_func = random_color_func self.layout_ = [] # Handle mask if self.mask is not None: self.width = self.mask.shape[1] self.height = self.mask.shape[0] if self.mask.dtype == bool: self.boolean_mask = self.mask.astype(np.uint8) * 255 elif self.mask.ndim == 3: # White pixels (all channels == 255) are blocked self.boolean_mask = np.where( np.all(self.mask[:, :, :3] == 255, axis=-1), 255, 0 ).astype(np.uint8) else: self.boolean_mask = self.mask.astype(np.uint8) else: self.boolean_mask = np.zeros((self.height, self.width), dtype=np.uint8) # Initialize C++ grid (>0 = occupied) self.grid = IntegralGrid(self.boolean_mask, self.height, self.width) def generate_from_frequencies(self, frequencies, max_font_size=None): """Generate word cloud from a dict of {word: frequency}. Parameters ---------- frequencies : dict max_font_size : int or None Override self.max_font_size for this call (used internally for the automatic font-size estimation). """ sorted_freq = sorted(frequencies.items(), key=lambda x: x[1], reverse=True) if not sorted_freq: raise ValueError("Need at least 1 word to generate a word cloud.") sorted_freq = sorted_freq[:self.max_words] # Normalize so the top word = 1.0 max_freq = float(sorted_freq[0][1]) sorted_freq = [(w, f / max_freq) for w, f in sorted_freq] # --- repeat: pad list up to max_words with down-weighted copies --- if self.repeat and len(sorted_freq) < self.max_words: import math times_extend = math.ceil(self.max_words / len(sorted_freq)) - 1 base = list(sorted_freq) downweight = base[-1][1] for i in range(times_extend): factor = downweight ** (i + 1) sorted_freq.extend([(w, f * factor) for w, f in base]) sorted_freq = sorted_freq[:self.max_words] self.words_ = dict(sorted_freq) # --- auto max_font_size estimation (mirrors ref) --- effective_max = max_font_size if max_font_size is not None else self.max_font_size if effective_max is None: if len(sorted_freq) == 1: effective_max = self.height else: # Trial run with just the first 2 words to estimate a good max size. # We must reinitialize the grid after so it is clean for the real run. _repeat_bak = self.repeat self.repeat = False self.generate_from_frequencies(dict(sorted_freq[:2]), max_font_size=self.height) self.repeat = _repeat_bak sizes = [s for _, s, *_ in self.layout_] try: effective_max = int(2 * sizes[0] * sizes[1] / (sizes[0] + sizes[1])) except (IndexError, ZeroDivisionError): effective_max = sizes[0] if sizes else self.height # Reinitialize the C++ grid so the trial run does not consume space self.grid = IntegralGrid(self.boolean_mask, self.height, self.width) rs = self.random_state # No PIL image needed during placement — C++ canvas handles collision. # The PIL image is constructed lazily in to_image(). self.layout_ = [] font_size = int(effective_max) last_freq = 1.0 # E1: font object cache {size -> ImageFont} font_cache: dict = {} def _get_font(size): if size not in font_cache: try: font_cache[size] = ImageFont.truetype(self.font_path, size) except IOError: font_cache[size] = ImageFont.load_default() return font_cache[size] def _query(qh, qw): return self.grid.query_direct(qh, qw, rs.randint(0, 2**31)) # v4: Ref-like linear step-down placement with bitmap occupancy # After each word placement, stamp glyph bitmap into C++ canvas # and rebuild integral for pixel-accurate collision detection. # No PIL image drawn during placement — to_image() renders later. # Dummy draw for textbbox measurement _measure_img = Image.new("L", (1, 1)) _measure_draw = ImageDraw.Draw(_measure_img) for idx, (word, freq) in enumerate(sorted_freq): if freq == 0: continue # Relative-scaling font size adjustment (mirrors ref logic) rs_val = self.relative_scaling if rs_val != 0: font_size = int(round( (rs_val * (freq / float(last_freq)) + (1 - rs_val)) * font_size )) if rs.random() < self.prefer_horizontal: orientation = None else: orientation = Image.ROTATE_90 tried_other_orientation = False while True: if font_size < self.min_font_size: break font = _get_font(font_size) transposed = ImageFont.TransposedFont(font, orientation=orientation) bbox = _measure_draw.textbbox((0, 0), word, font=transposed) tw = bbox[2] - bbox[0] th = bbox[3] - bbox[1] qh = th + self.margin qw = tw + self.margin pos = _query(qh, qw) if pos is not None: break # No position found — try alternate orientation, then reduce size if not tried_other_orientation and self.prefer_horizontal < 1: orientation = Image.ROTATE_90 if orientation is None else None tried_other_orientation = True else: font_size -= self.font_step orientation = None tried_other_orientation = False if font_size < self.min_font_size: # Canvas full — no more words can fit break y, x = pos # Adjust position for margin (like ref: x,y += margin // 2) draw_x = x + self.margin // 2 draw_y = y + self.margin // 2 # Get glyph bitmap and stamp into C++ canvas font = _get_font(font_size) transposed = ImageFont.TransposedFont(font, orientation=orientation) glyph_mask = transposed.getmask(word, mode="L") gw, gh = glyph_mask.size glyph_arr = np.frombuffer(bytes(glyph_mask), dtype=np.uint8).reshape(gh, gw) # Stamp glyph into C++ canvas + rebuild integral self.grid.stamp_and_rebuild(glyph_arr, gh, gw, draw_y, draw_x) color = self.color_func( word=word, font_size=font_size, position=(y, x), orientation=orientation, font_path=self.font_path, random_state=rs, ) self.layout_.append((word, font_size, (y, x), orientation, color)) last_freq = freq return self def generate(self, text): """Generate word cloud from raw text (calls process_text + generate_from_frequencies).""" return self.generate_from_text(text) def generate_from_text(self, text): """Process *text* into word frequencies, then generate the word cloud.""" words = self.process_text(text) self.generate_from_frequencies(words) return self def process_text(self, text): """Tokenize *text* and return ``{word: count}`` after filtering. Applies regexp splitting, stopword removal, number/length filters, plural normalization, and optional bigram collocation detection. """ min_len = self.min_word_length pattern = r"\w[\w']+" if min_len <= 1 else r"\w[\w']+" regexp = self.regexp if self.regexp is not None else pattern words = re.findall(regexp, text) # Strip possessive 's words = [w[:-2] if w.lower().endswith("'s") else w for w in words] if not self.include_numbers: words = [w for w in words if not w.isdigit()] if self.min_word_length: words = [w for w in words if len(w) >= self.min_word_length] stopwords_lower = {s.lower() for s in self.stopwords} if self.collocations: word_counts = unigrams_and_bigrams( words, stopwords_lower, normalize_plurals=self.normalize_plurals, collocation_threshold=self.collocation_threshold, ) else: words = [w for w in words if w.lower() not in stopwords_lower] word_counts, _ = process_tokens(words, self.normalize_plurals) self.words_ = word_counts return word_counts def to_image(self): """Render the layout to a PIL Image, respecting *scale* and *contour*.""" s = self.scale out_w = int(self.width * s) out_h = int(self.height * s) img = Image.new(self.mode, (out_w, out_h), self.background_color) draw = ImageDraw.Draw(img) for word, size, (y, x), orient, color in self.layout_: try: font = ImageFont.truetype(self.font_path, int(size * s)) except Exception: font = ImageFont.load_default() transposed_font = ImageFont.TransposedFont(font, orientation=orient) draw.text((int(x * s), int(y * s)), word, font=transposed_font, fill=color) return self._draw_contour(img) def _draw_contour(self, img): """Draw mask contour on *img* if contour_width > 0.""" if self.mask is None or self.contour_width == 0: return img # Build boolean mask: True where drawing area (not blocked) if self.mask.ndim == 3: blocked = np.all(self.mask[:, :, :3] == 255, axis=-1) else: blocked = self.mask == 255 mask_uint8 = (~blocked).astype(np.uint8) * 255 contour = Image.fromarray(mask_uint8) contour = contour.resize(img.size) contour = contour.filter(ImageFilter.FIND_EDGES) contour_arr = np.array(contour) # Zero out border pixels so edges aren't drawn at image boundary contour_arr[[0, -1], :] = 0 contour_arr[:, [0, -1]] = 0 # Gaussian blur controls perceived width (divide by 10 for sub-pixel) radius = self.contour_width / 10 contour = Image.fromarray(contour_arr) contour = contour.filter(ImageFilter.GaussianBlur(radius=radius)) contour_arr = np.array(contour) > 0 contour_3d = np.dstack([contour_arr] * 3) result = np.array(img.convert("RGB")) * ~contour_3d if self.contour_color != "black": color_img = Image.new("RGB", img.size, self.contour_color) result = result + np.array(color_img) * contour_3d out = Image.fromarray(result.astype(np.uint8)) if self.mode == "RGBA": out = out.convert("RGBA") return out def to_array(self, copy=None): """Return the word cloud as a numpy ndarray (H x W x channels).""" image = self.to_image() if copy is None: return np.asarray(image) try: return np.asarray(image, copy=copy) except TypeError: return np.asarray(image) def __array__(self, copy=None): return self.to_array(copy=copy) def to_file(self, filename): """Save to *filename* and return self (for chaining).""" img = self.to_image() img.save(filename, optimize=True) return self def recolor(self, random_state=None, color_func=None, colormap=None): """Re-apply colors to the current layout without regenerating it. Parameters ---------- random_state : int, Random, or None color_func : callable or None colormap : str or matplotlib colormap or None """ if isinstance(random_state, int): random_state = Random(random_state) elif random_state is None: random_state = Random() if color_func is None: if colormap is not None: color_func = colormap_color_func(colormap) else: color_func = self.color_func self.layout_ = [ (word, font_size, position, orientation, color_func(word=word, font_size=font_size, position=position, orientation=orientation, font_path=self.font_path, random_state=random_state)) for word, font_size, position, orientation, _ in self.layout_ ] return self def to_svg(self, filename=None): """Export as SVG with scale, correct rotation transforms and XML escaping. Parameters ---------- filename : str or None If given, write to this file. Otherwise return the SVG string. """ from xml.sax import saxutils s = self.scale out_w = int(self.width * s) out_h = int(self.height * s) # Derive font metadata from the actual font file try: _font_probe = ImageFont.truetype(self.font_path, 12) raw_family, raw_style = _font_probe.getname() except Exception: raw_family, raw_style = "sans-serif", "Regular" raw_style_lower = raw_style.lower() font_weight = "bold" if "bold" in raw_style_lower else "normal" if "italic" in raw_style_lower: font_style = "italic" elif "oblique" in raw_style_lower: font_style = "oblique" else: font_style = "normal" font_family = repr(raw_family) lines = [ f'', f'', ] if self.background_color is not None: lines.append( f'' ) for word, size, (y, x), orient, color in self.layout_: scaled_size = int(size * s) try: font = ImageFont.truetype(self.font_path, scaled_size) except Exception: font = ImageFont.load_default() (size_x, size_y), (offset_x, offset_y) = font.font.getsize(word) ascent, _ = font.getmetrics() min_x = -offset_x max_x = size_x - offset_x max_y = ascent - offset_y sx = int(x * s) sy = int(y * s) if orient == Image.ROTATE_90: tx = sx + max_y ty = sy + max_x - min_x transform = f"translate({tx},{ty}) rotate(-90)" else: tx = sx + min_x ty = sy + max_y transform = f"translate({tx},{ty})" lines.append( f'{saxutils.escape(word)}' ) lines.append("") svg_str = "\n".join(lines) if filename is not None: with open(filename, "w", encoding="utf-8") as f: f.write(svg_str) return svg_str