完善部署并训练5个新发型 + 换发型集成文档
部署修复： - 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：换发型完整流程、服务架构、资源依赖、训练方法、集成步骤
This commit is contained in:
@@ -0,0 +1,121 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
|
||||
def flip_back(output_flipped, matched_parts):
|
||||
'''
|
||||
ouput_flipped: numpy.ndarray(batch_size, num_joints, height, width)
|
||||
'''
|
||||
assert output_flipped.ndim == 4,\
|
||||
'output_flipped should be [batch_size, num_joints, height, width]'
|
||||
|
||||
output_flipped = output_flipped[:, :, :, ::-1]
|
||||
|
||||
for pair in matched_parts:
|
||||
tmp = output_flipped[:, pair[0], :, :].copy()
|
||||
output_flipped[:, pair[0], :, :] = output_flipped[:, pair[1], :, :]
|
||||
output_flipped[:, pair[1], :, :] = tmp
|
||||
|
||||
return output_flipped
|
||||
|
||||
|
||||
def fliplr_joints(joints, joints_vis, width, matched_parts):
|
||||
"""
|
||||
flip coords
|
||||
"""
|
||||
# Flip horizontal
|
||||
joints[:, 0] = width - joints[:, 0] - 1
|
||||
|
||||
# Change left-right parts
|
||||
for pair in matched_parts:
|
||||
joints[pair[0], :], joints[pair[1], :] = \
|
||||
joints[pair[1], :], joints[pair[0], :].copy()
|
||||
joints_vis[pair[0], :], joints_vis[pair[1], :] = \
|
||||
joints_vis[pair[1], :], joints_vis[pair[0], :].copy()
|
||||
|
||||
return joints*joints_vis, joints_vis
|
||||
|
||||
|
||||
def transform_preds(coords, center, scale, output_size):
|
||||
target_coords = np.zeros(coords.shape)
|
||||
trans = get_affine_transform(center, scale, 0, output_size, inv=1)
|
||||
for p in range(coords.shape[0]):
|
||||
target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
|
||||
return target_coords
|
||||
|
||||
|
||||
def get_affine_transform(
|
||||
center, scale, rot, output_size,
|
||||
shift=np.array([0, 0], dtype=np.float32), inv=0
|
||||
):
|
||||
if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
|
||||
print(scale)
|
||||
scale = np.array([scale, scale])
|
||||
|
||||
scale_tmp = scale * 200.0
|
||||
src_w = scale_tmp[0]
|
||||
dst_w = output_size[0]
|
||||
dst_h = output_size[1]
|
||||
|
||||
rot_rad = np.pi * rot / 180
|
||||
src_dir = get_dir([0, src_w * -0.5], rot_rad)
|
||||
dst_dir = np.array([0, dst_w * -0.5], np.float32)
|
||||
|
||||
src = np.zeros((3, 2), dtype=np.float32)
|
||||
dst = np.zeros((3, 2), dtype=np.float32)
|
||||
src[0, :] = center + scale_tmp * shift
|
||||
src[1, :] = center + src_dir + scale_tmp * shift
|
||||
dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
|
||||
dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
|
||||
|
||||
src[2:, :] = get_3rd_point(src[0, :], src[1, :])
|
||||
dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
|
||||
|
||||
if inv:
|
||||
trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
|
||||
else:
|
||||
trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
|
||||
|
||||
return trans
|
||||
|
||||
|
||||
def affine_transform(pt, t):
|
||||
new_pt = np.array([pt[0], pt[1], 1.]).T
|
||||
new_pt = np.dot(t, new_pt)
|
||||
return new_pt[:2]
|
||||
|
||||
|
||||
def get_3rd_point(a, b):
|
||||
direct = a - b
|
||||
return b + np.array([-direct[1], direct[0]], dtype=np.float32)
|
||||
|
||||
|
||||
def get_dir(src_point, rot_rad):
|
||||
sn, cs = np.sin(rot_rad), np.cos(rot_rad)
|
||||
|
||||
src_result = [0, 0]
|
||||
src_result[0] = src_point[0] * cs - src_point[1] * sn
|
||||
src_result[1] = src_point[0] * sn + src_point[1] * cs
|
||||
|
||||
return src_result
|
||||
|
||||
|
||||
def crop(img, center, scale, output_size, rot=0):
|
||||
trans = get_affine_transform(center, scale, rot, output_size)
|
||||
|
||||
dst_img = cv2.warpAffine(
|
||||
img, trans, (int(output_size[0]), int(output_size[1])),
|
||||
flags=cv2.INTER_LINEAR
|
||||
)
|
||||
|
||||
return dst_img
|
||||
@@ -0,0 +1,203 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import os
|
||||
import logging
|
||||
import time
|
||||
from collections import namedtuple
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.optim as optim
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def create_logger(cfg, cfg_name, phase='train'):
|
||||
root_output_dir = Path(cfg.OUTPUT_DIR)
|
||||
# set up logger
|
||||
if not root_output_dir.exists():
|
||||
print('=> creating {}'.format(root_output_dir))
|
||||
root_output_dir.mkdir()
|
||||
|
||||
dataset = cfg.DATASET.DATASET + '_' + cfg.DATASET.HYBRID_JOINTS_TYPE \
|
||||
if cfg.DATASET.HYBRID_JOINTS_TYPE else cfg.DATASET.DATASET
|
||||
dataset = dataset.replace(':', '_')
|
||||
model = cfg.MODEL.NAME
|
||||
cfg_name = os.path.basename(cfg_name).split('.')[0]
|
||||
|
||||
final_output_dir = root_output_dir / dataset / model / cfg_name
|
||||
|
||||
print('=> creating {}'.format(final_output_dir))
|
||||
final_output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
time_str = time.strftime('%Y-%m-%d-%H-%M')
|
||||
log_file = '{}_{}_{}.log'.format(cfg_name, time_str, phase)
|
||||
final_log_file = final_output_dir / log_file
|
||||
head = '%(asctime)-15s %(message)s'
|
||||
logging.basicConfig(filename=str(final_log_file),
|
||||
format=head)
|
||||
logger = logging.getLogger()
|
||||
logger.setLevel(logging.INFO)
|
||||
console = logging.StreamHandler()
|
||||
logging.getLogger('').addHandler(console)
|
||||
|
||||
tensorboard_log_dir = Path(cfg.LOG_DIR) / dataset / model / \
|
||||
(cfg_name + '_' + time_str)
|
||||
|
||||
print('=> creating {}'.format(tensorboard_log_dir))
|
||||
tensorboard_log_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
return logger, str(final_output_dir), str(tensorboard_log_dir)
|
||||
|
||||
|
||||
def get_optimizer(cfg, model):
|
||||
optimizer = None
|
||||
if cfg.TRAIN.OPTIMIZER == 'sgd':
|
||||
optimizer = optim.SGD(
|
||||
model.parameters(),
|
||||
lr=cfg.TRAIN.LR,
|
||||
momentum=cfg.TRAIN.MOMENTUM,
|
||||
weight_decay=cfg.TRAIN.WD,
|
||||
nesterov=cfg.TRAIN.NESTEROV
|
||||
)
|
||||
elif cfg.TRAIN.OPTIMIZER == 'adam':
|
||||
optimizer = optim.Adam(
|
||||
model.parameters(),
|
||||
lr=cfg.TRAIN.LR
|
||||
)
|
||||
|
||||
return optimizer
|
||||
|
||||
|
||||
def save_checkpoint(states, is_best, output_dir,
|
||||
filename='checkpoint.pth'):
|
||||
torch.save(states, os.path.join(output_dir, filename))
|
||||
if is_best and 'state_dict' in states:
|
||||
torch.save(states['best_state_dict'],
|
||||
os.path.join(output_dir, 'model_best.pth'))
|
||||
|
||||
|
||||
def get_model_summary(model, *input_tensors, item_length=26, verbose=False):
|
||||
"""
|
||||
:param model:
|
||||
:param input_tensors:
|
||||
:param item_length:
|
||||
:return:
|
||||
"""
|
||||
|
||||
summary = []
|
||||
|
||||
ModuleDetails = namedtuple(
|
||||
"Layer", ["name", "input_size", "output_size", "num_parameters", "multiply_adds"])
|
||||
hooks = []
|
||||
layer_instances = {}
|
||||
|
||||
def add_hooks(module):
|
||||
|
||||
def hook(module, input, output):
|
||||
class_name = str(module.__class__.__name__)
|
||||
|
||||
instance_index = 1
|
||||
if class_name not in layer_instances:
|
||||
layer_instances[class_name] = instance_index
|
||||
else:
|
||||
instance_index = layer_instances[class_name] + 1
|
||||
layer_instances[class_name] = instance_index
|
||||
|
||||
layer_name = class_name + "_" + str(instance_index)
|
||||
|
||||
params = 0
|
||||
|
||||
if class_name.find("Conv") != -1 or class_name.find("BatchNorm") != -1 or \
|
||||
class_name.find("Linear") != -1:
|
||||
for param_ in module.parameters():
|
||||
params += param_.view(-1).size(0)
|
||||
|
||||
flops = "Not Available"
|
||||
if class_name.find("Conv") != -1 and hasattr(module, "weight"):
|
||||
flops = (
|
||||
torch.prod(
|
||||
torch.LongTensor(list(module.weight.data.size()))) *
|
||||
torch.prod(
|
||||
torch.LongTensor(list(output.size())[2:]))).item()
|
||||
elif isinstance(module, nn.Linear):
|
||||
flops = (torch.prod(torch.LongTensor(list(output.size()))) \
|
||||
* input[0].size(1)).item()
|
||||
|
||||
if isinstance(input[0], list):
|
||||
input = input[0]
|
||||
if isinstance(output, list):
|
||||
output = output[0]
|
||||
|
||||
summary.append(
|
||||
ModuleDetails(
|
||||
name=layer_name,
|
||||
input_size=list(input[0].size()),
|
||||
output_size=list(output.size()),
|
||||
num_parameters=params,
|
||||
multiply_adds=flops)
|
||||
)
|
||||
|
||||
if not isinstance(module, nn.ModuleList) \
|
||||
and not isinstance(module, nn.Sequential) \
|
||||
and module != model:
|
||||
hooks.append(module.register_forward_hook(hook))
|
||||
|
||||
model.eval()
|
||||
model.apply(add_hooks)
|
||||
|
||||
space_len = item_length
|
||||
|
||||
model(*input_tensors)
|
||||
for hook in hooks:
|
||||
hook.remove()
|
||||
|
||||
details = ''
|
||||
if verbose:
|
||||
details = "Model Summary" + \
|
||||
os.linesep + \
|
||||
"Name{}Input Size{}Output Size{}Parameters{}Multiply Adds (Flops){}".format(
|
||||
' ' * (space_len - len("Name")),
|
||||
' ' * (space_len - len("Input Size")),
|
||||
' ' * (space_len - len("Output Size")),
|
||||
' ' * (space_len - len("Parameters")),
|
||||
' ' * (space_len - len("Multiply Adds (Flops)"))) \
|
||||
+ os.linesep + '-' * space_len * 5 + os.linesep
|
||||
|
||||
params_sum = 0
|
||||
flops_sum = 0
|
||||
for layer in summary:
|
||||
params_sum += layer.num_parameters
|
||||
if layer.multiply_adds != "Not Available":
|
||||
flops_sum += layer.multiply_adds
|
||||
if verbose:
|
||||
details += "{}{}{}{}{}{}{}{}{}{}".format(
|
||||
layer.name,
|
||||
' ' * (space_len - len(layer.name)),
|
||||
layer.input_size,
|
||||
' ' * (space_len - len(str(layer.input_size))),
|
||||
layer.output_size,
|
||||
' ' * (space_len - len(str(layer.output_size))),
|
||||
layer.num_parameters,
|
||||
' ' * (space_len - len(str(layer.num_parameters))),
|
||||
layer.multiply_adds,
|
||||
' ' * (space_len - len(str(layer.multiply_adds)))) \
|
||||
+ os.linesep + '-' * space_len * 5 + os.linesep
|
||||
|
||||
details += os.linesep \
|
||||
+ "Total Parameters: {:,}".format(params_sum) \
|
||||
+ os.linesep + '-' * space_len * 5 + os.linesep
|
||||
details += "Total Multiply Adds (For Convolution and Linear Layers only): {:,} GFLOPs".format(flops_sum/(1024**3)) \
|
||||
+ os.linesep + '-' * space_len * 5 + os.linesep
|
||||
details += "Number of Layers" + os.linesep
|
||||
for layer in layer_instances:
|
||||
details += "{} : {} layers ".format(layer, layer_instances[layer])
|
||||
|
||||
return details
|
||||
@@ -0,0 +1,141 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import torchvision
|
||||
import cv2
|
||||
|
||||
from core.inference import get_max_preds
|
||||
|
||||
|
||||
def save_batch_image_with_joints(batch_image, batch_joints, batch_joints_vis,
|
||||
file_name, nrow=8, padding=2):
|
||||
'''
|
||||
batch_image: [batch_size, channel, height, width]
|
||||
batch_joints: [batch_size, num_joints, 3],
|
||||
batch_joints_vis: [batch_size, num_joints, 1],
|
||||
}
|
||||
'''
|
||||
grid = torchvision.utils.make_grid(batch_image, nrow, padding, True)
|
||||
ndarr = grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()
|
||||
ndarr = ndarr.copy()
|
||||
|
||||
nmaps = batch_image.size(0)
|
||||
xmaps = min(nrow, nmaps)
|
||||
ymaps = int(math.ceil(float(nmaps) / xmaps))
|
||||
height = int(batch_image.size(2) + padding)
|
||||
width = int(batch_image.size(3) + padding)
|
||||
k = 0
|
||||
for y in range(ymaps):
|
||||
for x in range(xmaps):
|
||||
if k >= nmaps:
|
||||
break
|
||||
joints = batch_joints[k]
|
||||
joints_vis = batch_joints_vis[k]
|
||||
|
||||
for joint, joint_vis in zip(joints, joints_vis):
|
||||
joint[0] = x * width + padding + joint[0]
|
||||
joint[1] = y * height + padding + joint[1]
|
||||
if joint_vis[0]:
|
||||
cv2.circle(ndarr, (int(joint[0]), int(joint[1])), 2, [255, 0, 0], 2)
|
||||
k = k + 1
|
||||
cv2.imwrite(file_name, ndarr)
|
||||
|
||||
|
||||
def save_batch_heatmaps(batch_image, batch_heatmaps, file_name,
|
||||
normalize=True):
|
||||
'''
|
||||
batch_image: [batch_size, channel, height, width]
|
||||
batch_heatmaps: ['batch_size, num_joints, height, width]
|
||||
file_name: saved file name
|
||||
'''
|
||||
if normalize:
|
||||
batch_image = batch_image.clone()
|
||||
min = float(batch_image.min())
|
||||
max = float(batch_image.max())
|
||||
|
||||
batch_image.add_(-min).div_(max - min + 1e-5)
|
||||
|
||||
batch_size = batch_heatmaps.size(0)
|
||||
num_joints = batch_heatmaps.size(1)
|
||||
heatmap_height = batch_heatmaps.size(2)
|
||||
heatmap_width = batch_heatmaps.size(3)
|
||||
|
||||
grid_image = np.zeros((batch_size*heatmap_height,
|
||||
(num_joints+1)*heatmap_width,
|
||||
3),
|
||||
dtype=np.uint8)
|
||||
|
||||
preds, maxvals = get_max_preds(batch_heatmaps.detach().cpu().numpy())
|
||||
|
||||
for i in range(batch_size):
|
||||
image = batch_image[i].mul(255)\
|
||||
.clamp(0, 255)\
|
||||
.byte()\
|
||||
.permute(1, 2, 0)\
|
||||
.cpu().numpy()
|
||||
heatmaps = batch_heatmaps[i].mul(255)\
|
||||
.clamp(0, 255)\
|
||||
.byte()\
|
||||
.cpu().numpy()
|
||||
|
||||
resized_image = cv2.resize(image,
|
||||
(int(heatmap_width), int(heatmap_height)))
|
||||
|
||||
height_begin = heatmap_height * i
|
||||
height_end = heatmap_height * (i + 1)
|
||||
for j in range(num_joints):
|
||||
cv2.circle(resized_image,
|
||||
(int(preds[i][j][0]), int(preds[i][j][1])),
|
||||
1, [0, 0, 255], 1)
|
||||
heatmap = heatmaps[j, :, :]
|
||||
colored_heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
|
||||
masked_image = colored_heatmap*0.7 + resized_image*0.3
|
||||
cv2.circle(masked_image,
|
||||
(int(preds[i][j][0]), int(preds[i][j][1])),
|
||||
1, [0, 0, 255], 1)
|
||||
|
||||
width_begin = heatmap_width * (j+1)
|
||||
width_end = heatmap_width * (j+2)
|
||||
grid_image[height_begin:height_end, width_begin:width_end, :] = \
|
||||
masked_image
|
||||
# grid_image[height_begin:height_end, width_begin:width_end, :] = \
|
||||
# colored_heatmap*0.7 + resized_image*0.3
|
||||
|
||||
grid_image[height_begin:height_end, 0:heatmap_width, :] = resized_image
|
||||
|
||||
cv2.imwrite(file_name, grid_image)
|
||||
|
||||
|
||||
def save_debug_images(config, input, meta, target, joints_pred, output,
|
||||
prefix):
|
||||
if not config.DEBUG.DEBUG:
|
||||
return
|
||||
|
||||
if config.DEBUG.SAVE_BATCH_IMAGES_GT:
|
||||
save_batch_image_with_joints(
|
||||
input, meta['joints'], meta['joints_vis'],
|
||||
'{}_gt.jpg'.format(prefix)
|
||||
)
|
||||
if config.DEBUG.SAVE_BATCH_IMAGES_PRED:
|
||||
save_batch_image_with_joints(
|
||||
input, joints_pred, meta['joints_vis'],
|
||||
'{}_pred.jpg'.format(prefix)
|
||||
)
|
||||
if config.DEBUG.SAVE_HEATMAPS_GT:
|
||||
save_batch_heatmaps(
|
||||
input, target, '{}_hm_gt.jpg'.format(prefix)
|
||||
)
|
||||
if config.DEBUG.SAVE_HEATMAPS_PRED:
|
||||
save_batch_heatmaps(
|
||||
input, output, '{}_hm_pred.jpg'.format(prefix)
|
||||
)
|
||||
@@ -0,0 +1,70 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import os
|
||||
import zipfile
|
||||
import xml.etree.ElementTree as ET
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
_im_zfile = []
|
||||
_xml_path_zip = []
|
||||
_xml_zfile = []
|
||||
|
||||
|
||||
def imread(filename, flags=cv2.IMREAD_COLOR):
|
||||
global _im_zfile
|
||||
path = filename
|
||||
pos_at = path.index('@')
|
||||
if pos_at == -1:
|
||||
print("character '@' is not found from the given path '%s'"%(path))
|
||||
assert 0
|
||||
path_zip = path[0: pos_at]
|
||||
path_img = path[pos_at + 2:]
|
||||
if not os.path.isfile(path_zip):
|
||||
print("zip file '%s' is not found"%(path_zip))
|
||||
assert 0
|
||||
for i in range(len(_im_zfile)):
|
||||
if _im_zfile[i]['path'] == path_zip:
|
||||
data = _im_zfile[i]['zipfile'].read(path_img)
|
||||
return cv2.imdecode(np.frombuffer(data, np.uint8), flags)
|
||||
|
||||
_im_zfile.append({
|
||||
'path': path_zip,
|
||||
'zipfile': zipfile.ZipFile(path_zip, 'r')
|
||||
})
|
||||
data = _im_zfile[-1]['zipfile'].read(path_img)
|
||||
|
||||
return cv2.imdecode(np.frombuffer(data, np.uint8), flags)
|
||||
|
||||
|
||||
def xmlread(filename):
|
||||
global _xml_path_zip
|
||||
global _xml_zfile
|
||||
path = filename
|
||||
pos_at = path.index('@')
|
||||
if pos_at == -1:
|
||||
print("character '@' is not found from the given path '%s'"%(path))
|
||||
assert 0
|
||||
path_zip = path[0: pos_at]
|
||||
path_xml = path[pos_at + 2:]
|
||||
if not os.path.isfile(path_zip):
|
||||
print("zip file '%s' is not found"%(path_zip))
|
||||
assert 0
|
||||
for i in xrange(len(_xml_path_zip)):
|
||||
if _xml_path_zip[i] == path_zip:
|
||||
data = _xml_zfile[i].open(path_xml)
|
||||
return ET.fromstring(data.read())
|
||||
_xml_path_zip.append(path_zip)
|
||||
print("read new xml file '%s'"%(path_zip))
|
||||
_xml_zfile.append(zipfile.ZipFile(path_zip, 'r'))
|
||||
data = _xml_zfile[-1].open(path_xml)
|
||||
return ET.fromstring(data.read())
|
||||
Reference in New Issue
Block a user