Agent/examples/knowledge_extract/ai-portrait-realism/scripts/group/collage.py

"""
把 result.json 渲染成拼图: 每张最多 12 格, 按总数均匀分配。
每格 = 完整(contain, 不裁剪)封面 + 序号徽章 + 标题。
输出到 collages/page-NN.jpg + INDEX.md。
"""
import json
import math
from pathlib import Path

from PIL import Image, ImageDraw, ImageFont

OUT_DIR = Path('/Users/sunlit/Profile/analysis/ai-portrait-realism')
RESULT = OUT_DIR / 'result.json'
COLLAGE_DIR = OUT_DIR / 'collages'
COLLAGE_DIR.mkdir(exist_ok=True)

# ── 视觉参数 ────────────────────────
MAX_PER_PAGE = 12
COLS = 4                    # 4 列网格 (4x3 = 12 max)
TILE_IMG_SIZE = 380         # 每格图片区边长 (正方形 contain 容器)
TILE_TEXT_HEIGHT = 96       # 每格文字区高度
PADDING = 16                # 格间距
PAGE_HEADER_H = 64          # 页眉高度

BG_COLOR = (245, 245, 247)
TEXT_CARD_BG = (255, 255, 255)
TEXT_COLOR = (29, 29, 31)
SUBTLE_COLOR = (134, 134, 139)
INDEX_BG = (220, 60, 60)
INDEX_FG = (255, 255, 255)
LETTERBOX_BG = (29, 29, 31)  # 与前端 contain 背景一致

CATEGORY_COLOR = {
    'realism': (52, 199, 89),
    'studio': (0, 113, 227),
    'portrait_art': (175, 82, 222),
}

FONT_CANDIDATES = [
    '/System/Library/Fonts/Hiragino Sans GB.ttc',
    '/System/Library/Fonts/STHeiti Medium.ttc',
    '/System/Library/Fonts/Supplemental/Arial Unicode.ttf',
]


def load_font(size):
    for path in FONT_CANDIDATES:
        try:
            return ImageFont.truetype(path, size)
        except (OSError, IOError):
            continue
    return ImageFont.load_default()


def split_evenly(total, max_per):
    """27 → [9,9,9]; 24 → [12,12]; 13 → [7,6]"""
    n_pages = math.ceil(total / max_per) or 1
    base = total // n_pages
    extra = total % n_pages
    return [base + 1] * extra + [base] * (n_pages - extra)


