146 lines
5.7 KiB
Python
146 lines
5.7 KiB
Python
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):
|
|
if not config.AUTO_EXPAND_CANVAS:
|
|
return base_w, base_h
|
|
|
|
effective_fill = config.TARGET_FILL_RATIO if config.TARGET_FILL_RATIO > 0 else max(config.MIN_ACCEPT_FILL_RATIO, 0.82)
|
|
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.2
|
|
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)
|
|
new_w = ((new_w // 100) + 1) * 100
|
|
new_h = ((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
|