hair/face_analysis/annotation.py

"""标注图层生成（透明底 RGBA PNG，仅标注、不含人物）。

规格（技术方案 §6）：线/字色 #FFFFFF、字体 10pt、线宽 1pt、透明底。
- 四庭水平分界线：numpy 向量化渐变消失（中间亮、两侧渐隐）。
- 四庭 cm 数值：图片左侧。
- 七眼标注：眼宽/两眼间距/脸宽，虚线带箭头，标签上下穿插。
中文字体用打包的思源黑体绝对路径加载，缺字体直接抛错（不静默降级成方块）。
"""
import os

import numpy as np
from PIL import Image, ImageDraw, ImageFont

FONT_PATH = os.path.join(os.path.dirname(__file__), "fonts", "NotoSansCJKsc-Regular.otf")
FONT_SIZE = 10
LINE_COLOR = (255, 255, 255, 255)   # #FFFFFF 100%
LINE_WIDTH = 1


def _load_font():
    if not os.path.isfile(FONT_PATH):
        raise FileNotFoundError(f"中文字体缺失：{FONT_PATH}（请按 OFFLINE_ASSETS.md 放置）")
    return ImageFont.truetype(FONT_PATH, FONT_SIZE)


def draw_gradient_horizontal_line(buf, cx, cy, color=LINE_COLOR, half_length=None):
    """在 RGBA numpy 缓冲 buf 上，以 (cx,cy) 为中心画向两侧渐变消失的水平线。

    numpy 向量化：一次性算整行 alpha，避免逐像素 draw.point。
    """
    h, w = buf.shape[:2]
    cy = int(round(cy)); cx = int(round(cx))
    if not (0 <= cy < h):
        return
    half = half_length or (w // 3)
    xs = np.arange(w)
    dist = np.abs(xs - cx)
    alpha = np.clip(1.0 - dist / half, 0.0, 1.0) * color[3]
    mask = alpha > 0
    row = buf[cy]
    row[mask, 0] = color[0]
    row[mask, 1] = color[1]
    row[mask, 2] = color[2]
    row[mask, 3] = np.maximum(row[mask, 3], alpha[mask].astype(np.uint8))


def draw_gradient_vertical_line(buf, cx, y0, y1, color=LINE_COLOR, fade=None):
    """在 RGBA numpy 缓冲 buf 上画一条竖线，两端渐变消失（中间实、上下淡）。"""
    h, w = buf.shape[:2]
    cx = int(round(cx))
    if not (0 <= cx < w):
        return
    y0, y1 = int(round(y0)), int(round(y1))
    y0, y1 = max(0, min(y0, y1)), min(h - 1, max(y0, y1))
    if y1 <= y0:
        return
    ys = np.arange(y0, y1 + 1)
    span = y1 - y0
    fade = fade or max(1, span // 5)               # 仅两端 ~1/5 段渐隐
    d = np.minimum(ys - y0, y1 - ys)               # 到最近端点的距离
    alpha = np.clip(d / fade, 0.0, 1.0) * color[3]
    col = buf[y0:y1 + 1, cx]
    m = alpha > 0
    col[m, 0] = color[0]
    col[m, 1] = color[1]
    col[m, 2] = color[2]
    col[m, 3] = np.maximum(col[m, 3], alpha[m].astype(np.uint8))


def draw_dashed_line_with_arrows(draw, x1, y1, x2, y2, color=LINE_COLOR,
                                 dash_len=6, gap_len=4, arrow_size=5):
    """两点间画虚线，两端带箭头（等腰三角）。"""
    total = ((x2 - x1) ** 2 + (y2 - y1) ** 2) ** 0.5
    if total == 0:
        return
    dx = (x2 - x1) / total
    dy = (y2 - y1) / total
    pos = 0.0
    while pos < total:
        seg_end = min(pos + dash_len, total)
        draw.line([(x1 + dx * pos, y1 + dy * pos),
                   (x1 + dx * seg_end, y1 + dy * seg_end)], fill=color, width=LINE_WIDTH)
        pos += dash_len + gap_len
    # 法向量（用于箭头两翼张开）
    nx, ny = -dy, dx
    for (ex, ey, sdx, sdy) in [(x1, y1, dx, dy), (x2, y2, -dx, -dy)]:
        p1 = (ex + sdx * arrow_size + nx * arrow_size * 0.6,
              ey + sdy * arrow_size + ny * arrow_size * 0.6)
        p2 = (ex + sdx * arrow_size - nx * arrow_size * 0.6,
              ey + sdy * arrow_size - ny * arrow_size * 0.6)
        draw.line([p1, (ex, ey)], fill=color, width=LINE_WIDTH)
        draw.line([p2, (ex, ey)], fill=color, width=LINE_WIDTH)


def _text_size(draw, text, font):
    bbox = draw.textbbox((0, 0), text, font=font)
    return bbox[2] - bbox[0], bbox[3] - bbox[1]


_LINE_NAMES = {
    "hair_top": "头顶",
    "hairline": "发际线",
    "brow_center": "眉心",
    "nose_bottom": "鼻翼下缘",
    "chin_tip": "下巴尖",
}


def create_annotated_image(image_bgr, measure_result):
    """生成标注图层 PNG（透明底 RGBA，尺寸同原图）。返回 PIL.Image。

    布局：
    - 横向 5 条分界线（渐变消失）：头顶/发际线/眉心/鼻翼下缘/下巴尖，左侧标线名 +
      四庭 cm 数值（顶庭/上庭/中庭/下庭，各段中点）。
    - 纵向 6 条线（渐变消失）：左脸颊/左眼外角/左眼内角/右眼内角/右眼外角/右脸颊，
      把脸宽切成 5 段；每段宽度在图的**顶部和底部**各标一次（只标 `X.XXcm`，不写名）。
    """
    h, w = image_bgr.shape[:2]
    v = measure_result.vertical
    pc = measure_result.px_per_cm

    buf = np.zeros((h, w, 4), dtype=np.uint8)

    # --- 1. 横向 5 条分界线（渐变） ---
    order = ["hair_top", "hairline", "brow_center", "nose_bottom", "chin_tip"]
    ys = [v[name][1] for name in order]
    cx_line = v["brow_center"][0]
    for cy in ys:
        draw_gradient_horizontal_line(buf, cx_line, cy)

    # --- 2. 纵向 6 条线（渐变），切出 5 段七眼宽度 ---
    pts = measure_result.eyes["points"]
    seven_keys = ["left_cheek", "left_outer", "left_inner",
                  "right_inner", "right_outer", "right_cheek"]
    xs = sorted(pts[k][0] for k in seven_keys)   # 自左向右
    # 顶/底各两行错开放置，避免窄段(眼宽~2cm)数字横向重叠
    row_h = FONT_SIZE + 2
    y_top0, y_top1 = 6, 6 + row_h
    y_bot0 = h - FONT_SIZE - 6
    y_bot1 = y_bot0 - row_h
    yv0 = y_top1 + FONT_SIZE + 4                  # 竖线避开顶部两行标签
    yv1 = y_bot1 - 4                              # 避开底部两行标签
    for vx in xs:
        draw_gradient_vertical_line(buf, vx, yv0, yv1)

    canvas = Image.fromarray(buf, mode="RGBA")
    draw = ImageDraw.Draw(canvas)
    font = _load_font()

    # --- 3. 横线左侧：线名 + 四庭 cm 数值 ---
    left_margin = 8
    court_indent = 30
    court_cm = [measure_result.top_cm, measure_result.upper_cm,
                measure_result.middle_cm, measure_result.lower_cm]
    court_name = ["顶庭", "上庭", "中庭", "下庭"]
    for i, name in enumerate(order):
        draw.text((left_margin, ys[i] - FONT_SIZE / 2), _LINE_NAMES[name],
                  fill=LINE_COLOR, font=font)
    for i in range(4):
        y_mid = (ys[i] + ys[i + 1]) / 2 - FONT_SIZE / 2
        draw.text((left_margin + court_indent, y_mid),
                  f"{court_name[i]} {court_cm[i]:.2f}cm", fill=LINE_COLOR, font=font)

    # --- 4. 七眼每段宽度：顶部 + 底部各标一次（只标数字 cm，相邻段上下错行） ---
    for i in range(len(xs) - 1):
        seg_cm = (xs[i + 1] - xs[i]) / pc
        cx_seg = (xs[i] + xs[i + 1]) / 2
        text = f"{seg_cm:.2f}cm"
        tw, _ = _text_size(draw, text, font)
        ty_top = y_top0 if i % 2 == 0 else y_top1
        ty_bot = y_bot0 if i % 2 == 0 else y_bot1
        draw.text((cx_seg - tw / 2, ty_top), text, fill=LINE_COLOR, font=font)
        draw.text((cx_seg - tw / 2, ty_bot), text, fill=LINE_COLOR, font=font)

    return canvas


if __name__ == "__main__":
    import sys
    import time
    import cv2
    from face_analysis.detector import detector
    from face_analysis.measure import measure_face

    path = sys.argv[1] if len(sys.argv) > 1 else "tests/fixtures/frontal.jpg"
    out = sys.argv[2] if len(sys.argv) > 2 else "tests/output/annotated.png"
    img = cv2.imread(path)
    if img is None:
        print(f"无法读取图片: {path}")
        sys.exit(1)
    h, w = img.shape[:2]
    lms = detector.detect(img)
    if lms is None:
        print("未检出人脸")
        sys.exit(1)
    mask = None
    try:
        from face_analysis.hair_segmenter import get_segmenter
        mask = get_segmenter().segment_hair(img)
    except Exception as e:  # noqa: BLE001
        print(f"[warn] 分割不可用，回退方案 A：{e}")

    result = measure_face(lms, mask, w, h)
    t0 = time.time()
    canvas = create_annotated_image(img, result)
    dt = time.time() - t0
    os.makedirs(os.path.dirname(out), exist_ok=True)
    canvas.save(out)
    arr = np.asarray(canvas)
    print(f"saved {out} mode={canvas.mode} size={canvas.size} "
          f"transparent={bool((arr[:,:,3]==0).any())} opaque={bool((arr[:,:,3]>0).any())} "
          f"elapsed={dt*1000:.1f}ms")