feat(接口2): 移植head3d发际线管线 + 接口2实现方案文档
- 从head3d复制发际线检测管线到 hairline/ 包:MediaPipe Tasks + SegFormer分割 + 17锚点射线检测 + 502点mesh(face_ext.obj)+UV - 复制模型:face_landmarker.task(3.7MB)、SegFormer config/preprocessor (model.safetensors 340MB 单独下载中) - 新增 docs/接口2-C端生发-技术实现方案.md:第一步=发际线曲线叠加预览图, 新增gender必填参数,按性别贴图数量输出(female5/male4),hairline_type英文key, 服务端cv2逐三角形warp渲染器(head3d只有浏览器端Three.js渲染) - 接口文档.md 接口2章节同步:gender参数、输出语义、错误码说明 - hairline_texture/ 9张发际线贴图入库
This commit is contained in:
@@ -0,0 +1,42 @@
|
||||
"""Face parsing via HuggingFace SegFormer (jonathandinu/face-parsing).
|
||||
|
||||
First call downloads ~150 MB of weights into the HF cache.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import numpy as np
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from . import constants as C
|
||||
|
||||
if TYPE_CHECKING: # avoid hard import at module load
|
||||
pass
|
||||
|
||||
|
||||
class FaceParser:
|
||||
"""Wraps a SegFormer face-parsing model and returns a (H, W) int label map."""
|
||||
|
||||
def __init__(self, device: str | None = None):
|
||||
import torch
|
||||
from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
|
||||
|
||||
self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
||||
self.processor = SegformerImageProcessor.from_pretrained(C.HF_FACE_PARSER_MODEL)
|
||||
self.model = SegformerForSemanticSegmentation.from_pretrained(C.HF_FACE_PARSER_MODEL)
|
||||
self.model.to(self.device).eval()
|
||||
self._torch = torch
|
||||
|
||||
def parse(self, image_rgb: np.ndarray) -> np.ndarray:
|
||||
"""image_rgb: (H, W, 3) uint8. Returns (H, W) int64 of class indices."""
|
||||
from PIL import Image
|
||||
|
||||
H, W = image_rgb.shape[:2]
|
||||
pil = Image.fromarray(image_rgb)
|
||||
inputs = self.processor(images=pil, return_tensors="pt").to(self.device)
|
||||
with self._torch.no_grad():
|
||||
logits = self.model(**inputs).logits # (1, C, h, w)
|
||||
# Upsample to original resolution
|
||||
up = self._torch.nn.functional.interpolate(
|
||||
logits, size=(H, W), mode="bilinear", align_corners=False
|
||||
)
|
||||
labels = up.argmax(dim=1).squeeze(0).to("cpu").numpy().astype(np.int32)
|
||||
return labels
|
||||
Reference in New Issue
Block a user