yu-zhi-ran/scripts/cover_generator.py

#!/usr/bin/env python3
"""
封面图生成器：为三平台文章生成文字封面图

依赖: Pillow + Noto Sans CJK 字体
用法: python3 scripts/cover_generator.py --topic-id A07
"""

import argparse, datetime, json, logging, sys, os
from pathlib import Path
from typing import Dict, Optional
from PIL import Image, ImageDraw, ImageFont

PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(PROJECT_ROOT / "platform" / "backend"))

from db_helper import get_topic_by_id

LOGS_DIR = PROJECT_ROOT / "automation" / "logs"
OUTPUT_DIR = PROJECT_ROOT / "automation" / "images" / "generated"
TODAY = datetime.datetime.now().strftime("%Y-%m-%d")

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s',
                    handlers=[logging.FileHandler(LOGS_DIR / f"cover_{TODAY}.log"), logging.StreamHandler()])
logger = logging.getLogger(__name__)

# Find a suitable CJK font
_FONT_CANDIDATES = [
    "/usr/share/fonts/google-noto-cjk/NotoSansCJK-Bold.ttc",
    "/usr/share/fonts/google-noto-cjk/NotoSansCJK-Medium.ttc",
    "/usr/share/fonts/google-noto-cjk/NotoSansCJK-Regular.ttc",
    "/usr/share/fonts/opentype/noto/NotoSansCJK-Bold.ttc",
    "/usr/share/fonts/truetype/noto/NotoSansCJK-Bold.ttc",
]
_FONT_PATH = None
for fp in _FONT_CANDIDATES:
    if Path(fp).exists():
        _FONT_PATH = fp
        break

if not _FONT_PATH:
    import subprocess
    r = subprocess.run(["fc-match", "-f", "%{file}", "sans:lang=zh"], capture_output=True, text=True)
    if r.returncode == 0 and Path(r.stdout.strip()).exists():
        _FONT_PATH = r.stdout.strip()

COVER_SIZES = {
    "zhihu": (1200, 630),
    "wechat": (900, 500),
    "xiaohongshu": (1080, 1440),
}

PLATFORM_COLORS = {
    "zhihu": { "bg_start": "#0066cc", "bg_end": "#004d99", "title": "#ffffff", "accent": "#66b3ff", "watermark": "知乎专栏" },
    "wechat": { "bg_start": "#07c160", "bg_end": "#06ad56", "title": "#ffffff", "accent": "#b3ffd9", "watermark": "微信公众号" },
    "xiaohongshu": { "bg_start": "#ff2442", "bg_end": "#cc1a35", "title": "#ffffff", "accent": "#ffb3c1", "watermark": "小红书" },
}


def _get_font(size: int):
    if _FONT_PATH:
        return ImageFont.truetype(_FONT_PATH, size)
    return ImageFont.load_default()


def _wrap_text(text: str, max_chars: int = 6) -> str:
    """Split long title into lines"""
    lines = []
    while text:
        if len(text) <= max_chars:
            lines.append(text)
            break
        idx = text.rfind("，", 0, max_chars + 1)
        if idx <= 0:
            idx = text.rfind(" ", 0, max_chars + 1)
        if idx <= 0:
            idx = max_chars
        lines.append(text[:idx])
        text = text[idx:]
    return "\n".join(lines)


