xsl
db32fa12e5
- 从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张发际线贴图入库
43 lines
1.6 KiB
Python
43 lines
1.6 KiB
Python
"""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
|