feat(worker): 接口1 四庭七眼测量真实实现（替换 Mock）
worker 侧从 Mock 替换为真实算法： - face_analysis 包：detector(MediaPipe 478点) / pose(solvePnP 姿态) / calibration(虹膜直径法) / hair_segmenter+bisenet_model(方案B 头发分割) / measure(方案A兜底+B/A决策+七眼+换算) / annotation(numpy渐变线+中文标注) - app.py：/api/v1/face/measure 接真实实现，返回 annotated_image_base64 （不落盘不拼URL，落盘由网关做）；加 X-Internal-Token 鉴权、/health 就绪态、 可配置分辨率门槛、异常兜底 - 部署：start.sh/run_worker.sh/hair-worker.service 监听 8187；worker_config 示例 - 测试 tests/：Tier1合成真值<1e-6 + Tier2缩放不变 + Tier3叠加 + 错误码集成 + 数值回归，pytest 24 项全绿 - 文档补实测基线表 + RTX5090/torch 说明 注：worker 为 RTX 5090(sm_120)，pinned torch 2.2.2(cu121) 只到 sm_90， BiSeNet 已自动回退 CPU（方案B 正常）；要用 GPU 需换 torch cu128(≥2.7)。 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Claude Opus 4.8
@@ -0,0 +1,166 @@
|
||||
"""方案 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")
|
||||
Reference in New Issue
Block a user