hair/face_analysis/hair_segmenter.py

"""方案 B：BiSeNet 头发分割 + 发际线/头顶定位。

加载 face-parsing BiSeNet（19 类，hair=17），对整图做像素级语义分割得到头发
mask，再沿面部中轴线扫描得到真实发际线与头顶。GPU 可用时走 CUDA，否则 CPU。
单例加载权重，避免每请求重载。详见技术方案 §1.4 / §4.0。
"""
import os

import cv2
import numpy as np

# ⚠️ torch / torchvision / BiSeNet 仅在 HairSegmenter.__init__ 内惰性导入，
#    使本模块的纯 numpy 函数 locate_hairline_by_segmentation 可在无 torch/GPU
#    的环境（如 Tier-1 合成几何测试、方案 A only 降级版）被安全导入。

_WEIGHTS = os.path.join(os.path.dirname(__file__), "weights", "79999_iter.pth")
HAIR_CLASS = 17          # CelebAMask-HQ 19 类中 hair 的索引
N_CLASSES = 19
_INPUT_SIZE = 512        # BiSeNet 推理输入边长


def _select_device(torch):
    """选择推理设备：优先 CUDA，但实测一次小算子确认当前 GPU 架构被本 torch 支持。

    场景：本机为 RTX 5090（sm_120/Blackwell），而 torch 2.2.2+cu121 仅编译到 sm_90，
    .cuda() 会在执行时抛 "no kernel image is available"。此处用一次小 matmul 探测，
    失败则回退 CPU（BiSeNet CPU 推理 ~0.3–1s/张，方案B 仍可用）。
    换装支持 sm_120 的 torch（cu128）后会自动改用 GPU，无需改代码。
    可用环境变量 FORCE_CPU=1 强制 CPU。
    """
    import os as _os
    if _os.getenv("FORCE_CPU") == "1" or not torch.cuda.is_available():
        return torch.device("cpu")
    try:
        _ = (torch.zeros(8, 8, device="cuda") @ torch.zeros(8, 8, device="cuda")).cpu()
        return torch.device("cuda")
    except Exception:  # noqa: BLE001  GPU 架构不被支持 → 回退 CPU
        return torch.device("cpu")


class HairSegmenter:
    """BiSeNet 头发分割封装。建议经 get_segmenter() 取单例。"""

    def __init__(self, weights_path=_WEIGHTS):
        import torch
        import torchvision.transforms as transforms

        from face_analysis.bisenet_model import BiSeNet

        self._torch = torch
        self.device = _select_device(torch)
        self.net = BiSeNet(n_classes=N_CLASSES)
        state = torch.load(weights_path, map_location="cpu")
        self.net.load_state_dict(state)
        self.net.to(self.device)
        self.net.eval()
        self._to_tensor = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        ])

    def segment_hair(self, image_bgr):
        """返回 hair_mask（H×W bool，True=头发），尺寸同输入原图。"""
        torch = self._torch
        h, w = image_bgr.shape[:2]
        rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
        resized = cv2.resize(rgb, (_INPUT_SIZE, _INPUT_SIZE),
                             interpolation=cv2.INTER_LINEAR)
        inp = self._to_tensor(resized).unsqueeze(0).to(self.device)
        with torch.no_grad():
            out = self.net(inp)[0]                        # 主输出 (1, C, 512, 512)
        parsing = out.squeeze(0).argmax(0).cpu().numpy()  # (512, 512) 类别图
        hair_small = (parsing == HAIR_CLASS).astype(np.uint8)
        # 还原到原图尺寸（最近邻保持类别边界）
        hair_mask = cv2.resize(hair_small, (w, h), interpolation=cv2.INTER_NEAREST)
        return hair_mask.astype(bool)


_segmenter = None


def get_segmenter():
    """惰性单例：首次调用时加载权重（并占用显存），后续复用。"""
    global _segmenter
    if _segmenter is None:
        _segmenter = HairSegmenter()
    return _segmenter


def locate_hairline_by_segmentation(hair_mask, brow_center_x, image_height):
    """从头发 mask 定位发际线与头顶。

    Args:
        hair_mask: H×W bool/uint8，True=头发。
        brow_center_x: 面部中轴线 x（像素）。
        image_height: 图高（保留参数，便于后续边界判断）。
    Returns:
        (hairline_y, hair_top_y) 像素坐标；失败返回 None（交给方案 A 兜底）。
    """
    if hair_mask is None:
        return None
    mask = np.asarray(hair_mask).astype(bool)
    if mask.sum() == 0:
        return None

    w = mask.shape[1]
    cx = int(round(brow_center_x))
    cx = max(0, min(cx, w - 1))
    # 中轴线附近窄列带（±3px）求稳，避免单列噪声
    band = mask[:, max(0, cx - 3): min(w, cx + 4)]
    col = band.any(axis=1)
    hair_rows = np.where(col)[0]
    if hair_rows.size == 0:
        return None

    # 发际线：中轴线列带上头发区域最靠下的行（头发→皮肤交界，y 向下为正）
    hairline_y = int(hair_rows.max())
    # 头顶：整张头发 mask 的最高点（最小 y），用全图更鲁棒
    top_rows = np.where(mask.any(axis=1))[0]
    hair_top_y = int(top_rows.min())

    # 合理性校验：头顶必须严格在发际线上方
    if hair_top_y >= hairline_y:
        return None
    return hairline_y, hair_top_y


if __name__ == "__main__":
    import sys
    from face_analysis.detector import detector
    from face_analysis.calibration import normalized_to_pixel
    from face_analysis.face_mesh_landmarks import GLABELLA_9, GLABELLA_151

    path = sys.argv[1] if len(sys.argv) > 1 else "tests/fixtures/frontal.jpg"
    img = cv2.imread(path)
    if img is None:
        print(f"无法读取图片: {path}")
        sys.exit(1)
    h, w = img.shape[:2]
    segmenter = get_segmenter()
    print("device:", segmenter.device)
    mask = segmenter.segment_hair(img)
    print("hair pixels:", int(mask.sum()))

    lms = detector.detect(img)
    if lms is None:
        print("未检出人脸，跳过定位")
        sys.exit(0)
    lm = lms.landmark
    bx = (normalized_to_pixel(lm[GLABELLA_9], w, h)[0]
          + normalized_to_pixel(lm[GLABELLA_151], w, h)[0]) / 2
    by = (normalized_to_pixel(lm[GLABELLA_9], w, h)[1]
          + normalized_to_pixel(lm[GLABELLA_151], w, h)[1]) / 2
    res = locate_hairline_by_segmentation(mask, bx, h)
    if res is None:
        print("定位失败（将回退方案 A）")
    else:
        hairline_y, hair_top_y = res
        print(f"brow_y={by:.1f} hairline_y={hairline_y} hair_top_y={hair_top_y}")
        print("自洽校验 hair_top_y < hairline_y < brow_y:",
              hair_top_y < hairline_y < by)

    # dump mask 预览
    os.makedirs("tests/output", exist_ok=True)
    cv2.imwrite("tests/output/hair_mask.png", (mask.astype(np.uint8) * 255))
    print("mask 预览已存 tests/output/hair_mask.png")