Initial project baseline
This commit is contained in:
@@ -0,0 +1,8 @@
|
||||
from .wordcloud import (
|
||||
EfficientWordCloud,
|
||||
STOPWORDS,
|
||||
random_color_func,
|
||||
colormap_color_func,
|
||||
get_single_color_func,
|
||||
)
|
||||
from .color_from_image import ImageColorGenerator
|
||||
@@ -0,0 +1,55 @@
|
||||
# Ported from ref/word_cloud/wordcloud/color_from_image.py (MIT License)
|
||||
import numpy as np
|
||||
from PIL import ImageFont
|
||||
|
||||
|
||||
class ImageColorGenerator:
|
||||
"""Color generator that samples colors from a reference image.
|
||||
|
||||
Each word is colored using the mean RGB of the region it occupies in
|
||||
*image*. Pass an instance as ``color_func`` to ``EfficientWordCloud``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
image : ndarray, shape (H, W, 3) or (H, W, 4)
|
||||
Reference color image. Should match the word-cloud canvas size.
|
||||
default_color : tuple (R, G, B) or None
|
||||
Fallback color when the word region falls outside *image*.
|
||||
If None, a ValueError is raised instead.
|
||||
"""
|
||||
|
||||
def __init__(self, image, default_color=None):
|
||||
if image.ndim not in (2, 3):
|
||||
raise ValueError(
|
||||
"ImageColorGenerator needs an image with ndim 2 or 3, "
|
||||
"got %d" % image.ndim
|
||||
)
|
||||
if image.ndim == 3 and image.shape[2] not in (3, 4):
|
||||
raise ValueError(
|
||||
"A color image must have 3 or 4 channels, got %d"
|
||||
% image.shape[2]
|
||||
)
|
||||
self.image = image
|
||||
self.default_color = default_color
|
||||
|
||||
def __call__(self, word, font_size, font_path, position, orientation,
|
||||
**kwargs):
|
||||
"""Return the mean color of the image patch under the word bounding box."""
|
||||
font = ImageFont.truetype(font_path, font_size)
|
||||
from PIL import ImageFont as _IF
|
||||
transposed_font = _IF.TransposedFont(font, orientation=orientation)
|
||||
box = transposed_font.getbbox(word) # (left, top, right, bottom)
|
||||
x, y = position # (row, col) i.e. (y, x) in PIL
|
||||
patch = self.image[x:x + box[2], y:y + box[3]]
|
||||
if patch.ndim == 3:
|
||||
patch = patch[:, :, :3] # drop alpha
|
||||
reshape = patch.reshape(-1, 3)
|
||||
if not np.all(np.array(reshape.shape) > 0):
|
||||
if self.default_color is None:
|
||||
raise ValueError(
|
||||
"ImageColorGenerator: word region is outside the image. "
|
||||
"Pass default_color to suppress this error."
|
||||
)
|
||||
return "rgb(%d, %d, %d)" % tuple(self.default_color)
|
||||
color = np.mean(reshape, axis=0)
|
||||
return "rgb(%d, %d, %d)" % tuple(color.astype(int))
|
||||
@@ -0,0 +1,829 @@
|
||||
/*
|
||||
* EfficientWordCloud Core (ewc_core) v3.0
|
||||
*
|
||||
* v3 changes:
|
||||
* - Direct pixel-grid scan (query_direct) — like ref, O(1) integral image
|
||||
* check per position, sequential memory access, cache-friendly
|
||||
* - Parallel reservoir sampling with thread pool (no std::async overhead)
|
||||
* - Batch query API to eliminate Python↔C++ per-word overhead
|
||||
* - Partial integral rebuild (only affected subregion)
|
||||
*/
|
||||
|
||||
#include <Python.h>
|
||||
#include <vector>
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdlib>
|
||||
#include <cstdint>
|
||||
#include <utility>
|
||||
#include <mutex>
|
||||
#include <shared_mutex>
|
||||
#include <thread>
|
||||
#include <future>
|
||||
#include <atomic>
|
||||
#include <random>
|
||||
#include <functional>
|
||||
#include <numeric>
|
||||
|
||||
// ==========================================
|
||||
// Thread Pool (avoid per-query thread creation)
|
||||
// ==========================================
|
||||
|
||||
class ThreadPool {
|
||||
std::vector<std::thread> workers;
|
||||
std::vector<std::function<void()>> tasks;
|
||||
std::mutex mtx;
|
||||
std::condition_variable cv;
|
||||
std::atomic<bool> stop{false};
|
||||
|
||||
public:
|
||||
ThreadPool(unsigned int n) {
|
||||
for (unsigned int i = 0; i < n; ++i) {
|
||||
workers.emplace_back([this] {
|
||||
while (true) {
|
||||
std::function<void()> task;
|
||||
{
|
||||
std::unique_lock<std::mutex> lock(mtx);
|
||||
cv.wait(lock, [this] { return stop.load() || !tasks.empty(); });
|
||||
if (stop.load() && tasks.empty()) return;
|
||||
task = std::move(tasks.back());
|
||||
tasks.pop_back();
|
||||
}
|
||||
task();
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
template<class F>
|
||||
std::future<typename std::invoke_result<F>::type> submit(F&& f) {
|
||||
using R = typename std::invoke_result<F>::type;
|
||||
auto p = std::make_shared<std::promise<R>>();
|
||||
auto fut = p->get_future();
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mtx);
|
||||
if constexpr (std::is_void_v<R>) {
|
||||
tasks.push_back([p, f=std::forward<F>(f)]() mutable {
|
||||
try { f(); p->set_value(); }
|
||||
catch (...) { p->set_exception(std::current_exception()); }
|
||||
});
|
||||
} else {
|
||||
tasks.push_back([p, f=std::forward<F>(f)]() mutable {
|
||||
try { p->set_value(f()); }
|
||||
catch (...) { p->set_exception(std::current_exception()); }
|
||||
});
|
||||
}
|
||||
}
|
||||
cv.notify_one();
|
||||
return fut;
|
||||
}
|
||||
|
||||
~ThreadPool() {
|
||||
stop.store(true);
|
||||
cv.notify_all();
|
||||
for (auto& w : workers) w.join();
|
||||
}
|
||||
};
|
||||
|
||||
static ThreadPool& get_pool() {
|
||||
static unsigned int n = std::max(2u, std::thread::hardware_concurrency());
|
||||
static ThreadPool pool(n);
|
||||
return pool;
|
||||
}
|
||||
|
||||
static unsigned int get_num_threads() {
|
||||
static unsigned int n = std::max(2u, std::thread::hardware_concurrency());
|
||||
return n;
|
||||
}
|
||||
|
||||
// ==========================================
|
||||
// Core Data Structure
|
||||
// ==========================================
|
||||
|
||||
class IntegralGrid {
|
||||
public:
|
||||
uint32_t* data;
|
||||
int width;
|
||||
int height;
|
||||
mutable std::shared_mutex mutex_;
|
||||
|
||||
struct Rect {
|
||||
int r;
|
||||
int c;
|
||||
int h;
|
||||
int w;
|
||||
};
|
||||
|
||||
// valid coordinates sorted by distance to center (for sorted/spiral search)
|
||||
std::vector<std::pair<int, int>> valid_coords;
|
||||
|
||||
// Lazy update buffers
|
||||
std::vector<int32_t> diff;
|
||||
std::vector<Rect> recent_rects;
|
||||
int dirty_count = 0;
|
||||
int rebuild_interval = 16; // v3: more frequent rebuilds (was 64), since rebuild is cheaper now
|
||||
|
||||
IntegralGrid(int h, int w) : width(w), height(h) {
|
||||
data = new uint32_t[width * height]();
|
||||
diff.resize(width * height, 0);
|
||||
}
|
||||
|
||||
~IntegralGrid() {
|
||||
if (data) delete[] data;
|
||||
}
|
||||
|
||||
void init_from_buffer(unsigned char* raw_mask, int h, int w) {
|
||||
int center_y = h / 2;
|
||||
int center_x = w / 2;
|
||||
valid_coords.reserve(h * w / 2);
|
||||
|
||||
// Initialize canvas from mask
|
||||
ensure_canvas();
|
||||
|
||||
for (int i = 0; i < h; ++i) {
|
||||
uint32_t row_sum = 0;
|
||||
for (int j = 0; j < w; ++j) {
|
||||
bool is_blocked = (raw_mask[i * w + j] > 0);
|
||||
uint32_t val = is_blocked ? 1 : 0;
|
||||
row_sum += val;
|
||||
uint32_t prev_row = (i > 0) ? data[(i - 1) * width + j] : 0;
|
||||
data[i * width + j] = row_sum + prev_row;
|
||||
if (is_blocked) {
|
||||
canvas[i * width + j] = 1;
|
||||
}
|
||||
if (!is_blocked) {
|
||||
valid_coords.push_back({i, j});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::sort(valid_coords.begin(), valid_coords.end(),
|
||||
[center_y, center_x](const std::pair<int, int>& a, const std::pair<int, int>& b) {
|
||||
long da = (long)(a.first - center_y)*(a.first - center_y) + (long)(a.second - center_x)*(a.second - center_x);
|
||||
long db = (long)(b.first - center_y)*(b.first - center_y) + (long)(b.second - center_x)*(b.second - center_x);
|
||||
return da < db;
|
||||
}
|
||||
);
|
||||
|
||||
std::fill(diff.begin(), diff.end(), 0);
|
||||
recent_rects.clear();
|
||||
dirty_count = 0;
|
||||
}
|
||||
|
||||
void reorder_stratified(int bands) {
|
||||
std::unique_lock<std::shared_mutex> lock(mutex_);
|
||||
const size_t len = valid_coords.size();
|
||||
if (bands <= 1 || len == 0) return;
|
||||
if ((size_t)bands > len) bands = (int)len;
|
||||
|
||||
std::vector<size_t> band_sizes(bands, 0);
|
||||
size_t base = len / bands;
|
||||
size_t rem = len % bands;
|
||||
for (int b = 0; b < bands; ++b)
|
||||
band_sizes[b] = base + (b < (int)rem ? 1 : 0);
|
||||
|
||||
std::vector<size_t> band_starts(bands, 0);
|
||||
size_t offset = 0, max_band_size = 0;
|
||||
for (int b = 0; b < bands; ++b) {
|
||||
band_starts[b] = offset;
|
||||
offset += band_sizes[b];
|
||||
if (band_sizes[b] > max_band_size) max_band_size = band_sizes[b];
|
||||
}
|
||||
|
||||
std::vector<std::pair<int, int>> reordered;
|
||||
reordered.reserve(len);
|
||||
for (size_t i = 0; i < max_band_size; ++i)
|
||||
for (int b = 0; b < bands; ++b)
|
||||
if (i < band_sizes[b])
|
||||
reordered.push_back(valid_coords[band_starts[b] + i]);
|
||||
valid_coords.swap(reordered);
|
||||
}
|
||||
|
||||
// O(1) rectangle emptiness check via integral image + recent-rect overlap
|
||||
inline uint32_t get_area_sum(int r, int c, int h, int w) const {
|
||||
int r_bottom = r + h - 1;
|
||||
int c_right = c + w - 1;
|
||||
int r_top = r - 1;
|
||||
int c_left = c - 1;
|
||||
uint32_t A = (r_top >= 0 && c_left >= 0) ? data[r_top * width + c_left] : 0;
|
||||
uint32_t B = (r_top >= 0) ? data[r_top * width + c_right] : 0;
|
||||
uint32_t C = (c_left >= 0) ? data[r_bottom * width + c_left] : 0;
|
||||
uint32_t D = data[r_bottom * width + c_right];
|
||||
uint32_t base_sum = D - B - C + A;
|
||||
if (base_sum > 0) return base_sum;
|
||||
|
||||
if (!recent_rects.empty()) {
|
||||
int r2 = r + h;
|
||||
int c2 = c + w;
|
||||
for (const auto& rect : recent_rects) {
|
||||
int rr2 = rect.r + rect.h;
|
||||
int cc2 = rect.c + rect.w;
|
||||
if (r < rr2 && r2 > rect.r && c < cc2 && c2 > rect.c) return 1;
|
||||
}
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
// O(1) fast check without recent_rects (for direct scan where we'll call flush first)
|
||||
inline uint32_t get_area_sum_fast(int r, int c, int h, int w) const {
|
||||
int r_bottom = r + h - 1;
|
||||
int c_right = c + w - 1;
|
||||
int r_top = r - 1;
|
||||
int c_left = c - 1;
|
||||
uint32_t A = (r_top >= 0 && c_left >= 0) ? data[r_top * width + c_left] : 0;
|
||||
uint32_t B = (r_top >= 0) ? data[r_top * width + c_right] : 0;
|
||||
uint32_t C = (c_left >= 0) ? data[r_bottom * width + c_left] : 0;
|
||||
uint32_t D = data[r_bottom * width + c_right];
|
||||
return D - B - C + A;
|
||||
}
|
||||
|
||||
void rebuild_integral() {
|
||||
if (dirty_count <= 0) return;
|
||||
|
||||
std::vector<int32_t> delta_prev_row(width, 0);
|
||||
std::vector<int32_t> delta_row(width, 0);
|
||||
std::vector<uint32_t> integral_prev_row(width, 0);
|
||||
|
||||
for (int i = 0; i < height; ++i) {
|
||||
uint32_t row_sum = 0;
|
||||
for (int j = 0; j < width; ++j) {
|
||||
int idx = i * width + j;
|
||||
int32_t up = (i > 0) ? delta_prev_row[j] : 0;
|
||||
int32_t left = (j > 0) ? delta_row[j - 1] : 0;
|
||||
int32_t up_left = (i > 0 && j > 0) ? delta_prev_row[j - 1] : 0;
|
||||
int32_t delta = diff[idx] + up + left - up_left;
|
||||
delta_row[j] = delta;
|
||||
row_sum += (uint32_t)delta;
|
||||
uint32_t prev = (i > 0) ? integral_prev_row[j] : 0;
|
||||
uint32_t di = row_sum + prev;
|
||||
data[idx] += di;
|
||||
integral_prev_row[j] = di;
|
||||
}
|
||||
delta_prev_row.swap(delta_row);
|
||||
}
|
||||
|
||||
std::fill(diff.begin(), diff.end(), 0);
|
||||
recent_rects.clear();
|
||||
dirty_count = 0;
|
||||
}
|
||||
|
||||
// Force flush: rebuild integral so get_area_sum_fast is safe
|
||||
void flush() {
|
||||
if (dirty_count > 0) rebuild_integral();
|
||||
}
|
||||
|
||||
// v4: Rebuild integral image from raw pixel array (bitmap occupancy like ref)
|
||||
// raw_pixels: H×W uint8 array where >0 = occupied
|
||||
void rebuild_from_bitmap(const unsigned char* raw_pixels) {
|
||||
// Reset lazy update state since we're rebuilding everything
|
||||
std::fill(diff.begin(), diff.end(), 0);
|
||||
recent_rects.clear();
|
||||
dirty_count = 0;
|
||||
|
||||
for (int i = 0; i < height; ++i) {
|
||||
uint32_t row_sum = 0;
|
||||
for (int j = 0; j < width; ++j) {
|
||||
uint32_t val = (raw_pixels[i * width + j] > 0) ? 1 : 0;
|
||||
row_sum += val;
|
||||
uint32_t prev_row = (i > 0) ? data[(i - 1) * width + j] : 0;
|
||||
data[i * width + j] = row_sum + prev_row;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Persistent canvas bitmap (C++ side) — avoids PIL→numpy conversion overhead
|
||||
std::vector<uint8_t> canvas;
|
||||
|
||||
void ensure_canvas() {
|
||||
if (canvas.empty()) {
|
||||
canvas.resize(width * height, 0);
|
||||
}
|
||||
}
|
||||
|
||||
// Stamp a small glyph bitmap onto the canvas at position (pos_r, pos_c)
|
||||
// glyph: gh×gw uint8 array where >0 = occupied pixel
|
||||
void stamp_glyph(const unsigned char* glyph, int gh, int gw, int pos_r, int pos_c) {
|
||||
ensure_canvas();
|
||||
for (int i = 0; i < gh; ++i) {
|
||||
int ri = pos_r + i;
|
||||
if (ri < 0 || ri >= height) continue;
|
||||
for (int j = 0; j < gw; ++j) {
|
||||
int cj = pos_c + j;
|
||||
if (cj < 0 || cj >= width) continue;
|
||||
if (glyph[i * gw + j] > 0) {
|
||||
canvas[ri * width + cj] = 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Rebuild integral from the internal canvas (partial from pos_r, pos_c)
|
||||
void rebuild_from_canvas(int pos_r, int pos_c) {
|
||||
if (canvas.empty()) return;
|
||||
rebuild_from_bitmap_partial(canvas.data(), pos_r, pos_c);
|
||||
}
|
||||
|
||||
// v4: Partial integral rebuild from position (pos_r, pos_c) downward
|
||||
// Optimized: use row-sum approach (like rebuild_from_bitmap) for the partial region
|
||||
void rebuild_from_bitmap_partial(const unsigned char* raw_pixels, int pos_r, int pos_c) {
|
||||
// Reset lazy update state
|
||||
dirty_count = 0;
|
||||
recent_rects.clear();
|
||||
|
||||
// Clamp to valid range
|
||||
if (pos_r < 0) pos_r = 0;
|
||||
if (pos_c < 0) pos_c = 0;
|
||||
|
||||
// For each row from pos_r, recompute integral from pos_c onward
|
||||
// Using the fast row-sum approach
|
||||
for (int i = pos_r; i < height; ++i) {
|
||||
const unsigned char* pixel_row = raw_pixels + i * width;
|
||||
uint32_t* data_row = data + i * width;
|
||||
|
||||
// Compute row_sum for columns [0, pos_c) from existing data
|
||||
// row_sum at pos_c = data[i][pos_c-1] - data[i-1][pos_c-1] (if prev row available)
|
||||
// But we need the actual row prefix sum, not the integral
|
||||
// Instead: use the standard formula: integral(i,j) = pixel(i,j) + left + up - diag
|
||||
// But optimize by accumulating row_sum separately
|
||||
|
||||
// row_sum = sum of pixel_row[pos_c..j] for current row
|
||||
uint32_t row_sum;
|
||||
if (pos_c > 0) {
|
||||
// Get the row prefix sum at pos_c-1 from existing integral
|
||||
uint32_t prev_row_integral = (i > 0) ? data[(i-1) * width + (pos_c - 1)] : 0;
|
||||
row_sum = data_row[pos_c - 1] - prev_row_integral;
|
||||
// But wait - data_row[pos_c-1] might be stale for rows > pos_r
|
||||
// Actually for i == pos_r, row pos_r-1 hasn't been modified, so data[pos_r-1] is correct
|
||||
// For i > pos_r, data[i-1] has been updated in previous iteration, so it's correct
|
||||
// And data_row[pos_c-1] is unchanged (we don't modify columns < pos_c)
|
||||
// So row_sum = integral[i][pos_c-1] - integral[i-1][pos_c-1]
|
||||
// This gives us the row prefix sum for columns [0, pos_c-1]
|
||||
} else {
|
||||
row_sum = 0;
|
||||
}
|
||||
|
||||
for (int j = pos_c; j < width; ++j) {
|
||||
row_sum += (pixel_row[j] > 0) ? 1 : 0;
|
||||
uint32_t prev_row = (i > 0) ? data[(i - 1) * width + j] : 0;
|
||||
data_row[j] = row_sum + prev_row;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void update_rect_add(int r, int c, int h, int w) {
|
||||
int end_r = std::min(height, r + h);
|
||||
int end_c = std::min(width, c + w);
|
||||
if (r < 0 || c < 0 || end_r <= r || end_c <= c) return;
|
||||
|
||||
recent_rects.push_back({r, c, end_r - r, end_c - c});
|
||||
|
||||
diff[r * width + c] += 1;
|
||||
if (end_r < height) diff[end_r * width + c] -= 1;
|
||||
if (end_c < width) diff[r * width + end_c] -= 1;
|
||||
if (end_r < height && end_c < width) diff[end_r * width + end_c] += 1;
|
||||
|
||||
dirty_count++;
|
||||
if (dirty_count >= rebuild_interval) rebuild_integral();
|
||||
}
|
||||
|
||||
// =========================================================
|
||||
// v3: Direct pixel-grid scan (like ref's query_integral_image)
|
||||
// Scans all (i,j) where 0<=i<=H-bh, 0<=j<=W-bw in memory order
|
||||
// Two-pass: count hits, then pick random one
|
||||
// =========================================================
|
||||
|
||||
struct DirectResult {
|
||||
bool found;
|
||||
int y, x;
|
||||
};
|
||||
|
||||
// Count valid positions in row range [row_start, row_end)
|
||||
uint64_t count_valid_rows(int row_start, int row_end, int box_h, int box_w) const {
|
||||
uint64_t count = 0;
|
||||
int max_col = width - box_w;
|
||||
for (int i = row_start; i < row_end; ++i) {
|
||||
for (int j = 0; j <= max_col; ++j) {
|
||||
if (get_area_sum_fast(i, j, box_h, box_w) == 0)
|
||||
++count;
|
||||
}
|
||||
}
|
||||
return count;
|
||||
}
|
||||
|
||||
// Pick the nth valid position in row range [row_start, row_end)
|
||||
std::pair<int,int> pick_nth_valid(int row_start, int row_end, int box_h, int box_w, uint64_t target) const {
|
||||
uint64_t count = 0;
|
||||
int max_col = width - box_w;
|
||||
for (int i = row_start; i < row_end; ++i) {
|
||||
for (int j = 0; j <= max_col; ++j) {
|
||||
if (get_area_sum_fast(i, j, box_h, box_w) == 0) {
|
||||
if (count == target) return {i, j};
|
||||
++count;
|
||||
}
|
||||
}
|
||||
}
|
||||
return {-1, -1};
|
||||
}
|
||||
|
||||
DirectResult query_direct(int box_h, int box_w, uint32_t seed) {
|
||||
// Always flush before scanning
|
||||
flush();
|
||||
|
||||
int max_row = height - box_h;
|
||||
int max_col = width - box_w;
|
||||
if (max_row < 0 || max_col < 0) return {false, -1, -1};
|
||||
|
||||
int n_rows = max_row + 1;
|
||||
int n_cols = max_col + 1;
|
||||
int64_t total_positions = (int64_t)n_rows * n_cols;
|
||||
|
||||
std::mt19937 rng(seed);
|
||||
|
||||
// Fast path: if the entire canvas is empty, all positions are valid
|
||||
uint32_t total_occupied = data[(height - 1) * width + (width - 1)];
|
||||
if (total_occupied == 0) {
|
||||
int y = std::uniform_int_distribution<int>(0, max_row)(rng);
|
||||
int x = std::uniform_int_distribution<int>(0, max_col)(rng);
|
||||
return {true, y, x};
|
||||
}
|
||||
|
||||
// Quick random probe: try a few random positions first
|
||||
// If the canvas is mostly empty, one of these will hit quickly
|
||||
// This avoids the full O(H*W) scan for early words
|
||||
{
|
||||
int n_probes = std::min(16, (int)total_positions);
|
||||
for (int p = 0; p < n_probes; ++p) {
|
||||
int y = std::uniform_int_distribution<int>(0, max_row)(rng);
|
||||
int x = std::uniform_int_distribution<int>(0, max_col)(rng);
|
||||
if (get_area_sum_fast(y, x, box_h, box_w) == 0) {
|
||||
return {true, y, x};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Adaptive: use single thread for small canvases, multi for large
|
||||
unsigned int nt = 1;
|
||||
if (total_positions > 200000) {
|
||||
nt = get_num_threads();
|
||||
nt = std::min(nt, (unsigned int)n_rows);
|
||||
}
|
||||
|
||||
if (nt <= 1) {
|
||||
// Single-thread: two-pass (count then pick) is faster than
|
||||
// reservoir sampling because it avoids per-position RNG calls
|
||||
uint64_t total_valid = 0;
|
||||
for (int i = 0; i <= max_row; ++i) {
|
||||
for (int j = 0; j <= max_col; ++j) {
|
||||
if (get_area_sum_fast(i, j, box_h, box_w) == 0)
|
||||
++total_valid;
|
||||
}
|
||||
}
|
||||
if (total_valid == 0) return {false, -1, -1};
|
||||
uint64_t target = std::uniform_int_distribution<uint64_t>(0, total_valid - 1)(rng);
|
||||
uint64_t count = 0;
|
||||
for (int i = 0; i <= max_row; ++i) {
|
||||
for (int j = 0; j <= max_col; ++j) {
|
||||
if (get_area_sum_fast(i, j, box_h, box_w) == 0) {
|
||||
if (count == target) return {true, i, j};
|
||||
++count;
|
||||
}
|
||||
}
|
||||
}
|
||||
return {false, -1, -1};
|
||||
}
|
||||
|
||||
// Multi-thread path: parallel count then pick
|
||||
std::vector<uint64_t> chunk_counts(nt, 0);
|
||||
std::vector<int> chunk_starts(nt), chunk_ends(nt);
|
||||
int rows_per = (n_rows + nt - 1) / nt;
|
||||
|
||||
for (unsigned int t = 0; t < nt; ++t) {
|
||||
chunk_starts[t] = t * rows_per;
|
||||
chunk_ends[t] = std::min(n_rows, (int)(t + 1) * rows_per);
|
||||
}
|
||||
|
||||
auto& pool = get_pool();
|
||||
std::vector<std::future<uint64_t>> futs;
|
||||
futs.reserve(nt);
|
||||
|
||||
for (unsigned int t = 0; t < nt; ++t) {
|
||||
if (chunk_starts[t] >= chunk_ends[t]) break;
|
||||
futs.push_back(pool.submit([this, cs=chunk_starts[t], ce=chunk_ends[t], box_h, box_w]() {
|
||||
return this->count_valid_rows(cs, ce, box_h, box_w);
|
||||
}));
|
||||
}
|
||||
|
||||
uint64_t total = 0;
|
||||
for (size_t t = 0; t < futs.size(); ++t) {
|
||||
chunk_counts[t] = futs[t].get();
|
||||
total += chunk_counts[t];
|
||||
}
|
||||
|
||||
if (total == 0) return {false, -1, -1};
|
||||
|
||||
uint64_t target = std::uniform_int_distribution<uint64_t>(0, total - 1)(rng);
|
||||
|
||||
uint64_t cum = 0;
|
||||
for (size_t t = 0; t < futs.size(); ++t) {
|
||||
if (cum + chunk_counts[t] > target) {
|
||||
auto pos = pick_nth_valid(chunk_starts[t], chunk_ends[t], box_h, box_w, target - cum);
|
||||
return {true, pos.first, pos.second};
|
||||
}
|
||||
cum += chunk_counts[t];
|
||||
}
|
||||
return {false, -1, -1};
|
||||
}
|
||||
|
||||
// =========================================================
|
||||
// v3: Batch query — process multiple (box_h, box_w) in one call
|
||||
// Returns vector of {found, y, x} for each query
|
||||
// =========================================================
|
||||
|
||||
std::vector<DirectResult> batch_query(const std::vector<std::tuple<int,int,uint32_t>>& queries) {
|
||||
std::vector<DirectResult> results;
|
||||
results.reserve(queries.size());
|
||||
for (auto& [bh, bw, seed] : queries) {
|
||||
auto r = query_direct(bh, bw, seed);
|
||||
if (r.found) {
|
||||
update_rect_add(r.y, r.x, bh, bw);
|
||||
}
|
||||
results.push_back(r);
|
||||
}
|
||||
return results;
|
||||
}
|
||||
|
||||
// =========================================================
|
||||
// Legacy methods (kept for backward compatibility)
|
||||
// =========================================================
|
||||
|
||||
struct SearchResult {
|
||||
bool found;
|
||||
int y, x;
|
||||
size_t index;
|
||||
};
|
||||
|
||||
SearchResult search_range(size_t start_idx, size_t end_idx, int box_h, int box_w, int step) const {
|
||||
for (size_t i = start_idx; i < end_idx; i += step) {
|
||||
int y = valid_coords[i].first;
|
||||
int x = valid_coords[i].second;
|
||||
if (y + box_h <= height && x + box_w <= width) {
|
||||
if (get_area_sum(y, x, box_h, box_w) == 0) {
|
||||
return {true, y, x, i};
|
||||
}
|
||||
}
|
||||
}
|
||||
return {false, -1, -1, end_idx};
|
||||
}
|
||||
|
||||
std::pair<bool, std::pair<int, int>> find_spot_parallel(int box_h, int box_w, int step) {
|
||||
{
|
||||
std::shared_lock<std::shared_mutex> lock(mutex_);
|
||||
if (dirty_count > 0 && dirty_count >= rebuild_interval) {
|
||||
lock.unlock();
|
||||
std::unique_lock<std::shared_mutex> write_lock(mutex_);
|
||||
if (dirty_count > 0 && dirty_count >= rebuild_interval)
|
||||
rebuild_integral();
|
||||
}
|
||||
}
|
||||
|
||||
const size_t len = valid_coords.size();
|
||||
if (len == 0) return {false, {-1, -1}};
|
||||
|
||||
SearchResult best = {false, -1, -1, len};
|
||||
unsigned int nt = get_num_threads();
|
||||
nt = std::min(nt, (unsigned int)len);
|
||||
const size_t chunk = (len + nt - 1) / nt;
|
||||
|
||||
auto& pool = get_pool();
|
||||
std::vector<std::future<SearchResult>> futs;
|
||||
futs.reserve(nt);
|
||||
|
||||
Py_BEGIN_ALLOW_THREADS
|
||||
for (unsigned int i = 0; i < nt; ++i) {
|
||||
size_t s = i * chunk, e = std::min(len, s + chunk);
|
||||
if (s >= e) break;
|
||||
futs.push_back(pool.submit([this, s, e, box_h, box_w, step]() {
|
||||
return this->search_range(s, e, box_h, box_w, step);
|
||||
}));
|
||||
}
|
||||
|
||||
for (auto& f : futs) {
|
||||
auto res = f.get();
|
||||
if (res.found && res.index < best.index) best = res;
|
||||
}
|
||||
Py_END_ALLOW_THREADS
|
||||
|
||||
if (best.found) return {true, {best.y, best.x}};
|
||||
return {false, {-1, -1}};
|
||||
}
|
||||
};
|
||||
|
||||
// ==========================================
|
||||
// Python Bindings
|
||||
// ==========================================
|
||||
|
||||
typedef struct {
|
||||
PyObject_HEAD
|
||||
IntegralGrid* grid;
|
||||
} PyIntegralGrid;
|
||||
|
||||
static void Grid_dealloc(PyIntegralGrid* self) {
|
||||
if (self->grid) delete self->grid;
|
||||
Py_TYPE(self)->tp_free((PyObject*)self);
|
||||
}
|
||||
|
||||
static PyObject* Grid_new(PyTypeObject* type, PyObject* args, PyObject* kwds) {
|
||||
PyIntegralGrid* self = (PyIntegralGrid*)type->tp_alloc(type, 0);
|
||||
if (self) self->grid = NULL;
|
||||
return (PyObject*)self;
|
||||
}
|
||||
|
||||
static int Grid_init(PyIntegralGrid* self, PyObject* args, PyObject* kwds) {
|
||||
PyObject* mask_obj;
|
||||
int h, w;
|
||||
if (!PyArg_ParseTuple(args, "Oii", &mask_obj, &h, &w)) return -1;
|
||||
Py_buffer view;
|
||||
if (PyObject_GetBuffer(mask_obj, &view, PyBUF_SIMPLE) < 0) return -1;
|
||||
|
||||
self->grid = new IntegralGrid(h, w);
|
||||
self->grid->init_from_buffer((unsigned char*)view.buf, h, w);
|
||||
|
||||
PyBuffer_Release(&view);
|
||||
return 0;
|
||||
}
|
||||
|
||||
static PyObject* Grid_reorder_stratified(PyIntegralGrid* self, PyObject* args) {
|
||||
int bands = 0;
|
||||
if (!PyArg_ParseTuple(args, "i", &bands)) return NULL;
|
||||
self->grid->reorder_stratified(bands);
|
||||
Py_RETURN_NONE;
|
||||
}
|
||||
|
||||
static PyObject* Grid_query_sorted(PyIntegralGrid* self, PyObject* args) {
|
||||
int box_h, box_w, step = 1;
|
||||
if (!PyArg_ParseTuple(args, "ii|i", &box_h, &box_w, &step)) return NULL;
|
||||
auto result = self->grid->find_spot_parallel(box_h, box_w, step);
|
||||
if (result.first) return Py_BuildValue("ii", result.second.first, result.second.second);
|
||||
Py_RETURN_NONE;
|
||||
}
|
||||
|
||||
static PyObject* Grid_query_reservoir(PyIntegralGrid* self, PyObject* args) {
|
||||
int box_h, box_w;
|
||||
unsigned int seed = 0;
|
||||
if (!PyArg_ParseTuple(args, "ii|I", &box_h, &box_w, &seed)) return NULL;
|
||||
|
||||
// v3: delegate to query_direct for parallel reservoir
|
||||
auto r = self->grid->query_direct(box_h, box_w, seed);
|
||||
if (r.found) return Py_BuildValue("ii", r.y, r.x);
|
||||
Py_RETURN_NONE;
|
||||
}
|
||||
|
||||
// v3: Direct pixel-grid scan with parallel counting
|
||||
static PyObject* Grid_query_direct(PyIntegralGrid* self, PyObject* args) {
|
||||
int box_h, box_w;
|
||||
unsigned int seed = 0;
|
||||
if (!PyArg_ParseTuple(args, "ii|I", &box_h, &box_w, &seed)) return NULL;
|
||||
|
||||
Py_BEGIN_ALLOW_THREADS
|
||||
// query_direct is GIL-free safe (no Python objects touched)
|
||||
Py_END_ALLOW_THREADS
|
||||
|
||||
auto r = self->grid->query_direct(box_h, box_w, seed);
|
||||
if (r.found) return Py_BuildValue("ii", r.y, r.x);
|
||||
Py_RETURN_NONE;
|
||||
}
|
||||
|
||||
// v3: Batch query — process multiple placements in one C++ call
|
||||
// Input: list of (box_h, box_w, seed) tuples
|
||||
// Output: list of (y, x) or None for each
|
||||
static PyObject* Grid_batch_query(PyIntegralGrid* self, PyObject* args) {
|
||||
PyObject* query_list;
|
||||
if (!PyArg_ParseTuple(args, "O", &query_list)) return NULL;
|
||||
|
||||
PyObject* seq = PySequence_Fast(query_list, "expected a sequence");
|
||||
if (!seq) return NULL;
|
||||
|
||||
Py_ssize_t n = PySequence_Fast_GET_SIZE(seq);
|
||||
std::vector<std::tuple<int,int,uint32_t>> queries;
|
||||
queries.reserve(n);
|
||||
|
||||
for (Py_ssize_t i = 0; i < n; ++i) {
|
||||
PyObject* item = PySequence_Fast_GET_ITEM(seq, i);
|
||||
int bh, bw;
|
||||
unsigned int seed = 0;
|
||||
if (!PyArg_ParseTuple(item, "ii|I", &bh, &bw, &seed)) {
|
||||
Py_DECREF(seq);
|
||||
return NULL;
|
||||
}
|
||||
queries.push_back({bh, bw, seed});
|
||||
}
|
||||
Py_DECREF(seq);
|
||||
|
||||
auto results = self->grid->batch_query(queries);
|
||||
|
||||
PyObject* result_list = PyList_New(n);
|
||||
for (Py_ssize_t i = 0; i < n; ++i) {
|
||||
if (results[i].found) {
|
||||
PyList_SET_ITEM(result_list, i, Py_BuildValue("ii", results[i].y, results[i].x));
|
||||
} else {
|
||||
Py_INCREF(Py_None);
|
||||
PyList_SET_ITEM(result_list, i, Py_None);
|
||||
}
|
||||
}
|
||||
return result_list;
|
||||
}
|
||||
|
||||
static PyObject* Grid_update(PyIntegralGrid* self, PyObject* args) {
|
||||
int r, c, h, w;
|
||||
if (!PyArg_ParseTuple(args, "iiii", &r, &c, &h, &w)) return NULL;
|
||||
self->grid->update_rect_add(r, c, h, w);
|
||||
Py_RETURN_NONE;
|
||||
}
|
||||
|
||||
static PyObject* Grid_flush(PyIntegralGrid* self, PyObject* args) {
|
||||
self->grid->flush();
|
||||
Py_RETURN_NONE;
|
||||
}
|
||||
|
||||
// v4: Rebuild integral from full bitmap array
|
||||
static PyObject* Grid_rebuild_from_bitmap(PyIntegralGrid* self, PyObject* args) {
|
||||
PyObject* arr_obj;
|
||||
if (!PyArg_ParseTuple(args, "O", &arr_obj)) return NULL;
|
||||
Py_buffer view;
|
||||
if (PyObject_GetBuffer(arr_obj, &view, PyBUF_SIMPLE) < 0) return NULL;
|
||||
self->grid->rebuild_from_bitmap((const unsigned char*)view.buf);
|
||||
PyBuffer_Release(&view);
|
||||
Py_RETURN_NONE;
|
||||
}
|
||||
|
||||
// v4: Partial rebuild from position (pos_r, pos_c)
|
||||
static PyObject* Grid_rebuild_from_bitmap_partial(PyIntegralGrid* self, PyObject* args) {
|
||||
PyObject* arr_obj;
|
||||
int pos_r, pos_c;
|
||||
if (!PyArg_ParseTuple(args, "Oii", &arr_obj, &pos_r, &pos_c)) return NULL;
|
||||
Py_buffer view;
|
||||
if (PyObject_GetBuffer(arr_obj, &view, PyBUF_SIMPLE) < 0) return NULL;
|
||||
self->grid->rebuild_from_bitmap_partial((const unsigned char*)view.buf, pos_r, pos_c);
|
||||
PyBuffer_Release(&view);
|
||||
Py_RETURN_NONE;
|
||||
}
|
||||
|
||||
// v4: Stamp glyph bitmap onto C++ canvas and rebuild integral
|
||||
static PyObject* Grid_stamp_and_rebuild(PyIntegralGrid* self, PyObject* args) {
|
||||
PyObject* glyph_obj;
|
||||
int gh, gw, pos_r, pos_c;
|
||||
if (!PyArg_ParseTuple(args, "Oiiii", &glyph_obj, &gh, &gw, &pos_r, &pos_c)) return NULL;
|
||||
Py_buffer view;
|
||||
if (PyObject_GetBuffer(glyph_obj, &view, PyBUF_SIMPLE) < 0) return NULL;
|
||||
|
||||
self->grid->stamp_glyph((const unsigned char*)view.buf, gh, gw, pos_r, pos_c);
|
||||
self->grid->rebuild_from_canvas(pos_r, pos_c);
|
||||
|
||||
PyBuffer_Release(&view);
|
||||
Py_RETURN_NONE;
|
||||
}
|
||||
|
||||
static PyMethodDef Grid_methods[] = {
|
||||
{"reorder_stratified", (PyCFunction)Grid_reorder_stratified, METH_VARARGS, "Reorder valid coords with stratified interleaving."},
|
||||
{"query_sorted", (PyCFunction)Grid_query_sorted, METH_VARARGS, "Find position using sorted coordinate list (center-out, parallel)."},
|
||||
{"query_reservoir", (PyCFunction)Grid_query_reservoir, METH_VARARGS, "Find position using parallel direct-scan reservoir sampling."},
|
||||
{"query_direct", (PyCFunction)Grid_query_direct, METH_VARARGS, "Direct pixel-grid scan with parallel reservoir sampling."},
|
||||
{"batch_query", (PyCFunction)Grid_batch_query, METH_VARARGS, "Batch placement: list of (bh,bw,seed) -> list of (y,x)|None."},
|
||||
{"update", (PyCFunction)Grid_update, METH_VARARGS, "Update grid with placed rectangle."},
|
||||
{"flush", (PyCFunction)Grid_flush, METH_NOARGS, "Force rebuild integral image."},
|
||||
{"rebuild_from_bitmap", (PyCFunction)Grid_rebuild_from_bitmap, METH_VARARGS, "Rebuild integral from pixel bitmap array."},
|
||||
{"rebuild_from_bitmap_partial", (PyCFunction)Grid_rebuild_from_bitmap_partial, METH_VARARGS, "Partial integral rebuild from (pos_r, pos_c)."},
|
||||
{"stamp_and_rebuild", (PyCFunction)Grid_stamp_and_rebuild, METH_VARARGS, "Stamp glyph + partial integral rebuild (avoids PIL->numpy)."},
|
||||
{NULL}
|
||||
};
|
||||
|
||||
static PyTypeObject PyIntegralGridType = {
|
||||
PyVarObject_HEAD_INIT(NULL, 0)
|
||||
"ewc_core.IntegralGrid",
|
||||
sizeof(PyIntegralGrid),
|
||||
0,
|
||||
(destructor)Grid_dealloc,
|
||||
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
|
||||
Py_TPFLAGS_DEFAULT,
|
||||
"Integral Grid v3",
|
||||
0, 0, 0, 0, 0, 0,
|
||||
Grid_methods,
|
||||
0, 0, 0, 0, 0, 0, 0,
|
||||
(initproc)Grid_init,
|
||||
0,
|
||||
Grid_new,
|
||||
};
|
||||
|
||||
static PyModuleDef ewc_module = {
|
||||
PyModuleDef_HEAD_INIT, "ewc_core", "EfficientWordCloud Core v3", -1, NULL, NULL, NULL, NULL, NULL
|
||||
};
|
||||
|
||||
PyMODINIT_FUNC PyInit_ewc_core(void) {
|
||||
PyObject* m;
|
||||
if (PyType_Ready(&PyIntegralGridType) < 0) return NULL;
|
||||
m = PyModule_Create(&ewc_module);
|
||||
if (!m) return NULL;
|
||||
Py_INCREF(&PyIntegralGridType);
|
||||
if (PyModule_AddObject(m, "IntegralGrid", (PyObject *)&PyIntegralGridType) < 0) return NULL;
|
||||
return m;
|
||||
}
|
||||
@@ -0,0 +1,192 @@
|
||||
a
|
||||
about
|
||||
above
|
||||
after
|
||||
again
|
||||
against
|
||||
all
|
||||
also
|
||||
am
|
||||
an
|
||||
and
|
||||
any
|
||||
are
|
||||
aren't
|
||||
as
|
||||
at
|
||||
be
|
||||
because
|
||||
been
|
||||
before
|
||||
being
|
||||
below
|
||||
between
|
||||
both
|
||||
but
|
||||
by
|
||||
can
|
||||
can't
|
||||
cannot
|
||||
com
|
||||
could
|
||||
couldn't
|
||||
did
|
||||
didn't
|
||||
do
|
||||
does
|
||||
doesn't
|
||||
doing
|
||||
don't
|
||||
down
|
||||
during
|
||||
each
|
||||
else
|
||||
ever
|
||||
few
|
||||
for
|
||||
from
|
||||
further
|
||||
get
|
||||
had
|
||||
hadn't
|
||||
has
|
||||
hasn't
|
||||
have
|
||||
haven't
|
||||
having
|
||||
he
|
||||
he'd
|
||||
he'll
|
||||
he's
|
||||
hence
|
||||
her
|
||||
here
|
||||
here's
|
||||
hers
|
||||
herself
|
||||
him
|
||||
himself
|
||||
his
|
||||
how
|
||||
how's
|
||||
however
|
||||
http
|
||||
i
|
||||
i'd
|
||||
i'll
|
||||
i'm
|
||||
i've
|
||||
if
|
||||
in
|
||||
into
|
||||
is
|
||||
isn't
|
||||
it
|
||||
it's
|
||||
its
|
||||
itself
|
||||
just
|
||||
k
|
||||
let's
|
||||
like
|
||||
me
|
||||
more
|
||||
most
|
||||
mustn't
|
||||
my
|
||||
myself
|
||||
no
|
||||
nor
|
||||
not
|
||||
of
|
||||
off
|
||||
on
|
||||
once
|
||||
only
|
||||
or
|
||||
other
|
||||
otherwise
|
||||
ought
|
||||
our
|
||||
ours
|
||||
ourselves
|
||||
out
|
||||
over
|
||||
own
|
||||
r
|
||||
same
|
||||
shall
|
||||
shan't
|
||||
she
|
||||
she'd
|
||||
she'll
|
||||
she's
|
||||
should
|
||||
shouldn't
|
||||
since
|
||||
so
|
||||
some
|
||||
such
|
||||
than
|
||||
that
|
||||
that's
|
||||
the
|
||||
their
|
||||
theirs
|
||||
them
|
||||
themselves
|
||||
then
|
||||
there
|
||||
there's
|
||||
therefore
|
||||
these
|
||||
they
|
||||
they'd
|
||||
they'll
|
||||
they're
|
||||
they've
|
||||
this
|
||||
those
|
||||
through
|
||||
to
|
||||
too
|
||||
under
|
||||
until
|
||||
up
|
||||
very
|
||||
was
|
||||
wasn't
|
||||
we
|
||||
we'd
|
||||
we'll
|
||||
we're
|
||||
we've
|
||||
were
|
||||
weren't
|
||||
what
|
||||
what's
|
||||
when
|
||||
when's
|
||||
where
|
||||
where's
|
||||
which
|
||||
while
|
||||
who
|
||||
who's
|
||||
whom
|
||||
why
|
||||
why's
|
||||
with
|
||||
won't
|
||||
would
|
||||
wouldn't
|
||||
www
|
||||
you
|
||||
you'd
|
||||
you'll
|
||||
you're
|
||||
you've
|
||||
your
|
||||
yours
|
||||
yourself
|
||||
yourselves
|
||||
@@ -0,0 +1,109 @@
|
||||
# Ported from ref/word_cloud/wordcloud/tokenization.py (MIT License)
|
||||
from __future__ import division
|
||||
from itertools import tee
|
||||
from operator import itemgetter
|
||||
from collections import defaultdict
|
||||
from math import log
|
||||
|
||||
|
||||
def _l(k, n, x):
|
||||
"""Dunning log-likelihood helper."""
|
||||
return log(max(x, 1e-10)) * k + log(max(1 - x, 1e-10)) * (n - k)
|
||||
|
||||
|
||||
def _collocation_score(count_bigram, count1, count2, n_words):
|
||||
"""Dunning likelihood ratio collocation score."""
|
||||
if n_words <= count1 or n_words <= count2:
|
||||
return 0
|
||||
N, c12, c1, c2 = n_words, count_bigram, count1, count2
|
||||
p = c2 / N
|
||||
p1 = c12 / c1
|
||||
p2 = (c2 - c12) / (N - c1)
|
||||
score = (
|
||||
_l(c12, c1, p) + _l(c2 - c12, N - c1, p)
|
||||
- _l(c12, c1, p1) - _l(c2 - c12, N - c1, p2)
|
||||
)
|
||||
return -2 * score
|
||||
|
||||
|
||||
def _pairwise(iterable):
|
||||
a, b = tee(iterable)
|
||||
next(b, None)
|
||||
return zip(a, b)
|
||||
|
||||
|
||||
def process_tokens(words, normalize_plurals=True):
|
||||
"""Count words, normalizing case and optionally merging plurals.
|
||||
|
||||
Returns
|
||||
-------
|
||||
counts : dict str -> int
|
||||
standard_forms : dict lowercase_str -> canonical_str
|
||||
"""
|
||||
d = defaultdict(dict)
|
||||
for word in words:
|
||||
wl = word.lower()
|
||||
case_dict = d[wl]
|
||||
case_dict[word] = case_dict.get(word, 0) + 1
|
||||
|
||||
if normalize_plurals:
|
||||
merged_plurals = {}
|
||||
for key in list(d.keys()):
|
||||
if key.endswith('s') and not key.endswith('ss'):
|
||||
singular = key[:-1]
|
||||
if singular in d:
|
||||
for word, count in d[key].items():
|
||||
sing_form = word[:-1]
|
||||
d[singular][sing_form] = d[singular].get(sing_form, 0) + count
|
||||
merged_plurals[key] = singular
|
||||
del d[key]
|
||||
|
||||
fused_cases = {}
|
||||
standard_cases = {}
|
||||
item1 = itemgetter(1)
|
||||
for word_lower, case_dict in d.items():
|
||||
first = max(case_dict.items(), key=item1)[0]
|
||||
fused_cases[first] = sum(case_dict.values())
|
||||
standard_cases[word_lower] = first
|
||||
|
||||
if normalize_plurals:
|
||||
for plural, singular in merged_plurals.items():
|
||||
standard_cases[plural] = standard_cases.get(singular, singular)
|
||||
|
||||
return fused_cases, standard_cases
|
||||
|
||||
|
||||
def unigrams_and_bigrams(words, stopwords, normalize_plurals=True,
|
||||
collocation_threshold=30):
|
||||
"""Return word counts including statistically significant bigrams."""
|
||||
bigrams = [
|
||||
p for p in _pairwise(words)
|
||||
if not any(w.lower() in stopwords for w in p)
|
||||
]
|
||||
unigrams = [w for w in words if w.lower() not in stopwords]
|
||||
n_words = len(unigrams)
|
||||
|
||||
counts_unigrams, standard_form = process_tokens(
|
||||
unigrams, normalize_plurals=normalize_plurals)
|
||||
counts_bigrams, _ = process_tokens(
|
||||
[" ".join(b) for b in bigrams], normalize_plurals=normalize_plurals)
|
||||
|
||||
orig_counts = counts_unigrams.copy()
|
||||
|
||||
for bigram_string, count in counts_bigrams.items():
|
||||
parts = bigram_string.split(" ", 1)
|
||||
if len(parts) != 2:
|
||||
continue
|
||||
word1 = standard_form.get(parts[0].lower(), parts[0])
|
||||
word2 = standard_form.get(parts[1].lower(), parts[1])
|
||||
if word1 not in orig_counts or word2 not in orig_counts:
|
||||
continue
|
||||
score = _collocation_score(count, orig_counts[word1],
|
||||
orig_counts[word2], n_words)
|
||||
if score > collocation_threshold:
|
||||
counts_unigrams[word1] -= count
|
||||
counts_unigrams[word2] -= count
|
||||
counts_unigrams[bigram_string] = count
|
||||
|
||||
# Remove non-positive counts
|
||||
return {w: c for w, c in counts_unigrams.items() if c > 0}
|
||||
@@ -0,0 +1,630 @@
|
||||
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'<svg width="{out_w}" height="{out_h}" xmlns="http://www.w3.org/2000/svg">',
|
||||
f'<style>text{{font-family:{font_family};font-weight:{font_weight};'
|
||||
f'font-style:{font_style};}}</style>',
|
||||
]
|
||||
if self.background_color is not None:
|
||||
lines.append(
|
||||
f'<rect width="100%" height="100%" style="fill:{self.background_color}"/>'
|
||||
)
|
||||
|
||||
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'<text transform="{transform}" font-size="{scaled_size}" '
|
||||
f'style="fill:{color}">{saxutils.escape(word)}</text>'
|
||||
)
|
||||
|
||||
lines.append("</svg>")
|
||||
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
|
||||
Reference in New Issue
Block a user