xsl
ce445b64a6
部署修复： - torch.load 增加 weights_only=False patch，兼容 PyTorch 2.6+ 加载旧权重 - OSS 改为懒加载，本地用 output_format=base64 无需配凭证即可启动 - 补全被 gitignore 误排除的必需代码：core/models/layers/data、models/layers/data、keypoints/lib - webui 训练命令 --xformers 改 --sdpa（修复 xformers 无 CUDA 支持报错） 功能调整： - hair_grow_service 端口改 8899、preview 路由修复（send_file） - list_hairstyles 增加发型白名单，测试页只展示当前5个发型 新增脚本： - train_lora_parallel.py：直接调 kohya 并行训练 LoRA（绕过 photo_service 串行限制） - train_hairstyles_parallel.py / train_batch_stepC.py：批量训练辅助脚本 - scripts/sync_data_to_server.sh：大文件断点续传到云服务器 文档： - docs/换发型集成文档.md：换发型完整流程、服务架构、资源依赖、训练方法、集成步骤
165 lines
4.7 KiB
Python
165 lines
4.7 KiB
Python
# ------------------------------------------------------------------------------
|
|
# Copyright (c) Microsoft
|
|
# Licensed under the MIT License.
|
|
# Modified from py-faster-rcnn (https://github.com/rbgirshick/py-faster-rcnn)
|
|
# ------------------------------------------------------------------------------
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
import numpy as np
|
|
|
|
def py_nms_wrapper(thresh):
|
|
def _nms(dets):
|
|
return nms(dets, thresh)
|
|
return _nms
|
|
|
|
|
|
def nms(dets, thresh):
|
|
"""
|
|
greedily select boxes with high confidence and overlap with current maximum <= thresh
|
|
rule out overlap >= thresh
|
|
:param dets: [[x1, y1, x2, y2 score]]
|
|
:param thresh: retain overlap < thresh
|
|
:return: indexes to keep
|
|
"""
|
|
if dets.shape[0] == 0:
|
|
return []
|
|
|
|
x1 = dets[:, 0]
|
|
y1 = dets[:, 1]
|
|
x2 = dets[:, 2]
|
|
y2 = dets[:, 3]
|
|
scores = dets[:, 4]
|
|
|
|
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
|
|
order = scores.argsort()[::-1]
|
|
|
|
keep = []
|
|
while order.size > 0:
|
|
i = order[0]
|
|
keep.append(i)
|
|
xx1 = np.maximum(x1[i], x1[order[1:]])
|
|
yy1 = np.maximum(y1[i], y1[order[1:]])
|
|
xx2 = np.minimum(x2[i], x2[order[1:]])
|
|
yy2 = np.minimum(y2[i], y2[order[1:]])
|
|
|
|
w = np.maximum(0.0, xx2 - xx1 + 1)
|
|
h = np.maximum(0.0, yy2 - yy1 + 1)
|
|
inter = w * h
|
|
ovr = inter / (areas[i] + areas[order[1:]] - inter)
|
|
|
|
inds = np.where(ovr <= thresh)[0]
|
|
order = order[inds + 1]
|
|
|
|
return keep
|
|
|
|
|
|
def oks_iou(g, d, a_g, a_d, sigmas=None, in_vis_thre=None):
|
|
if not isinstance(sigmas, np.ndarray):
|
|
sigmas = np.array([.26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, 1.07, 1.07, .87, .87, .89, .89]) / 10.0
|
|
vars = (sigmas * 2) ** 2
|
|
xg = g[0::3]
|
|
yg = g[1::3]
|
|
vg = g[2::3]
|
|
ious = np.zeros((d.shape[0]))
|
|
for n_d in range(0, d.shape[0]):
|
|
xd = d[n_d, 0::3]
|
|
yd = d[n_d, 1::3]
|
|
vd = d[n_d, 2::3]
|
|
dx = xd - xg
|
|
dy = yd - yg
|
|
e = (dx ** 2 + dy ** 2) / vars / ((a_g + a_d[n_d]) / 2 + np.spacing(1)) / 2
|
|
if in_vis_thre is not None:
|
|
ind = list(vg > in_vis_thre) and list(vd > in_vis_thre)
|
|
e = e[ind]
|
|
ious[n_d] = np.sum(np.exp(-e)) / e.shape[0] if e.shape[0] != 0 else 0.0
|
|
return ious
|
|
|
|
|
|
def oks_nms(kpts_db, thresh, sigmas=None, in_vis_thre=None):
|
|
"""
|
|
greedily select boxes with high confidence and overlap with current maximum <= thresh
|
|
rule out overlap >= thresh, overlap = oks
|
|
:param kpts_db
|
|
:param thresh: retain overlap < thresh
|
|
:return: indexes to keep
|
|
"""
|
|
if len(kpts_db) == 0:
|
|
return []
|
|
|
|
scores = np.array([kpts_db[i]['score'] for i in range(len(kpts_db))])
|
|
kpts = np.array([kpts_db[i]['keypoints'].flatten() for i in range(len(kpts_db))])
|
|
areas = np.array([kpts_db[i]['area'] for i in range(len(kpts_db))])
|
|
|
|
order = scores.argsort()[::-1]
|
|
|
|
keep = []
|
|
while order.size > 0:
|
|
i = order[0]
|
|
keep.append(i)
|
|
|
|
oks_ovr = oks_iou(kpts[i], kpts[order[1:]], areas[i], areas[order[1:]], sigmas, in_vis_thre)
|
|
|
|
inds = np.where(oks_ovr <= thresh)[0]
|
|
order = order[inds + 1]
|
|
|
|
return keep
|
|
|
|
|
|
def rescore(overlap, scores, thresh, type='gaussian'):
|
|
assert overlap.shape[0] == scores.shape[0]
|
|
if type == 'linear':
|
|
inds = np.where(overlap >= thresh)[0]
|
|
scores[inds] = scores[inds] * (1 - overlap[inds])
|
|
else:
|
|
scores = scores * np.exp(- overlap**2 / thresh)
|
|
|
|
return scores
|
|
|
|
|
|
def soft_oks_nms(kpts_db, thresh, sigmas=None, in_vis_thre=None):
|
|
"""
|
|
greedily select boxes with high confidence and overlap with current maximum <= thresh
|
|
rule out overlap >= thresh, overlap = oks
|
|
:param kpts_db
|
|
:param thresh: retain overlap < thresh
|
|
:return: indexes to keep
|
|
"""
|
|
if len(kpts_db) == 0:
|
|
return []
|
|
|
|
scores = np.array([kpts_db[i]['score'] for i in range(len(kpts_db))])
|
|
kpts = np.array([kpts_db[i]['keypoints'].flatten() for i in range(len(kpts_db))])
|
|
areas = np.array([kpts_db[i]['area'] for i in range(len(kpts_db))])
|
|
|
|
order = scores.argsort()[::-1]
|
|
scores = scores[order]
|
|
|
|
# max_dets = order.size
|
|
max_dets = 20
|
|
keep = np.zeros(max_dets, dtype=np.intp)
|
|
keep_cnt = 0
|
|
while order.size > 0 and keep_cnt < max_dets:
|
|
i = order[0]
|
|
|
|
oks_ovr = oks_iou(kpts[i], kpts[order[1:]], areas[i], areas[order[1:]], sigmas, in_vis_thre)
|
|
|
|
order = order[1:]
|
|
scores = rescore(oks_ovr, scores[1:], thresh)
|
|
|
|
tmp = scores.argsort()[::-1]
|
|
order = order[tmp]
|
|
scores = scores[tmp]
|
|
|
|
keep[keep_cnt] = i
|
|
keep_cnt += 1
|
|
|
|
keep = keep[:keep_cnt]
|
|
|
|
return keep
|
|
# kpts_db = kpts_db[:keep_cnt]
|
|
|
|
# return kpts_db
|