def generate_cover(topic_id: str, platform: str) -> Optional[str]:
    """Generate cover image for a topic on a specific platform"""
    topic = get_topic_by_id(topic_id)
    if not topic:
        logger.error(f"Topic {topic_id} not found")
        return None

    title = topic.get("title", "")
    if not title:
        logger.warning(f"Topic {topic_id} has no title")
        return None

    size = COVER_SIZES.get(platform)
    if not size:
        logger.error(f"Unknown platform: {platform}")
        return None

    colors = PLATFORM_COLORS[platform]
    w, h = size
    img = Image.new("RGB", (w, h))
    draw = ImageDraw.Draw(img)

    for y in range(h):
        r = y / h
        for c in range(3):
            start = int(colors["bg_start"][1:3], 16) if c == 0 else (int(colors["bg_start"][3:5], 16) if c == 1 else int(colors["bg_start"][5:7], 16))
            end = int(colors["bg_end"][1:3], 16) if c == 0 else (int(colors["bg_end"][3:5], 16) if c == 1 else int(colors["bg_end"][5:7], 16))
            val = int(start + (end - start) * r)
            color = list(img.getpixel((0, y)) if y > 0 else (0, 0, 0))
            color[c] = val
            img.putpixel((0, y), tuple(color))
    for x in range(1, w):
        color = img.getpixel((0, y))
        draw.line([(x, y), (x, y)], fill=color)

    if platform == "xiaohongshu":
        title_font_size = 56
        wrapped = _wrap_text(title, max_chars=7)
        font = _get_font(title_font_size)
        lines = wrapped.split("\n")
        total_h = len(lines) * (title_font_size + 20)
        start_y = (h - total_h) // 2 - 40
        for i, line in enumerate(lines):
            bbox = draw.textbbox((0, 0), line, font=font)
            tw = bbox[2] - bbox[0]
            x = (w - tw) // 2
            y_pos = start_y + i * (title_font_size + 20)
            draw.text((x, y_pos), line, fill=colors["title"], font=font)

        font_small = _get_font(24)
        wm = colors["watermark"]
        bbox = draw.textbbox((0, 0), wm, font=font_small)
        ww = bbox[2] - bbox[0]
        draw.text(((w - ww) // 2, h - 80), wm, fill=colors["accent"], font=font_small)
    else:
        title_font_size = 52 if platform == "wechat" else 64
        font = _get_font(title_font_size)
        bbox = draw.textbbox((0, 0), title, font=font)
        tw = bbox[2] - bbox[0]
        if tw > w - 120:
            wrapped = _wrap_text(title, max_chars=8)
            lines = wrapped.split("\n")
            total_h = len(lines) * (title_font_size + 16)
            start_y = (h - total_h) // 2
            for i, line in enumerate(lines):
                bbox2 = draw.textbbox((0, 0), line, font=font)
                tw2 = bbox2[2] - bbox2[0]
                x = (w - tw2) // 2
                draw.text((x, start_y + i * (title_font_size + 16)), line, fill=colors["title"], font=font)
        else:
            x = (w - tw) // 2
            draw.text((x, h // 2 - title_font_size), title, fill=colors["title"], font=font)

        font_small = _get_font(20)
        wm = colors["watermark"]
        bbox = draw.textbbox((0, 0), wm, font=font_small)
        ww = bbox[2] - bbox[0]
        draw.text((w - ww - 20, h - 40), wm, fill=colors["accent"], font=font_small)

    out_dir = OUTPUT_DIR / TODAY
    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / f"cover_{platform}_{topic_id}.png"
    img.save(out_path, "PNG")
    relative_path = f"automation/images/generated/{TODAY}/cover_{platform}_{topic_id}.png"
    logger.info(f"Cover saved: {out_path}")
    return relative_path


def update_article_image(topic_id: str, platform: str, image_path: str):
    """Save cover path to Article.images JSON field"""
    try:
        from db_helper import save_cover_to_article
        save_cover_to_article(topic_id, platform, image_path)
    except Exception as e:
        logger.warning(f"Failed to update article image in DB: {e}")


def main():
    parser = argparse.ArgumentParser(description="生成文章封面图")
    parser.add_argument("--topic-id", required=True)
    parser.add_argument("--platform", choices=["zhihu", "wechat", "xiaohongshu", "all"], default="all")
    args = parser.parse_args()

    platforms = ["zhihu", "wechat", "xiaohongshu"] if args.platform == "all" else [args.platform]
    images = {}
    for plat in platforms:
        path = generate_cover(args.topic_id, plat)
        if path:
            images[plat] = path
            update_article_image(args.topic_id, plat, path)
    print(json.dumps(images, ensure_ascii=False))
    logger.info(f"Generated covers for {args.topic_id}: {list(images.keys())}")


if __name__ == "__main__":
    main()