200 lines
7.6 KiB
Python
200 lines
7.6 KiB
Python
"""接口4：用户面部特征分析。
|
||
|
||
- 眉形 / 面部年龄 / 动静类型 / 性别 / 基因风格：火山方舟豆包视觉模型
|
||
- face_shape（脸型）：本机 MediaPipe 分类（face/face_shape_classifier.py）覆盖，
|
||
不用豆包结果
|
||
|
||
⚠️ 仍依赖外网豆包（其余 5 项）。API Key 走配置/环境变量，不入 git。
|
||
网关本机需可 import face 包（opencv + mediapipe）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import json
|
||
import logging
|
||
import os
|
||
|
||
logger = logging.getLogger("hair.worker")
|
||
|
||
ARK_BASE_URL = os.getenv("ARK_BASE_URL", "https://ark.cn-beijing.volces.com/api/v3")
|
||
ARK_MODEL = os.getenv("ARK_MODEL", "doubao-seed-2-0-lite-260428")
|
||
|
||
# doubao 中文键 → 接口4 英文字段（脸型不走豆包，见 _local_face_shape）
|
||
_KEY_MAP = {
|
||
"眉形": "eyebrow_shape",
|
||
"面部年龄": "facial_age",
|
||
"动静类型": "dynamic_static_type",
|
||
"性别": "gender",
|
||
"基因风格": "gene_style",
|
||
}
|
||
|
||
# 豆包只问 5 项 + 是否有人脸；脸型由本地分类器给出
|
||
_PROMPT = (
|
||
"分析一下图片告诉我以下特征，只要答案，格式为json字符串，"
|
||
"图片是否有人脸（有人/没人） "
|
||
"眉形 "
|
||
"面部年龄（给出区间年龄） 动静类型（静态型/动态型） 性别（男/女） "
|
||
"基因风格（戏剧型/睿智型/自然型/古典型/优雅型/浪漫型/前卫型/少女型/少年型）"
|
||
)
|
||
|
||
_client = None # 缓存的 Ark client（api_key 变更时自动重建）
|
||
_client_key: str | None = None # _client 构建时使用的 api_key，用于检测配置变更
|
||
|
||
|
||
def _load_api_key() -> str | None:
|
||
"""ARK_API_KEY 环境变量优先；否则优先 gateway/config.json（接口4 已迁网关），
|
||
再回退 worker_config.json（兼容旧配置）。"""
|
||
key = os.getenv("ARK_API_KEY")
|
||
if key:
|
||
return key
|
||
base = os.path.dirname(__file__)
|
||
for cfg_name in ("gateway/config.json", "worker_config.json"):
|
||
cfg = os.path.join(base, cfg_name)
|
||
if os.path.isfile(cfg):
|
||
try:
|
||
with open(cfg, encoding="utf-8") as f:
|
||
v = json.load(f).get("ark_api_key")
|
||
if v:
|
||
return v
|
||
except Exception as e: # noqa: BLE001
|
||
logger.warning("读取 %s ark_api_key 失败：%s", cfg_name, e)
|
||
return None
|
||
|
||
|
||
def get_client():
|
||
"""返回 Ark client。
|
||
|
||
client 全局缓存，但每次都会重新读取 api_key —— 一旦配置（环境变量 /
|
||
worker_config.json / gateway/config.json 的 ark_api_key）发生变化，
|
||
自动重建 client。这样换 key 后无需重启进程。
|
||
"""
|
||
global _client, _client_key
|
||
key = _load_api_key()
|
||
if not key:
|
||
raise RuntimeError("缺少火山方舟 API Key（设 ARK_API_KEY 或 worker_config.json.ark_api_key）")
|
||
# client 未建、或 key 变了 → 重建
|
||
if _client is None or key != _client_key:
|
||
from volcenginesdkarkruntime import Ark
|
||
_client = Ark(base_url=ARK_BASE_URL, api_key=key)
|
||
_client_key = key
|
||
return _client
|
||
|
||
|
||
def _parse_json(text: str) -> dict:
|
||
"""去掉 ```json 包裹后解析。"""
|
||
s = text.strip()
|
||
if s.startswith("```"):
|
||
s = s.strip("`")
|
||
if s[:4].lower() == "json":
|
||
s = s[4:]
|
||
return json.loads(s.strip())
|
||
|
||
|
||
def _image_to_url(image_bytes: bytes = None, image_url: str = None) -> str:
|
||
"""优先用现成 URL；否则把字节转 base64 data URI（doubao 兼容）。"""
|
||
if image_url:
|
||
return image_url
|
||
fmt = "png" if image_bytes[:8] == b"\x89PNG\r\n\x1a\n" else "jpeg"
|
||
return f"data:image/{fmt};base64," + base64.b64encode(image_bytes).decode()
|
||
|
||
|
||
def _resolve_image_bytes(image_bytes: bytes = None, image_url: str = None) -> bytes:
|
||
"""本地分类器用：优先已有字节；仅有 URL 时下载。"""
|
||
if image_bytes:
|
||
return image_bytes
|
||
if not image_url:
|
||
raise ValueError("缺少图片数据")
|
||
if image_url.startswith("data:"):
|
||
# data URI
|
||
b64 = image_url.split(",", 1)[1] if "," in image_url else image_url
|
||
return base64.b64decode(b64)
|
||
import httpx
|
||
with httpx.Client(timeout=15.0, follow_redirects=True) as client:
|
||
r = client.get(image_url)
|
||
r.raise_for_status()
|
||
return r.content
|
||
|
||
|
||
def _local_face_shape(image_bytes: bytes = None, image_url: str = None) -> str:
|
||
"""MediaPipe 脸型分类，返回 display（含混合脸型描述）或主脸型。"""
|
||
import cv2
|
||
import numpy as np
|
||
from face.face_shape_classifier import classify_from_image
|
||
|
||
raw = _resolve_image_bytes(image_bytes, image_url)
|
||
bgr = cv2.imdecode(np.frombuffer(raw, np.uint8), cv2.IMREAD_COLOR)
|
||
if bgr is None:
|
||
raise ValueError("图片格式不支持，无法解码")
|
||
result = classify_from_image(bgr, return_details=True, return_annotated=False)
|
||
shape = result.get("display") or result["face_shape"]
|
||
logger.info(
|
||
"local face_shape=%s conf=%.3f",
|
||
shape,
|
||
float(result.get("confidence") or 0),
|
||
)
|
||
return shape
|
||
|
||
|
||
def analyze_features(image_bytes: bytes = None, image_url: str = None):
|
||
"""豆包分析 5 项特征 + 本机 MediaPipe 覆盖 face_shape。
|
||
|
||
Returns: dict —— 6 个英文字段(face_shape/eyebrow_shape/facial_age/
|
||
dynamic_static_type/gender/gene_style)；**无人脸返回 None**（调用方据此判 1001）。
|
||
"""
|
||
url = _image_to_url(image_bytes, image_url)
|
||
resp = get_client().chat.completions.create(
|
||
model=ARK_MODEL,
|
||
messages=[{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "image_url", "image_url": {"url": url}},
|
||
{"type": "text", "text": _PROMPT},
|
||
],
|
||
}],
|
||
max_tokens=1024, # 限制输出长度，模型秒回
|
||
temperature=0, # 固定输出，无随机采样，提速+结果稳定
|
||
stream=False, # 关闭流式，单次返回结果更快
|
||
extra_body={
|
||
"thinking": {
|
||
"type": "disabled", # 彻底关闭深度思考模式，提速50%+
|
||
},
|
||
},
|
||
)
|
||
text = resp.choices[0].message.content
|
||
logger.info("doubao raw response (first 500 chars): %s", text[:500])
|
||
try:
|
||
raw = _parse_json(text) # doubao 原始中文字段
|
||
except (json.JSONDecodeError, ValueError) as e:
|
||
logger.error("doubao 返回非 JSON，原文: %s", text[:1000])
|
||
raise RuntimeError(f"豆包模型返回格式异常，无法解析为 JSON：{text[:200]}") from e
|
||
if not has_face(raw):
|
||
return None
|
||
# 豆包 5 项 + 本地脸型覆盖
|
||
feats = {en: raw[zh] for zh, en in _KEY_MAP.items() if zh in raw}
|
||
try:
|
||
feats["face_shape"] = _local_face_shape(image_bytes, image_url)
|
||
except ValueError as e:
|
||
logger.warning("本地脸型分类未检测到人脸: %s", e)
|
||
return None
|
||
except Exception as e: # noqa: BLE001
|
||
logger.exception("本地脸型分类失败")
|
||
raise RuntimeError(f"本地脸型分类失败：{e}") from e
|
||
return feats
|
||
|
||
|
||
def has_face(features: dict) -> bool:
|
||
"""据 doubao 的「图片是否有人脸」判断。"""
|
||
v = features.get("图片是否有人脸") or features.get("是否有人") or ""
|
||
return "没人" not in str(v) and "没有" not in str(v)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import sys
|
||
path = sys.argv[1] if len(sys.argv) > 1 else "tests/fixtures/frontal.jpg"
|
||
with open(path, "rb") as f:
|
||
feats = analyze_features(image_bytes=f.read())
|
||
if feats is None:
|
||
print("无人脸（1001）")
|
||
else:
|
||
print(json.dumps(feats, ensure_ascii=False, indent=2))
|