525 lines
21 KiB
Python
525 lines
21 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""区域生发服务（走换发型工作流，端口 8888）
|
||
|
||
工作流与换发型一致：infer_hairstyle_diy_jy → warpAffine → photo_service+LoRA → webui → 贴回。
|
||
mask = origin_matting ∪ new_matting ∪ 手绘mask（不减刘海）。
|
||
|
||
启动: python app.py
|
||
页面: http://<host>:8888
|
||
"""
|
||
import os
|
||
import sys
|
||
import time
|
||
import base64
|
||
import traceback
|
||
|
||
# 把 hair_service_sd 加入 path，使其模块可被导入
|
||
HAIR_SERVICE_DIR = "/home/ubuntu/change_hair/project/hair_service_sd"
|
||
sys.path.insert(0, HAIR_SERVICE_DIR)
|
||
# 切换到 hair_service_sd 目录，使 common.logger 能读到 config/configure.ini（相对路径）
|
||
os.chdir(HAIR_SERVICE_DIR)
|
||
os.environ.setdefault("HF_HUB_OFFLINE", "1")
|
||
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
|
||
os.environ.setdefault("CRYPTOGRAPHY_OPENSSL_NO_LEGACY", "1")
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from flask import Flask, request, jsonify, send_from_directory, send_file
|
||
from gevent import pywsgi
|
||
|
||
# PyTorch 2.6+ 默认 weights_only=True，无法加载含 numpy 对象的旧权重（yolov5l.pt 等）。
|
||
# 统一改回 False（权重均为本机自有可信文件）。
|
||
import torch
|
||
_orig_torch_load = torch.load
|
||
def _torch_load(*args, **kwargs):
|
||
kwargs.setdefault("weights_only", False)
|
||
return _orig_torch_load(*args, **kwargs)
|
||
torch.load = _torch_load
|
||
|
||
PORT = 8899
|
||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
|
||
app = Flask(__name__, static_folder="static", static_url_path="/static")
|
||
|
||
# ===== 模型懒加载（首次请求时初始化，避免启动卡死）=====
|
||
_hairstyle_process = None
|
||
_landmark_processor = None
|
||
|
||
|
||
def _get_models():
|
||
"""懒加载换发型全套模型"""
|
||
global _hairstyle_process, _landmark_processor
|
||
if _hairstyle_process is None:
|
||
print("[init] 加载换发型模型（首次，约30-90秒）...")
|
||
t0 = time.time()
|
||
import torch
|
||
from core.hairstyle_model import HairStyle_Model
|
||
from utils import landmark_processor
|
||
_hairstyle_process = HairStyle_Model(gpu=True, use_enhance=True)
|
||
_landmark_processor = landmark_processor
|
||
print(f"[init] 模型加载完成，耗时 {time.time()-t0:.1f}s")
|
||
return _hairstyle_process, _landmark_processor
|
||
|
||
|
||
@app.route("/")
|
||
def index():
|
||
return send_from_directory(os.path.join(BASE_DIR, "static"), "index.html")
|
||
|
||
|
||
@app.route("/swap")
|
||
def swap_page():
|
||
"""换发型测试页（独立，无遮罩）"""
|
||
return send_from_directory(os.path.join(BASE_DIR, "static"), "swap.html")
|
||
|
||
|
||
@app.route("/test_new")
|
||
def test_new_page():
|
||
"""新发型效果测试页（聚焦本次训练的30款发型）"""
|
||
return send_from_directory(os.path.join(BASE_DIR, "static"), "test_new.html")
|
||
|
||
|
||
@app.route("/debug")
|
||
def debug_page():
|
||
"""换发型调试页（全参数可视化）"""
|
||
return send_from_directory(os.path.join(BASE_DIR, "static"), "debug.html")
|
||
|
||
|
||
@app.route("/hairline")
|
||
def hairline_page():
|
||
"""发际线带重绘实验页"""
|
||
return send_from_directory(os.path.join(BASE_DIR, "static"), "hairline.html")
|
||
|
||
|
||
@app.route("/manual")
|
||
def manual_page():
|
||
"""手绘mask重绘测试页"""
|
||
return send_from_directory(os.path.join(BASE_DIR, "static"), "manual.html")
|
||
|
||
|
||
@app.route("/api/swap", methods=["POST"])
|
||
def api_swap_proxy():
|
||
"""代理换发型请求到 8801（避免前端跨域问题）"""
|
||
import requests as req
|
||
try:
|
||
data = request.json
|
||
resp = req.post("http://127.0.0.1:8801/api/swapHair/v1",
|
||
json=data, timeout=600)
|
||
return jsonify(resp.json())
|
||
except Exception as e:
|
||
return jsonify({"state": -1, "msg": f"换发型代理失败: {e}"}), 500
|
||
|
||
|
||
@app.route("/api/swap_viz", methods=["POST"])
|
||
def api_swap_viz():
|
||
"""换发型（带可视化中间产物）
|
||
|
||
入参 JSON:
|
||
img: 原图 base64
|
||
hair_id: 发型ID
|
||
is_hr: "true"/"false"
|
||
返回:
|
||
{state, result: 最终图b64, steps: [{title, desc, images:[{label, b64}]}]}
|
||
"""
|
||
t0 = time.time()
|
||
try:
|
||
data = request.json
|
||
img_b64 = data.get("img", "")
|
||
hair_id = data.get("hair_id", "")
|
||
is_hr = str(data.get("is_hr", "false")).lower() == "true"
|
||
strict_mask = data.get("strict_mask", False)
|
||
|
||
if not img_b64 or not hair_id:
|
||
return jsonify({"state": -1, "msg": "img 和 hair_id 不能为空"}), 400
|
||
|
||
img = _b64_to_ndarray(img_b64, color=True)
|
||
if img is None:
|
||
return jsonify({"state": -1, "msg": "img 解析失败"}), 400
|
||
|
||
print(f"[swap_viz] img={img.shape}, hair_id={hair_id}, is_hr={is_hr}, strict_mask={strict_mask}")
|
||
hairstyle_process, landmark_processor = _get_models()
|
||
|
||
from hair_swap_viz import hair_swap_viz
|
||
task_id = f"swapviz_{int(time.time()*1000)}"
|
||
result, steps = hair_swap_viz(
|
||
origin_img=img, hair_id=hair_id,
|
||
hairstyle_process=hairstyle_process,
|
||
landmark_processor=landmark_processor,
|
||
task_id=task_id, is_hr=is_hr, strict_mask=strict_mask)
|
||
|
||
_, buf = cv2.imencode(".jpg", result, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||
result_b64 = base64.b64encode(buf).decode("utf-8")
|
||
print(f"[swap_viz] 完成，总耗时 {time.time()-t0:.1f}s，{len(steps)}个步骤")
|
||
return jsonify({"state": 0, "result": result_b64, "steps": steps})
|
||
|
||
except Exception as e:
|
||
print(f"[swap_viz] 失败: {e}")
|
||
traceback.print_exc()
|
||
return jsonify({"state": -1, "msg": f"换发型失败: {e}"}), 500
|
||
|
||
|
||
@app.route("/api/swap_debug", methods=["POST"])
|
||
def api_swap_debug():
|
||
"""换发型调试接口（全参数可调 + 每步可视化）。
|
||
|
||
入参 JSON:
|
||
img: 原图 base64
|
||
hair_id: 发型ID
|
||
# 流程开关
|
||
cut_bang: bool, 是否减刘海圆（默认 true）
|
||
strict_mask: bool, 严格按mask贴回（默认 false）
|
||
seamless_blend: bool, 泊松融合消除接缝（默认 true，仅 strict_mask 时生效）
|
||
# 尺寸/对齐
|
||
is_hr: bool, 高清模式（默认 true）
|
||
dilate_kernel: [x,y], mask膨胀核（默认 [6,18]）
|
||
# SD 推理（仅 denoising 可调）
|
||
denoising_strength: float 0~1, 重绘强度（默认 0.6）
|
||
# 贴回/融合
|
||
blend_dilate: [x,y], strict贴回mask膨胀（默认 [5,5]）
|
||
seamless_dilate: [x,y], 泊松融合mask膨胀（默认 [9,9]）
|
||
返回:
|
||
{state, result, steps, params}
|
||
"""
|
||
t0 = time.time()
|
||
try:
|
||
d = request.json
|
||
img_b64 = d.get("img", "")
|
||
hair_id = d.get("hair_id", "")
|
||
if not img_b64 or not hair_id:
|
||
return jsonify({"state": -1, "msg": "img 和 hair_id 不能为空"}), 400
|
||
|
||
img = _b64_to_ndarray(img_b64, color=True)
|
||
if img is None:
|
||
return jsonify({"state": -1, "msg": "img 解析失败"}), 400
|
||
|
||
# 解析参数（带默认值 + 类型转换）
|
||
def g(k, default): return d.get(k, default)
|
||
dk = g("dilate_kernel", [6, 18])
|
||
bd = g("blend_dilate", [5, 5])
|
||
sd = g("seamless_dilate", [9, 9])
|
||
|
||
print(f"[swap_debug] hair_id={hair_id}, cut_bang={g('cut_bang',True)}, "
|
||
f"strict_mask={g('strict_mask',False)}, seamless_blend={g('seamless_blend',True)}, "
|
||
f"is_hr={g('is_hr',True)}, denoising={g('denoising_strength',0.6)}")
|
||
|
||
hairstyle_process, landmark_processor = _get_models()
|
||
from hair_swap_debug import hair_swap_debug
|
||
task_id = f"swapdbg_{int(time.time()*1000)}"
|
||
result, steps, params = hair_swap_debug(
|
||
origin_img=img, hair_id=hair_id,
|
||
hairstyle_process=hairstyle_process,
|
||
landmark_processor=landmark_processor,
|
||
task_id=task_id,
|
||
cut_bang=bool(g("cut_bang", True)),
|
||
strict_mask=bool(g("strict_mask", False)),
|
||
seamless_blend=bool(g("seamless_blend", True)),
|
||
is_hr=bool(g("is_hr", True)),
|
||
dilate_kernel=(int(dk[0]), int(dk[1])),
|
||
denoising_strength=float(g("denoising_strength", 0.6)),
|
||
blend_dilate=(int(bd[0]), int(bd[1])),
|
||
seamless_dilate=(int(sd[0]), int(sd[1])),
|
||
)
|
||
|
||
_, buf = cv2.imencode(".jpg", result, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||
result_b64 = base64.b64encode(buf).decode("utf-8")
|
||
print(f"[swap_debug] 完成，总耗时 {time.time()-t0:.1f}s")
|
||
return jsonify({"state": 0, "result": result_b64, "steps": steps, "params": params})
|
||
|
||
except Exception as e:
|
||
print(f"[swap_debug] 失败: {e}")
|
||
traceback.print_exc()
|
||
return jsonify({"state": -1, "msg": f"换发型失败: {e}"}), 500
|
||
|
||
|
||
@app.route("/api/swap_hairline", methods=["POST"])
|
||
def api_swap_hairline():
|
||
"""发际线带重绘接口（实验性，流程同 swap_debug，仅步骤③改为边界带）。
|
||
|
||
比 swap_debug 多一个参数：
|
||
method: str, 发际线mask方案，可选 boundary_band/mediapipe/landmark_1k/deeplab（默认 mediapipe）
|
||
band_width: int, 边界带宽度（形态学核大小，默认15，带宽≈2*band_width）
|
||
其余参数同 swap_debug。
|
||
"""
|
||
t0 = time.time()
|
||
try:
|
||
d = request.json
|
||
img_b64 = d.get("img", "")
|
||
hair_id = d.get("hair_id", "")
|
||
if not img_b64 or not hair_id:
|
||
return jsonify({"state": -1, "msg": "img 和 hair_id 不能为空"}), 400
|
||
|
||
img = _b64_to_ndarray(img_b64, color=True)
|
||
if img is None:
|
||
return jsonify({"state": -1, "msg": "img 解析失败"}), 400
|
||
|
||
def g(k, default): return d.get(k, default)
|
||
dk = g("dilate_kernel", [6, 18])
|
||
bd = g("blend_dilate", [5, 5])
|
||
sd = g("seamless_dilate", [9, 9])
|
||
|
||
print(f"[swap_hairline] hair_id={hair_id}, method={g('method','mediapipe')}, "
|
||
f"band_width={g('band_width',15)}, denoising={g('denoising_strength',0.6)}")
|
||
|
||
hairstyle_process, landmark_processor = _get_models()
|
||
from hair_swap_hairline import hair_swap_hairline
|
||
task_id = f"hairline_{int(time.time()*1000)}"
|
||
result, steps, params = hair_swap_hairline(
|
||
origin_img=img, hair_id=hair_id,
|
||
hairstyle_process=hairstyle_process,
|
||
landmark_processor=landmark_processor,
|
||
task_id=task_id,
|
||
method=str(g("method", "mediapipe")),
|
||
strict_mask=bool(g("strict_mask", False)),
|
||
seamless_blend=bool(g("seamless_blend", True)),
|
||
is_hr=bool(g("is_hr", True)),
|
||
dilate_kernel=(int(dk[0]), int(dk[1])),
|
||
denoising_strength=float(g("denoising_strength", 0.6)),
|
||
blend_dilate=(int(bd[0]), int(bd[1])),
|
||
seamless_dilate=(int(sd[0]), int(sd[1])),
|
||
band_width=int(g("band_width", 15)),
|
||
preview_only=bool(g("preview_only", False)),
|
||
height_ratio=float(g("height_ratio", 0.432)),
|
||
width_ratio=float(g("width_ratio", 0.144)),
|
||
corner_ratio=float(g("corner_ratio", 0.25)),
|
||
vertical_offset=float(g("vertical_offset", 0.0)),
|
||
refiner_switch_at=float(g("refiner_switch_at", 0.5)),
|
||
)
|
||
|
||
_, buf = cv2.imencode(".jpg", result, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||
result_b64 = base64.b64encode(buf).decode("utf-8")
|
||
print(f"[swap_hairline] 完成，总耗时 {time.time()-t0:.1f}s")
|
||
return jsonify({"state": 0, "result": result_b64, "steps": steps, "params": params})
|
||
|
||
except Exception as e:
|
||
print(f"[swap_hairline] 失败: {e}")
|
||
traceback.print_exc()
|
||
return jsonify({"state": -1, "msg": f"发际线带重绘失败: {e}"}), 500
|
||
|
||
|
||
@app.route("/api/swap_manual", methods=["POST"])
|
||
def api_swap_manual():
|
||
"""手绘mask版换发型接口（重绘区域完全由用户手绘mask决定）。
|
||
|
||
入参 JSON:
|
||
img: 原图 base64
|
||
mask: 用户手绘mask base64（白色=重绘区，与img同分辨率）
|
||
hair_id: 发型ID
|
||
其余参数同 swap_debug（strict_mask/seamless_blend/is_hr/dilate_kernel/
|
||
denoising_strength/blend_dilate/seamless_dilate）
|
||
返回: {state, result, steps, params}
|
||
"""
|
||
t0 = time.time()
|
||
try:
|
||
d = request.json
|
||
img_b64 = d.get("img", "")
|
||
mask_b64 = d.get("mask", "")
|
||
hair_id = d.get("hair_id", "")
|
||
if not img_b64 or not mask_b64:
|
||
return jsonify({"state": -1, "msg": "img 和 mask 不能为空（请先手绘重绘区域）"}), 400
|
||
if not hair_id:
|
||
return jsonify({"state": -1, "msg": "hair_id 不能为空"}), 400
|
||
|
||
img = _b64_to_ndarray(img_b64, color=True)
|
||
hand_mask = _b64_to_ndarray(mask_b64, color=False)
|
||
if img is None:
|
||
return jsonify({"state": -1, "msg": "img 解析失败"}), 400
|
||
if hand_mask is None:
|
||
return jsonify({"state": -1, "msg": "mask 解析失败"}), 400
|
||
|
||
def g(k, default): return d.get(k, default)
|
||
dk = g("dilate_kernel", [6, 18])
|
||
bd = g("blend_dilate", [5, 5])
|
||
sd = g("seamless_dilate", [9, 9])
|
||
|
||
print(f"[swap_manual] hair_id={hair_id}, img={img.shape}, mask={hand_mask.shape}, "
|
||
f"strict_mask={g('strict_mask',False)}, denoising={g('denoising_strength',0.6)}")
|
||
|
||
hairstyle_process, landmark_processor = _get_models()
|
||
from hair_swap_manual import hair_swap_manual
|
||
task_id = f"manual_{int(time.time()*1000)}"
|
||
result, steps, params = hair_swap_manual(
|
||
origin_img=img, hand_mask=hand_mask, hair_id=hair_id,
|
||
hairstyle_process=hairstyle_process,
|
||
landmark_processor=landmark_processor,
|
||
task_id=task_id,
|
||
strict_mask=bool(g("strict_mask", False)),
|
||
seamless_blend=bool(g("seamless_blend", True)),
|
||
is_hr=bool(g("is_hr", True)),
|
||
dilate_kernel=(int(dk[0]), int(dk[1])),
|
||
denoising_strength=float(g("denoising_strength", 0.6)),
|
||
blend_dilate=(int(bd[0]), int(bd[1])),
|
||
seamless_dilate=(int(sd[0]), int(sd[1])),
|
||
feather_px=int(g("feather_px", 0)),
|
||
enhance=bool(g("enhance", False)),
|
||
enhance_denoising=float(g("enhance_denoising", 0.35)),
|
||
)
|
||
|
||
_, buf = cv2.imencode(".jpg", result, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||
result_b64 = base64.b64encode(buf).decode("utf-8")
|
||
print(f"[swap_manual] 完成，总耗时 {time.time()-t0:.1f}s")
|
||
return jsonify({"state": 0, "result": result_b64, "steps": steps, "params": params})
|
||
|
||
except Exception as e:
|
||
print(f"[swap_manual] 失败: {e}")
|
||
traceback.print_exc()
|
||
return jsonify({"state": -1, "msg": f"手绘mask换发型失败: {e}"}), 500
|
||
|
||
|
||
def _b64_to_ndarray(b64_str, color=True):
|
||
if "," in b64_str and b64_str.startswith("data:"):
|
||
b64_str = b64_str.split(",", 1)[1]
|
||
data = base64.b64decode(b64_str)
|
||
flag = cv2.IMREAD_COLOR if color else cv2.IMREAD_GRAYSCALE
|
||
return cv2.imdecode(np.frombuffer(data, np.uint8), flag)
|
||
|
||
|
||
@app.route("/api/hairstyles")
|
||
def api_hairstyles():
|
||
"""返回可用发型列表（含性别，用于前端下拉框分组）"""
|
||
try:
|
||
from common.logger import config
|
||
hairstyle_dir = config.get('default', 'hairstyleDir')
|
||
train_dir = config.get('default', 'train_dir')
|
||
upload_dir = config.get('default', 'upload_train_dir')
|
||
from hair_grow_swap import list_hairstyles
|
||
styles = list_hairstyles(hairstyle_dir, train_dir, upload_dir)
|
||
boy = [s for s in styles if s["gender"] == "boy"]
|
||
girl = [s for s in styles if s["gender"] == "girl"]
|
||
print(f"[hairstyles] 共 {len(styles)} 个可用发型 (boy={len(boy)}, girl={len(girl)})")
|
||
return jsonify({"state": 0, "data": styles, "count": len(styles)})
|
||
except Exception as e:
|
||
traceback.print_exc()
|
||
return jsonify({"state": -1, "msg": str(e)}), 500
|
||
|
||
|
||
@app.route("/preview/<hair_id>")
|
||
def preview_img(hair_id):
|
||
"""直接返回预览图文件（比 base64 API 快，浏览器可缓存）"""
|
||
effect_path = os.path.join(BASE_DIR, "static", "previews", f"{hair_id}.jpg")
|
||
if os.path.exists(effect_path):
|
||
return send_from_directory(os.path.join(BASE_DIR, "static", "previews"), f"{hair_id}.jpg")
|
||
# 回退到 ref_rgb
|
||
try:
|
||
from common.logger import config
|
||
hairstyle_dir = config.get('default', 'hairstyleDir')
|
||
fallback = os.path.join(hairstyle_dir, hair_id, "ref_rgb_8uc3_768.png")
|
||
if os.path.exists(fallback):
|
||
return send_file(fallback)
|
||
except Exception:
|
||
pass
|
||
return ("", 404)
|
||
|
||
|
||
@app.route("/train_src/<hair_id>")
|
||
def train_src_img(hair_id):
|
||
"""返回发型的训练原图（hair_type_images/<hair_id>.jpg/.png）。
|
||
用于测试页展示发型真实样子，而非套在标准脸上的效果图。
|
||
"""
|
||
src_dir = "/home/ubuntu/change_hair/hair_type_images"
|
||
for ext in (".jpg", ".jpeg", ".png"):
|
||
path = os.path.join(src_dir, hair_id + ext)
|
||
if os.path.exists(path):
|
||
return send_from_directory(src_dir, hair_id + ext)
|
||
return ("", 404)
|
||
|
||
|
||
@app.route("/api/hairstyle_preview/<hair_id>")
|
||
def api_hairstyle_preview(hair_id):
|
||
"""返回某发型的预览图。
|
||
优先用生成的效果图 static/previews/<hair_id>.jpg（发型套在标准脸上的样子），
|
||
没有则回退到 first## 原始上传图。
|
||
"""
|
||
try:
|
||
# 1. 优先：生成的效果图（发型套在 boy/girl 标准脸上）
|
||
preview = None
|
||
effect_path = os.path.join(BASE_DIR, "static", "previews", f"{hair_id}.jpg")
|
||
if os.path.exists(effect_path):
|
||
preview = effect_path
|
||
else:
|
||
# 2. 回退：first## 原始上传图
|
||
from common.logger import config
|
||
upload_dir = config.get('default', 'upload_train_dir')
|
||
save_dir = os.path.join(upload_dir, hair_id)
|
||
if os.path.isdir(save_dir):
|
||
for name in os.listdir(save_dir):
|
||
if "first##" in name:
|
||
preview = os.path.join(save_dir, name)
|
||
break
|
||
# 3. 再回退：ref_rgb
|
||
if not preview or not os.path.exists(preview):
|
||
hairstyle_dir = config.get('default', 'hairstyleDir')
|
||
fallback = os.path.join(hairstyle_dir, hair_id, "ref_rgb_8uc3_768.png")
|
||
if os.path.exists(fallback):
|
||
preview = fallback
|
||
if not preview:
|
||
return jsonify({"state": -1, "msg": "无预览图"}), 404
|
||
img = cv2.imread(preview)
|
||
img = cv2.resize(img, (256, 256), interpolation=cv2.INTER_AREA)
|
||
_, buf = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, 85])
|
||
b64 = base64.b64encode(buf).decode()
|
||
return jsonify({"state": 0, "preview": b64})
|
||
except Exception as e:
|
||
return jsonify({"state": -1, "msg": str(e)}), 500
|
||
|
||
|
||
@app.route("/api/grow", methods=["POST"])
|
||
def api_grow():
|
||
"""生发接口（走换发型工作流）
|
||
|
||
入参 JSON:
|
||
img: 原图 base64
|
||
mask: 手绘遮罩 base64（白=生发区，与img同分辨率）
|
||
hair_id: 选择的发型ID（必填）
|
||
is_hr: 是否高清 "true"/"false"，默认 "true"
|
||
"""
|
||
t0 = time.time()
|
||
try:
|
||
data = request.json
|
||
img_b64 = data.get("img", "")
|
||
mask_b64 = data.get("mask", "")
|
||
hair_id = data.get("hair_id", "")
|
||
is_hr = str(data.get("is_hr", "true")).lower() == "true"
|
||
|
||
if not img_b64 or not mask_b64:
|
||
return jsonify({"state": -1, "msg": "img 和 mask 不能为空"}), 400
|
||
if not hair_id:
|
||
return jsonify({"state": -1, "msg": "hair_id 不能为空（请先选择发型）"}), 400
|
||
|
||
img = _b64_to_ndarray(img_b64, color=True)
|
||
mask = _b64_to_ndarray(mask_b64, color=False)
|
||
if img is None:
|
||
return jsonify({"state": -1, "msg": "img 解析失败"}), 400
|
||
if mask is None:
|
||
return jsonify({"state": -1, "msg": "mask 解析失败"}), 400
|
||
|
||
print(f"[grow] img={img.shape}, mask={mask.shape}, hair_id={hair_id}, is_hr={is_hr}")
|
||
|
||
# 加载模型（首次慢）
|
||
hairstyle_process, landmark_processor = _get_models()
|
||
|
||
# 调用生发（走换发型工作流）
|
||
from hair_grow_swap import hair_grow_swap
|
||
task_id = f"grow_{int(time.time()*1000)}"
|
||
result = hair_grow_swap(
|
||
origin_img=img, hand_mask=mask, hair_id=hair_id,
|
||
hairstyle_process=hairstyle_process,
|
||
landmark_processor=landmark_processor,
|
||
task_id=task_id, is_hr=is_hr)
|
||
|
||
_, buf = cv2.imencode(".jpg", result, [cv2.IMWRITE_JPEG_QUALITY, 95])
|
||
result_b64 = base64.b64encode(buf).decode("utf-8")
|
||
print(f"[grow] 完成，总耗时 {time.time()-t0:.1f}s")
|
||
return jsonify({"state": 0, "result": result_b64})
|
||
|
||
except Exception as e:
|
||
print(f"[grow] 失败: {e}")
|
||
traceback.print_exc()
|
||
return jsonify({"state": -1, "msg": f"生发失败: {e}"}), 500
|
||
|
||
|
||
if __name__ == "__main__":
|
||
print(f"生发服务启动，端口 {PORT}")
|
||
print(f"测试页面: http://0.0.0.0:{PORT}")
|
||
print(f"注意：首次请求会加载换发型模型（约30-60秒）")
|
||
server = pywsgi.WSGIServer(("0.0.0.0", PORT), app)
|
||
server.serve_forever()
|