save code
This commit is contained in:
@@ -0,0 +1,312 @@
|
||||
"""
|
||||
人脸贴图渲染核心：接受一张图 + 一个效果 ID，输出 PNG 字节。
|
||||
|
||||
流程（与 Android 端 vulkan/FaceApp.cpp 一致，CPU 实现）：
|
||||
1. MediaPipe FaceLandmarker → 478 normalized landmarks per face
|
||||
2. 取最大人脸（按 landmark bbox 面积）
|
||||
3. 把 obj 的 468 顶点位置改写为 landmarks[INDEX_MAP[i]]，得到 dst 顶点
|
||||
4. 对 mesh 的每个三角形：从 512×512 effect 贴图取 src 三角形，仿射 warp
|
||||
到 dst 三角形位置，按贴图 alpha 通道 alpha-blend 到原图副本上
|
||||
5. 输出 PNG（保持原图尺寸）
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image, ImageOps
|
||||
|
||||
from .index_map import INDEX_MAP
|
||||
from .mesh import FaceMesh, load_face_mesh
|
||||
|
||||
|
||||
class NoFaceError(Exception):
|
||||
"""图里检测不到人脸。"""
|
||||
|
||||
|
||||
class EffectNotFoundError(Exception):
|
||||
"""指定 effect_id 没有对应贴图。"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class RendererPaths:
|
||||
obj_path: Path
|
||||
effects_dir: Path
|
||||
|
||||
|
||||
class FaceRenderer:
|
||||
"""
|
||||
单例式渲染器：mesh / mediapipe 模型加载一次复用。
|
||||
|
||||
线程安全：MediaPipe FaceLandmarker 的 detect() 不保证多线程并发，
|
||||
用 _lock 串行化。FastAPI 单进程同步部署下，请求本身就串行，开销可忽略。
|
||||
"""
|
||||
|
||||
def __init__(self, paths: RendererPaths, max_image_bytes: int = 20 * 1024 * 1024):
|
||||
self.paths = paths
|
||||
self.max_image_bytes = max_image_bytes
|
||||
self._mesh: FaceMesh = load_face_mesh(paths.obj_path)
|
||||
if len(self._mesh.uvs) != len(INDEX_MAP):
|
||||
raise RuntimeError(
|
||||
f"mesh has {len(self._mesh.uvs)} vertices but INDEX_MAP has "
|
||||
f"{len(INDEX_MAP)}; obj/index_map.py 不匹配"
|
||||
)
|
||||
# 预转 numpy 加速
|
||||
self._index_map = np.asarray(INDEX_MAP, dtype=np.int32)
|
||||
|
||||
self._effect_cache: dict[int, np.ndarray] = {}
|
||||
self._cache_lock = threading.Lock()
|
||||
|
||||
# MediaPipe 延迟加载（首次请求时加载，缩短启动时间，便于本地调试）
|
||||
self._landmarker = None
|
||||
self._landmarker_lock = threading.Lock()
|
||||
|
||||
# ---------- public ----------
|
||||
|
||||
def list_effect_ids(self) -> list[int]:
|
||||
"""枚举 effects/<n>.png 中的 n。"""
|
||||
ids: list[int] = []
|
||||
if not self.paths.effects_dir.is_dir():
|
||||
return ids
|
||||
for p in self.paths.effects_dir.iterdir():
|
||||
if p.suffix.lower() != ".png":
|
||||
continue
|
||||
try:
|
||||
n = int(p.stem)
|
||||
except ValueError:
|
||||
continue
|
||||
if 1 <= n <= 99:
|
||||
ids.append(n)
|
||||
return sorted(ids)
|
||||
|
||||
def render(self, image_bytes: bytes, effect_id: int) -> bytes:
|
||||
if len(image_bytes) > self.max_image_bytes:
|
||||
raise ValueError("image too large")
|
||||
|
||||
effect_rgba = self._load_effect(effect_id) # (H,W,4) uint8
|
||||
|
||||
# 1. 解码 + 处理 EXIF orientation → RGB ndarray
|
||||
img_rgb = self._decode_input(image_bytes)
|
||||
h, w = img_rgb.shape[:2]
|
||||
|
||||
# 2. MediaPipe 检测，挑最大人脸
|
||||
landmarks_norm = self._detect_largest_face(img_rgb) # (N,2) ∈ [0,1]² (N>=468)
|
||||
|
||||
# 3. obj 顶点 → 像素坐标
|
||||
dst_xy = self._compute_dst_vertices(landmarks_norm, w, h)
|
||||
|
||||
# 4. warp + 混合
|
||||
out = img_rgb.copy()
|
||||
self._composite(out, effect_rgba, dst_xy)
|
||||
|
||||
# 5. 编码 PNG
|
||||
return self._encode_png(out)
|
||||
|
||||
# ---------- internals ----------
|
||||
|
||||
def _load_effect(self, effect_id: int) -> np.ndarray:
|
||||
if not (1 <= effect_id <= 99):
|
||||
raise EffectNotFoundError(f"effect_id out of range: {effect_id}")
|
||||
with self._cache_lock:
|
||||
cached = self._effect_cache.get(effect_id)
|
||||
if cached is not None:
|
||||
return cached
|
||||
png_path = self.paths.effects_dir / f"{effect_id}.png"
|
||||
if not png_path.is_file():
|
||||
raise EffectNotFoundError(f"effect {effect_id} not found")
|
||||
# 用 PIL 解码以保证 RGBA 一致；OpenCV imread 对部分 PNG 会丢 alpha
|
||||
img = Image.open(png_path).convert("RGBA")
|
||||
arr = np.array(img, dtype=np.uint8) # (H,W,4) RGBA
|
||||
with self._cache_lock:
|
||||
self._effect_cache[effect_id] = arr
|
||||
return arr
|
||||
|
||||
def _decode_input(self, image_bytes: bytes) -> np.ndarray:
|
||||
try:
|
||||
pil = Image.open(io.BytesIO(image_bytes))
|
||||
pil = ImageOps.exif_transpose(pil) # 处理手机 EXIF 旋转
|
||||
pil = pil.convert("RGB")
|
||||
except Exception as e:
|
||||
raise ValueError(f"invalid image: {e}") from e
|
||||
return np.array(pil, dtype=np.uint8)
|
||||
|
||||
def _ensure_landmarker(self):
|
||||
if self._landmarker is not None:
|
||||
return
|
||||
with self._landmarker_lock:
|
||||
if self._landmarker is not None:
|
||||
return
|
||||
# 延迟 import，避免单元测试时强依赖 mediapipe
|
||||
from mediapipe.tasks import python as mp_python
|
||||
from mediapipe.tasks.python import vision as mp_vision
|
||||
|
||||
model_path = self._find_model_file()
|
||||
base_options = mp_python.BaseOptions(model_asset_path=str(model_path))
|
||||
options = mp_vision.FaceLandmarkerOptions(
|
||||
base_options=base_options,
|
||||
running_mode=mp_vision.RunningMode.IMAGE,
|
||||
num_faces=5, # 检测最多 5 张，挑最大那张
|
||||
min_face_detection_confidence=0.5,
|
||||
min_face_presence_confidence=0.5,
|
||||
min_tracking_confidence=0.5,
|
||||
output_face_blendshapes=False,
|
||||
output_facial_transformation_matrixes=False,
|
||||
)
|
||||
self._landmarker = mp_vision.FaceLandmarker.create_from_options(options)
|
||||
|
||||
def _find_model_file(self) -> Path:
|
||||
"""优先用本服务 assets 下的，找不到再回退到仓库 app/src/main/assets。"""
|
||||
local = self.paths.obj_path.parent / "face_landmarker.task"
|
||||
if local.is_file():
|
||||
return local
|
||||
# 回退到仓库 app/src/main/assets/face_landmarker.task（Android SDK 用同一份）
|
||||
repo_root = self.paths.obj_path.resolve().parents[3]
|
||||
fallback = repo_root / "app" / "src" / "main" / "assets" / "face_landmarker.task"
|
||||
if fallback.is_file():
|
||||
return fallback
|
||||
raise FileNotFoundError(
|
||||
"face_landmarker.task not found; place it under "
|
||||
f"{local} or {fallback}"
|
||||
)
|
||||
|
||||
def _detect_largest_face(self, img_rgb: np.ndarray) -> np.ndarray:
|
||||
self._ensure_landmarker()
|
||||
import mediapipe as mp
|
||||
|
||||
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=img_rgb)
|
||||
with self._landmarker_lock:
|
||||
result = self._landmarker.detect(mp_image)
|
||||
|
||||
faces = result.face_landmarks or []
|
||||
if not faces:
|
||||
raise NoFaceError("no face detected")
|
||||
|
||||
# 选 bbox 面积最大的那张
|
||||
def bbox_area(landmarks) -> float:
|
||||
xs = [lm.x for lm in landmarks]
|
||||
ys = [lm.y for lm in landmarks]
|
||||
return (max(xs) - min(xs)) * (max(ys) - min(ys))
|
||||
|
||||
biggest = max(faces, key=bbox_area)
|
||||
# 转 ndarray，只取 xy
|
||||
n = len(biggest)
|
||||
arr = np.empty((n, 2), dtype=np.float32)
|
||||
for i, lm in enumerate(biggest):
|
||||
arr[i, 0] = lm.x
|
||||
arr[i, 1] = lm.y
|
||||
return arr
|
||||
|
||||
def _compute_dst_vertices(
|
||||
self, landmarks_norm: np.ndarray, w: int, h: int
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
landmarks_norm: (N, 2) ∈ [0,1]²
|
||||
返回 (468, 2) float32 像素坐标，对应 mesh 的 468 个顶点。
|
||||
"""
|
||||
n_landmarks = landmarks_norm.shape[0]
|
||||
if self._index_map.max() >= n_landmarks:
|
||||
raise RuntimeError(
|
||||
f"INDEX_MAP references landmark index {int(self._index_map.max())} "
|
||||
f"but mediapipe returned only {n_landmarks} landmarks"
|
||||
)
|
||||
picked = landmarks_norm[self._index_map] # (468, 2)
|
||||
out = np.empty_like(picked)
|
||||
out[:, 0] = picked[:, 0] * w
|
||||
out[:, 1] = picked[:, 1] * h
|
||||
return out.astype(np.float32)
|
||||
|
||||
def _composite(
|
||||
self,
|
||||
canvas_rgb: np.ndarray,
|
||||
effect_rgba: np.ndarray,
|
||||
dst_xy: np.ndarray,
|
||||
) -> None:
|
||||
"""
|
||||
在 canvas_rgb（uint8 H×W×3，原地修改）上对 mesh 每个三角形：
|
||||
从 effect_rgba 取对应 src 三角形 → warpAffine 到 dst → alpha 混合。
|
||||
"""
|
||||
eh, ew = effect_rgba.shape[:2]
|
||||
canvas_h, canvas_w = canvas_rgb.shape[:2]
|
||||
uvs = self._mesh.uvs # (468, 2)
|
||||
tris = self._mesh.triangles # (852, 3)
|
||||
|
||||
# 把 UV 一次性转成 effect 像素坐标
|
||||
src_pts_all = np.empty_like(uvs)
|
||||
src_pts_all[:, 0] = uvs[:, 0] * ew
|
||||
src_pts_all[:, 1] = uvs[:, 1] * eh
|
||||
|
||||
effect_rgb = effect_rgba[..., :3]
|
||||
effect_a = effect_rgba[..., 3]
|
||||
|
||||
for tri in tris:
|
||||
i0, i1, i2 = int(tri[0]), int(tri[1]), int(tri[2])
|
||||
src_tri = np.float32([src_pts_all[i0], src_pts_all[i1], src_pts_all[i2]])
|
||||
dst_tri = np.float32([dst_xy[i0], dst_xy[i1], dst_xy[i2]])
|
||||
|
||||
# 用 dst bbox 限定计算区域，避免对全画布 warp
|
||||
x_min = int(np.floor(dst_tri[:, 0].min()))
|
||||
y_min = int(np.floor(dst_tri[:, 1].min()))
|
||||
x_max = int(np.ceil(dst_tri[:, 0].max()))
|
||||
y_max = int(np.ceil(dst_tri[:, 1].max()))
|
||||
|
||||
# clip 到画布
|
||||
x_min_c = max(0, x_min)
|
||||
y_min_c = max(0, y_min)
|
||||
x_max_c = min(canvas_w, x_max)
|
||||
y_max_c = min(canvas_h, y_max)
|
||||
if x_max_c <= x_min_c or y_max_c <= y_min_c:
|
||||
continue
|
||||
|
||||
box_w = x_max_c - x_min_c
|
||||
box_h = y_max_c - y_min_c
|
||||
|
||||
# 构造从 src → 局部 bbox 坐标系（左上为原点）的仿射
|
||||
dst_tri_local = dst_tri - np.array([x_min_c, y_min_c], dtype=np.float32)
|
||||
M = cv2.getAffineTransform(src_tri, dst_tri_local)
|
||||
|
||||
# 给 RGB 和 alpha 分别 warp 到 bbox 大小
|
||||
warped_rgb = cv2.warpAffine(
|
||||
effect_rgb,
|
||||
M,
|
||||
(box_w, box_h),
|
||||
flags=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_REPLICATE,
|
||||
)
|
||||
warped_a = cv2.warpAffine(
|
||||
effect_a,
|
||||
M,
|
||||
(box_w, box_h),
|
||||
flags=cv2.INTER_LINEAR,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=0,
|
||||
)
|
||||
|
||||
# 限制只在三角形内部混合（防止 bbox 多余像素污染）
|
||||
tri_mask = np.zeros((box_h, box_w), dtype=np.uint8)
|
||||
cv2.fillConvexPoly(
|
||||
tri_mask,
|
||||
dst_tri_local.astype(np.int32),
|
||||
255,
|
||||
lineType=cv2.LINE_AA,
|
||||
)
|
||||
alpha = (
|
||||
warped_a.astype(np.float32) * (tri_mask.astype(np.float32) / 255.0)
|
||||
) / 255.0
|
||||
alpha = alpha[..., None] # (h,w,1)
|
||||
|
||||
roi = canvas_rgb[y_min_c:y_max_c, x_min_c:x_max_c].astype(np.float32)
|
||||
blended = roi * (1.0 - alpha) + warped_rgb.astype(np.float32) * alpha
|
||||
canvas_rgb[y_min_c:y_max_c, x_min_c:x_max_c] = np.clip(
|
||||
blended, 0, 255
|
||||
).astype(np.uint8)
|
||||
|
||||
def _encode_png(self, img_rgb: np.ndarray) -> bytes:
|
||||
pil = Image.fromarray(img_rgb)
|
||||
buf = io.BytesIO()
|
||||
pil.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
Reference in New Issue
Block a user