wordcloud/backend/core/mask.py

import os
from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw

from . import config
from .fonts import get_cached_font


def analyze_mask(mask):
    free = mask == 0
    free_area = int(np.sum(free))
    total_area = int(mask.size)
    free_ratio = (free_area / total_area) if total_area else 0.0

    rows = np.where(np.any(free, axis=1))[0]
    cols = np.where(np.any(free, axis=0))[0]
    bbox_fill_ratio = free_ratio
    bbox = None
    if rows.size and cols.size:
        y0, y1 = int(rows[0]), int(rows[-1])
        x0, x1 = int(cols[0]), int(cols[-1])
        bbox = (x0, y0, x1, y1)
        bbox_area = max(1, (x1 - x0 + 1) * (y1 - y0 + 1))
        bbox_fill_ratio = free_area / bbox_area

    return {
        "free_area": free_area,
        "free_ratio": free_ratio,
        "bbox": bbox,
        "bbox_fill_ratio": bbox_fill_ratio,
    }


def normalize_mask_for_fill(mask):
    if config.FILL_ON == "WHITE":
        return np.where(mask > 128, 0, 255).astype(np.uint8)
    return np.where(mask > 128, 255, 0).astype(np.uint8)


def calculate_dynamic_dimensions(base_w, base_h, num_words, avg_len=3, mask_stats=None):
    effective_fill = config.TARGET_FILL_RATIO if config.TARGET_FILL_RATIO > 0 else 0.45
    mask_fill_ratio = 0.5
    if mask_stats is not None:
        mask_fill_ratio = max(0.05, mask_stats["free_ratio"])

    area_per_word = (config.MIN_READABLE_HEIGHT_PX ** 2) * max(1.0, avg_len) * 1.05
    required_fillable_area = (num_words * area_per_word * max(1, config.N_REPETITIONS)) / max(effective_fill, 0.1)
    required_canvas_area = required_fillable_area / mask_fill_ratio
    current_area = base_w * base_h
    if required_canvas_area > current_area:
        scale_factor = (required_canvas_area / current_area) ** 0.5
        new_w = int(base_w * scale_factor)
        new_h = int(base_h * scale_factor)
        # Cap HD canvas to keep render/export under the speed budget.
        # 6000px 边长对激光/打印足够，再大收益很小但 SVG/PNG 成本陡增。
        max_edge = 6000
        if max(new_w, new_h) > max_edge:
            s = max_edge / max(new_w, new_h)
            new_w = int(new_w * s)
            new_h = int(new_h * s)
        new_w = max(100, ((new_w // 100) + 1) * 100)
        new_h = max(100, ((new_h // 100) + 1) * 100)
        print(f"[Auto-Size] 扩展画布: {base_w}x{base_h} -> {new_w}x{new_h}")
        return new_w, new_h
    return base_w, base_h


def prepare_mask(target_w, target_h):
    if config.MODE == "TEXT":
        img_mask_gen = Image.new("L", (target_w, target_h), 255)
        draw_mask = ImageDraw.Draw(img_mask_gen)
        font_size = min(config.MASK_FONT_SIZE, int(target_h * 0.75))
        font_mask = get_cached_font(config.MASK_FONT_PATH, font_size)
        bbox = draw_mask.textbbox((0, 0), config.MASK_TEXT, font=font_mask)
        text_w = bbox[2] - bbox[0]
        text_h = bbox[3] - bbox[1]
        x_pos = (target_w - text_w) // 2
        y_pos = (target_h - text_h) // 2
        draw_mask.text((x_pos, y_pos), config.MASK_TEXT, fill=0, font=font_mask)
        mask_hd = np.array(img_mask_gen)
        mask_hd = normalize_mask_for_fill(mask_hd)
        return mask_hd, (target_w, target_h), None
    if config.MODE == "IMAGE":
        if not os.path.exists(config.MASK_IMAGE_PATH):
            raise FileNotFoundError(f"找不到掩膜文件 {config.MASK_IMAGE_PATH}")
        img_raw = Image.open(config.MASK_IMAGE_PATH)
        if img_raw.mode in ('RGBA', 'LA') or (img_raw.mode == 'P' and 'transparency' in img_raw.info):
            img_bg = Image.new('RGB', img_raw.size, (255, 255, 255))
            if img_raw.mode == 'P':
                img_raw = img_raw.convert('RGBA')
            img_bg.paste(img_raw, mask=img_raw.split()[-1])
            img_src = img_bg.convert('L')
        else:
            img_src = img_raw.convert("L")
        src_w, src_h = img_src.size
        if config.IMAGE_CANVAS_MODE == "AUTO" or (target_w is None and target_h is None):
            final_w, final_h = src_w, src_h
        elif config.IMAGE_CANVAS_MODE == "WIDTH":
            final_w = target_w
            final_h = int(round(final_w * src_h / src_w))
        elif config.IMAGE_CANVAS_MODE == "HEIGHT":
            final_h = target_h
            final_w = int(round(final_h * src_w / src_h))
        else:
            final_w, final_h = target_w, target_h
        if (final_w, final_h) != (src_w, src_h):
            print(f"正在重采样掩膜: {src_w}x{src_h} -> {final_w}x{final_h} (LANCZOS)")
            img_src = img_src.resize((final_w, final_h), Image.Resampling.LANCZOS)
        threshold = 200
        img_src = img_src.point(lambda p: 255 if p > threshold else 0)

        # 自动填充边角区域为可填充（黑色）
        if config.FILL_CORNERS:
            arr = np.array(img_src)
            corner_h = int(final_h * config.CORNER_FILL_RATIO)
            corner_w = int(final_w * config.CORNER_FILL_RATIO)
            # 四个角落区域设为黑色（可填充）
            arr[:corner_h, :corner_w] = 0  # 左上
            arr[:corner_h, -corner_w:] = 0  # 右上
            arr[-corner_h:, :corner_w] = 0  # 左下
            arr[-corner_h:, -corner_w:] = 0  # 右下
            img_src = Image.fromarray(arr)
            print(f"[边角填充] 四角区域 {corner_w}x{corner_h} 已设为可填充")

        if config.SAVE_DEBUG_IMAGES:
            debug_dir = config.DEBUG_OUTPUT_DIR
            os.makedirs(debug_dir, exist_ok=True)
            img_src.save(str(Path(debug_dir) / "mask_src.png"))
        mask_hd = np.array(img_src)
        mask_hd = normalize_mask_for_fill(mask_hd)

        return mask_hd, (final_w, final_h), None

    raise ValueError(f"未知 MODE: {config.MODE}")


def apply_safe_padding(mask, padding_px=4, padding_ratio=0.003, max_padding=20):
    h, w = mask.shape
    padding = max(padding_px, int(min(h, w) * padding_ratio))
    padding = min(padding, max_padding)
    if padding <= 0:
        return mask
    mask[:padding, :] = 255
    mask[-padding:, :] = 255
    mask[:, :padding] = 255
    mask[:, -padding:] = 255
    return mask