def fit_image(src_path, target_w, target_h):
    """contain 缩放, 居中, letterbox 填充 — 完整显示, 不裁剪。"""
    try:
        img = Image.open(src_path).convert('RGB')
    except Exception:
        return None
    w, h = img.size
    scale = min(target_w / w, target_h / h)
    new_w = max(1, int(w * scale))
    new_h = max(1, int(h * scale))
    img = img.resize((new_w, new_h), Image.LANCZOS)
    canvas = Image.new('RGB', (target_w, target_h), LETTERBOX_BG)
    canvas.paste(img, ((target_w - new_w) // 2, (target_h - new_h) // 2))
    return canvas


def wrap_chinese(draw, text, font, max_width):
    """字符级换行, 中英文均适用。"""
    lines, cur = [], ''
    for ch in text:
        test = cur + ch
        bbox = draw.textbbox((0, 0), test, font=font)
        if (bbox[2] - bbox[0]) > max_width and cur:
            lines.append(cur)
            cur = ch
        else:
            cur = test
    if cur:
        lines.append(cur)
    return lines


def make_tile(item):
    img_w = TILE_IMG_SIZE
    text_h = TILE_TEXT_HEIGHT
    tile = Image.new('RGB', (img_w, img_w + text_h), BG_COLOR)

    # 上半: 封面 (contain, 完整)
    cover_path = OUT_DIR / item['cover'] if item.get('cover') else None
    img_canvas = fit_image(cover_path, img_w, img_w) if cover_path and cover_path.exists() else None
    if img_canvas is None:
        img_canvas = Image.new('RGB', (img_w, img_w), (200, 200, 205))
        d = ImageDraw.Draw(img_canvas)
        f = load_font(20)
        msg = '无封面'
        bbox = d.textbbox((0, 0), msg, font=f)
        bw, bh = bbox[2] - bbox[0], bbox[3] - bbox[1]
        d.text(((img_w - bw) // 2, (img_w - bh) // 2), msg, fill=(110, 110, 115), font=f)
    tile.paste(img_canvas, (0, 0))

    # 下半: 文字卡 (序号徽章 + 标题 + 分类点)
    text_canvas = Image.new('RGB', (img_w, text_h), TEXT_CARD_BG)
    draw = ImageDraw.Draw(text_canvas)

    # 序号徽章 (左上)
    pad = 12
    badge_h = 30
    badge_w = 56
    draw.rounded_rectangle((pad, pad, pad + badge_w, pad + badge_h),
                           radius=8, fill=INDEX_BG)
    badge_font = load_font(17)
    badge_text = f'#{item["index"]:02d}'
    bbox = draw.textbbox((0, 0), badge_text, font=badge_font)
    bw = bbox[2] - bbox[0]
    bh = bbox[3] - bbox[1]
    draw.text((pad + (badge_w - bw) // 2, pad + (badge_h - bh) // 2 - 2),
              badge_text, fill=INDEX_FG, font=badge_font)

    # 分类小圆点 (右上)
    cat_color = CATEGORY_COLOR.get(item.get('category', ''), (180, 180, 180))
    dot_r = 6
    dot_x = img_w - pad - dot_r * 2
    dot_y = pad + badge_h // 2 - dot_r
    draw.ellipse((dot_x, dot_y, dot_x + dot_r * 2, dot_y + dot_r * 2), fill=cat_color)

    # 标题 (徽章下方, 最多 2 行)
    title_font = load_font(15)
    title_x = pad
    title_y = pad + badge_h + 8
    title_max_w = img_w - 2 * pad
    title_lines = wrap_chinese(draw, item.get('title', ''), title_font, title_max_w)
    if len(title_lines) > 2:
        # 截到 2 行 + 省略号
        last = title_lines[1]
        while last:
            test = last[:-1] + '…'
            bbox = draw.textbbox((0, 0), test, font=title_font)
            if (bbox[2] - bbox[0]) <= title_max_w:
                title_lines = [title_lines[0], test]
                break
            last = last[:-1]
    for i, line in enumerate(title_lines[:2]):
        draw.text((title_x, title_y + i * 22), line, fill=TEXT_COLOR, font=title_font)

    tile.paste(text_canvas, (0, img_w))
    return tile


def make_page(items, page_num, total_pages):
    cols = min(COLS, len(items))
    rows = math.ceil(len(items) / cols)
    tile_w = TILE_IMG_SIZE
    tile_h = TILE_IMG_SIZE + TILE_TEXT_HEIGHT
    page_w = cols * tile_w + (cols + 1) * PADDING
    page_h = PAGE_HEADER_H + rows * tile_h + (rows + 1) * PADDING

    page = Image.new('RGB', (page_w, page_h), BG_COLOR)
    draw = ImageDraw.Draw(page)
    header_font = load_font(22)
    sub_font = load_font(13)

    title = f'AI 写真调研 — 第 {page_num} / {total_pages} 页'
    draw.text((PADDING * 2, 18), title, fill=TEXT_COLOR, font=header_font)
    sub = f'{len(items)} 条 (#{items[0]["index"]:02d} – #{items[-1]["index"]:02d})'
    draw.text((PADDING * 2, 44), sub, fill=SUBTLE_COLOR, font=sub_font)

    for i, item in enumerate(items):
        r, c = divmod(i, cols)
        x = PADDING + c * (tile_w + PADDING)
        y = PAGE_HEADER_H + PADDING + r * (tile_h + PADDING)
        tile = make_tile(item)
        page.paste(tile, (x, y))
    return page


def main():
    items = json.loads(RESULT.read_text(encoding='utf-8'))
    sizes = split_evenly(len(items), MAX_PER_PAGE)
    print(f'Total {len(items)} items → {len(sizes)} pages ({sizes})')

    pages_meta = []
    cursor = 0
    for page_num, n in enumerate(sizes, 1):
        page_items = items[cursor:cursor + n]
        cursor += n
        page_img = make_page(page_items, page_num, len(sizes))
        out_path = COLLAGE_DIR / f'page-{page_num:02d}.jpg'
        page_img.save(out_path, 'JPEG', quality=88, optimize=True)
        pages_meta.append({
            'page': page_num,
            'file': out_path.name,
            'size': page_img.size,
            'indices': [it['index'] for it in page_items],
        })
        print(f'  [OK] {out_path.name}  {page_img.size[0]}×{page_img.size[1]}px  ({len(page_items)} items)')

    # 索引
    md = ['# 拼图索引', '', f'总 {len(items)} 条 / {len(sizes)} 张拼图 / 每张最多 {MAX_PER_PAGE} 格', '']
    for m in pages_meta:
        md.append(f'## {m["file"]}')
        md.append(f'- 尺寸: {m["size"][0]}×{m["size"][1]}px')
        md.append(f'- 条目: {", ".join(f"#{i:02d}" for i in m["indices"])}')
        md.append('')
    (COLLAGE_DIR / 'INDEX.md').write_text('\n'.join(md), encoding='utf-8')

    print(f'\n=== {len(sizes)} pages → {COLLAGE_DIR}')


if __name__ == '__main__':
    main()