persona/scripts/portrait.py

#!/usr/bin/env python3
"""從參考照片裁出人物臉部，輸出成人格圖示用的正方形 PNG。

這支腳本是**選用的加值工具**：jsc-persona 本體只用 Node 內建模組，沒有它照樣能產生
向量人物形象。裝了 Pillow（＋可選的 OpenCV 動漫臉偵測）之後，圖示就能改用真實照片裁臉。

用法：
    python3 portrait.py --input <圖片> --list                     # 列出偵測到的所有臉
    python3 portrait.py --input <圖片> --output <out.png> [--size 512]
                        [--cascade <xml>] [--pick largest|leftmost|rightmost|<index>]
                        [--face x,y,w,h]                          # 直接指定裁切框

多角色的圖片一定要挑臉：`--list` 看有哪些，再用 `--pick` 或 `--face` 指定。
輸出（stdout）：一行 JSON。失敗時 ok=false，並附上 hint。
"""

import argparse
import json
import os
import sys


def emit(payload):
    print(json.dumps(payload, ensure_ascii=False))
    sys.exit(0)


def fail(reason, hint=None):
    emit({"ok": False, "reason": reason, "hint": hint})


def detect_faces(path, cascade_path):
    """回傳 (faces, method)；faces 是 [(x, y, w, h), ...]，偵測不到就回 ([], None)。"""
    try:
        import cv2
    except ImportError:
        return [], None

    candidates = []
    if cascade_path and os.path.exists(cascade_path):
        candidates.append((cascade_path, "anime-cascade"))
    builtin = getattr(getattr(cv2, "data", None), "haarcascades", "")
    if builtin:
        frontal = os.path.join(builtin, "haarcascade_frontalface_default.xml")
        if os.path.exists(frontal):
            candidates.append((frontal, "frontal-cascade"))

    try:
        image = cv2.imread(path)
        if image is None:
            return [], None
        gray = cv2.equalizeHist(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY))
    except Exception:
        return [], None

    # 由嚴到寬試幾組參數：先求準，找不到才放寬
    ladder = [(1.05, 5, 40), (1.05, 3, 32), (1.02, 2, 24)]
    for xml, method in candidates:
        try:
            clf = cv2.CascadeClassifier(xml)
            if clf.empty():
                continue
        except Exception:
            continue
        for sf, mn, ms in ladder:
            try:
                faces = clf.detectMultiScale(gray, scaleFactor=sf, minNeighbors=mn, minSize=(ms, ms))
            except Exception:
                continue
            if len(faces):
                return [tuple(int(v) for v in f) for f in faces], method
    return [], None


def choose(faces, pick):
    if not faces:
        return None
    if pick is None or pick == "largest":
        return max(faces, key=lambda f: f[2] * f[3])
    if pick == "leftmost":
        return min(faces, key=lambda f: f[0])
    if pick == "rightmost":
        return max(faces, key=lambda f: f[0] + f[2])
    if pick == "topmost":
        return min(faces, key=lambda f: f[1])
    try:
        return sorted(faces, key=lambda f: f[0])[int(pick)]
    except (ValueError, IndexError):
        return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", required=True)
    ap.add_argument("--output")
    ap.add_argument("--size", type=int, default=512)
    ap.add_argument("--cascade", default=None)
    ap.add_argument("--radius", type=float, default=0.22, help="圓角半徑（佔邊長比例）")
    ap.add_argument("--pick", default=None, help="largest|leftmost|rightmost|topmost|<由左至右的索引>")
    ap.add_argument("--face", default=None, help="直接指定臉的框 x,y,w,h")
    ap.add_argument("--list", action="store_true", help="只列出偵測到的臉，不輸出圖")
    args = ap.parse_args()

    try:
        from PIL import Image, ImageDraw
    except ImportError:
        fail("缺少 Pillow", "pip install pillow")

    try:
        img = Image.open(args.input).convert("RGB")
    except Exception as exc:
        fail(f"讀不到圖片：{exc}", "確認檔案完整；WebP 需要較新的 Pillow")

    W, H = img.size
    faces, method = detect_faces(args.input, args.cascade)

    if args.list:
        emit({
            "ok": True,
            "size": [W, H],
            "method": method,
            "faces": [
                {"index": i, "x": f[0], "y": f[1], "w": f[2], "h": f[3],
                 "center": [f[0] + f[2] // 2, f[1] + f[3] // 2]}
                for i, f in enumerate(sorted(faces, key=lambda f: f[0]))
            ],
        })

    if not args.output:
        fail("需要 --output（或用 --list 只看偵測結果）")

    if args.face:
        try:
            x, y, w, h = (int(v) for v in args.face.split(","))
        except ValueError:
            fail("--face 格式要是 x,y,w,h")
        method = "manual-box"
    else:
        chosen = choose(faces, args.pick)
        if chosen:
            x, y, w, h = chosen
        else:
            method = "heuristic-top-center"
            side = min(W, H)
            x, y, w, h = int(W / 2 - side * 0.25), int(min(H / 2, side * 0.30) - side * 0.25), \
                int(side * 0.5), int(side * 0.5)

    # 往外留邊，讓頭髮與肩膀進來一點，構圖才像頭像
    pad = max(w, h) * 0.55
    cx, cy = x + w / 2, y + h / 2 - h * 0.06
    side = min(max(w, h) + pad * 2, min(W, H))
    left = int(max(0, min(W - side, cx - side / 2)))
    top = int(max(0, min(H - side, cy - side / 2)))
    box = (left, top, int(left + side), int(top + side))

    face = img.crop(box).resize((args.size, args.size), Image.LANCZOS).convert("RGBA")
    radius = int(args.size * args.radius)
    mask = Image.new("L", (args.size, args.size), 0)
    ImageDraw.Draw(mask).rounded_rectangle([0, 0, args.size - 1, args.size - 1], radius=radius, fill=255)
    face.putalpha(mask)
    face.save(args.output, "PNG", optimize=True)

    emit({
        "ok": True,
        "method": method,
        "faces_found": len(faces),
        "source_size": [W, H],
        "face": [x, y, w, h],
        "box": list(box),
        "size": args.size,
        "output": args.output,
    })


if __name__ == "__main__":
    main()