初始化：换发型/换发色/训练发型服务
包含： - hair_service_sd: 主服务（换发型/换发色/生发，端口8801） - photo_service: LoRA调度+训练（端口32678） - hair_grow_service: 调试测试页（端口8888，含4个测试页） - 批量训练脚本（batch_train_hairstyles.py） - 发际线mask自动识别（hairline_mask.py，4种方案） - 手绘mask换发型（hair_swap_manual.py） - 文档：README.md + LARGE_FILES.md + docs/ 大文件（模型权重200G、训练数据123G）已排除，见 LARGE_FILES.md OSS/COS密钥已脱敏为环境变量，原文件备份在本地
This commit is contained in:
xsl
@@ -0,0 +1,353 @@
|
||||
import math
|
||||
import numpy as np
|
||||
from .model import PNet, RNet, ONet
|
||||
from .box_utils import nms, calibrate_box, get_image_boxes, convert_to_square, _preprocess
|
||||
import torch
|
||||
import cv2
|
||||
from .nms.py_cpu_nms import py_cpu_nms
|
||||
from utils import box_utils_Retina
|
||||
from .layers.functions.prior_box import PriorBox
|
||||
from .config import cfg_re50
|
||||
from .retinaface import RetinaFace
|
||||
|
||||
|
||||
def detect_faces(image, min_face_size=20.0, thresholds=[0.6, 0.7, 0.8],
|
||||
nms_thresholds=[0.7, 0.7, 0.7], gpu_id=0):
|
||||
device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
pnet, rnet, onet = PNet(), RNet(), ONet()
|
||||
pnet.to(device)
|
||||
rnet.to(device)
|
||||
onet.to(device)
|
||||
onet.eval()
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
(height, width, _) = image.shape
|
||||
min_length = min(height, width)
|
||||
min_detection_size = 12
|
||||
factor = 0.707 # sqrt(0.5)
|
||||
|
||||
scales = []
|
||||
m = min_detection_size / min_face_size
|
||||
min_length *= m
|
||||
|
||||
factor_count = 0
|
||||
while min_length > min_detection_size:
|
||||
scales.append(m * factor ** factor_count)
|
||||
min_length *= factor
|
||||
factor_count += 1
|
||||
|
||||
# STAGE 1
|
||||
bounding_boxes = []
|
||||
for s in scales: # run P-Net on different scales
|
||||
boxes = run_first_stage(image, pnet, scale=s, threshold=thresholds[0], gpu_id=gpu_id)
|
||||
bounding_boxes.append(boxes)
|
||||
bounding_boxes = [i for i in bounding_boxes if i is not None]
|
||||
bounding_boxes = np.vstack(bounding_boxes)
|
||||
|
||||
keep = nms(bounding_boxes[:, 0:5], nms_thresholds[0])
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes = calibrate_box(bounding_boxes[:, 0:5], bounding_boxes[:, 5:])
|
||||
bounding_boxes = convert_to_square(bounding_boxes)
|
||||
bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
|
||||
|
||||
# STAGE 2
|
||||
img_boxes = get_image_boxes(bounding_boxes, image, size=24)
|
||||
img_boxes = torch.from_numpy(img_boxes)
|
||||
img_boxes = img_boxes.to(device)
|
||||
output = rnet(img_boxes)
|
||||
offsets = output[0].to('cpu').data.numpy() # shape [n_boxes, 4]
|
||||
probs = output[1].to('cpu').data.numpy() # shape [n_boxes, 2]
|
||||
|
||||
keep = np.where(probs[:, 1] > thresholds[1])[0]
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
|
||||
offsets = offsets[keep]
|
||||
|
||||
keep = nms(bounding_boxes, nms_thresholds[1])
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes = calibrate_box(bounding_boxes, offsets[keep])
|
||||
bounding_boxes = convert_to_square(bounding_boxes)
|
||||
bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
|
||||
|
||||
# STAGE 3
|
||||
img_boxes = get_image_boxes(bounding_boxes, image, size=48)
|
||||
if len(img_boxes) == 0:
|
||||
return [], []
|
||||
img_boxes = torch.from_numpy(img_boxes)
|
||||
img_boxes = img_boxes.to(device)
|
||||
output = onet(img_boxes)
|
||||
landmarks = output[0].to('cpu').data.numpy() # shape [n_boxes, 10]
|
||||
offsets = output[1].to('cpu').data.numpy() # shape [n_boxes, 4]
|
||||
probs = output[2].to('cpu').data.numpy() # shape [n_boxes, 2]
|
||||
|
||||
keep = np.where(probs[:, 1] > thresholds[2])[0]
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
|
||||
offsets = offsets[keep]
|
||||
landmarks = landmarks[keep]
|
||||
|
||||
# compute landmark points
|
||||
width = bounding_boxes[:, 2] - bounding_boxes[:, 0] + 1.0
|
||||
height = bounding_boxes[:, 3] - bounding_boxes[:, 1] + 1.0
|
||||
xmin, ymin = bounding_boxes[:, 0], bounding_boxes[:, 1]
|
||||
landmarks[:, 0:5] = np.expand_dims(xmin, 1) + np.expand_dims(width, 1) * landmarks[:, 0:5]
|
||||
landmarks[:, 5:10] = np.expand_dims(ymin, 1) + np.expand_dims(height, 1) * landmarks[:, 5:10]
|
||||
|
||||
bounding_boxes = calibrate_box(bounding_boxes, offsets)
|
||||
keep = nms(bounding_boxes, nms_thresholds[2], mode='min')
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
landmarks = landmarks[keep]
|
||||
|
||||
return bounding_boxes, landmarks
|
||||
|
||||
|
||||
class RetinaFaceDetector(object):
|
||||
def __init__(self, gpu_id=None):
|
||||
self.gpu_id = gpu_id
|
||||
self.cfg = cfg_re50
|
||||
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
model = RetinaFace(cfg=self.cfg, phase='test')
|
||||
model = self.load_model(model, './weights/Resnet50_Final.pth', True)
|
||||
model.eval()
|
||||
self.net = model.to(self.device)
|
||||
self.resize = 1
|
||||
self.confidence_threshold = 0.02
|
||||
self.top_k = 5000
|
||||
self.nms_threshold = 0.4
|
||||
self.keep_top_k = 750
|
||||
|
||||
def remove_prefix(self, state_dict, prefix):
|
||||
# print('remove prefix \'{}\''.format(prefix))
|
||||
f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else x
|
||||
return {f(key): value for key, value in state_dict.items()}
|
||||
|
||||
def check_keys(self, model, pretrained_state_dict):
|
||||
ckpt_keys = set(pretrained_state_dict.keys())
|
||||
model_keys = set(model.state_dict().keys())
|
||||
used_pretrained_keys = model_keys & ckpt_keys
|
||||
# unused_pretrained_keys = ckpt_keys - model_keys
|
||||
# missing_keys = model_keys - ckpt_keys
|
||||
assert len(used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint'
|
||||
return True
|
||||
|
||||
def load_model(self, model, pretrained_path, load_to_cpu):
|
||||
# print('Loading pretrained model from {}'.format(pretrained_path))
|
||||
if load_to_cpu:
|
||||
pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage)
|
||||
else:
|
||||
device = torch.cuda.current_device()
|
||||
pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage.cuda(device))
|
||||
if "state_dict" in pretrained_dict.keys():
|
||||
pretrained_dict = self.remove_prefix(pretrained_dict['state_dict'], 'module.')
|
||||
else:
|
||||
pretrained_dict = self.remove_prefix(pretrained_dict, 'module.')
|
||||
self.check_keys(model, pretrained_dict)
|
||||
model.load_state_dict(pretrained_dict, strict=False)
|
||||
return model
|
||||
|
||||
def forward(self, img_raw, min_face_size=50):
|
||||
img_scale = 640 / max(img_raw.shape[0], img_raw.shape[1])
|
||||
img = cv2.resize(img_raw, (0, 0), fx=img_scale, fy=img_scale)
|
||||
img = np.float32(img)
|
||||
|
||||
im_height, im_width, _ = img.shape
|
||||
scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
|
||||
img -= (104, 117, 123)
|
||||
img = img.transpose(2, 0, 1)
|
||||
img = torch.from_numpy(img).unsqueeze(0)
|
||||
img = img.to(self.device)
|
||||
scale = scale.to(self.device)
|
||||
|
||||
loc, conf, landms = self.net(img) # forward pass
|
||||
|
||||
priorbox = PriorBox(self.cfg, image_size=(im_height, im_width))
|
||||
priors = priorbox.forward()
|
||||
priors = priors.to(self.device)
|
||||
prior_data = priors.data
|
||||
boxes = box_utils_Retina.decode(loc.data.squeeze(0), prior_data, self.cfg['variance'])
|
||||
|
||||
boxes = boxes * scale / self.resize
|
||||
boxes = boxes.cpu().numpy()
|
||||
scores = conf.squeeze(0).data.cpu().numpy()[:, 1]
|
||||
landms = box_utils_Retina.decode_landm(landms.data.squeeze(0), prior_data, self.cfg['variance'])
|
||||
scale1 = torch.Tensor([img.shape[3], img.shape[2], img.shape[3], img.shape[2],
|
||||
img.shape[3], img.shape[2], img.shape[3], img.shape[2],
|
||||
img.shape[3], img.shape[2]])
|
||||
scale1 = scale1.to(self.device)
|
||||
landms = landms * scale1 / self.resize
|
||||
landms = landms.cpu().numpy()
|
||||
|
||||
# ignore low scores
|
||||
inds = np.where(scores > self.confidence_threshold)[0]
|
||||
boxes = boxes[inds]
|
||||
landms = landms[inds]
|
||||
scores = scores[inds]
|
||||
|
||||
# keep top-K before NMS
|
||||
order = scores.argsort()[::-1][:self.top_k]
|
||||
boxes = boxes[order]
|
||||
landms = landms[order]
|
||||
scores = scores[order]
|
||||
|
||||
# do NMS
|
||||
dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
|
||||
keep = py_cpu_nms(dets, self.nms_threshold, min_face_size = min_face_size * img_scale)
|
||||
# keep = nms(dets, args.nms_threshold,force_cpu=args.cpu)
|
||||
dets = dets[keep, :]
|
||||
landms = landms[keep]
|
||||
|
||||
# keep top-K faster NMS
|
||||
dets = dets[:self.keep_top_k, :]
|
||||
landms = landms[:self.keep_top_k, :]
|
||||
|
||||
dets[:, :4] = dets[:, :4] / img_scale
|
||||
landms /= img_scale
|
||||
return dets, landms
|
||||
|
||||
|
||||
class MTCNNFaceDetector(object):
|
||||
def __init__(self, gpu_id=None):
|
||||
self.gpu_id = gpu_id
|
||||
self.device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
self.pnet, self.rnet, self.onet = PNet(), RNet(), ONet()
|
||||
self.pnet.to(self.device)
|
||||
self.rnet.to(self.device)
|
||||
self.onet.to(self.device)
|
||||
self.onet.eval()
|
||||
|
||||
def forward(self, image, min_face_size=20.0, thresholds=[0.6, 0.7, 0.8], nms_thresholds=[0.7, 0.7, 0.7]):
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
(height, width, _) = image.shape
|
||||
min_length = min(height, width)
|
||||
min_detection_size = 12
|
||||
factor = 0.707 # sqrt(0.5)
|
||||
|
||||
scales = []
|
||||
m = min_detection_size / min_face_size
|
||||
min_length *= m
|
||||
|
||||
factor_count = 0
|
||||
while min_length > min_detection_size:
|
||||
scales.append(m * factor ** factor_count)
|
||||
min_length *= factor
|
||||
factor_count += 1
|
||||
|
||||
# STAGE 1
|
||||
bounding_boxes = []
|
||||
for s in scales: # run P-Net on different scales
|
||||
boxes = run_first_stage(image, self.pnet, scale=s, threshold=thresholds[0], gpu_id=self.gpu_id)
|
||||
bounding_boxes.append(boxes)
|
||||
bounding_boxes = [i for i in bounding_boxes if i is not None]
|
||||
if len(bounding_boxes) == 0:
|
||||
return [], []
|
||||
bounding_boxes = np.vstack(bounding_boxes)
|
||||
|
||||
keep = nms(bounding_boxes[:, 0:5], nms_thresholds[0])
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes = calibrate_box(bounding_boxes[:, 0:5], bounding_boxes[:, 5:])
|
||||
bounding_boxes = convert_to_square(bounding_boxes)
|
||||
bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
|
||||
|
||||
# STAGE 2
|
||||
img_boxes = get_image_boxes(bounding_boxes, image, size=24)
|
||||
img_boxes = torch.from_numpy(img_boxes)
|
||||
img_boxes = img_boxes.to(self.device)
|
||||
output = self.rnet(img_boxes)
|
||||
offsets = output[0].to('cpu').data.numpy() # shape [n_boxes, 4]
|
||||
probs = output[1].to('cpu').data.numpy() # shape [n_boxes, 2]
|
||||
|
||||
keep = np.where(probs[:, 1] > thresholds[1])[0]
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
|
||||
offsets = offsets[keep]
|
||||
|
||||
keep = nms(bounding_boxes, nms_thresholds[1])
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes = calibrate_box(bounding_boxes, offsets[keep])
|
||||
bounding_boxes = convert_to_square(bounding_boxes)
|
||||
bounding_boxes[:, 0:4] = np.round(bounding_boxes[:, 0:4])
|
||||
|
||||
# STAGE 3
|
||||
img_boxes = get_image_boxes(bounding_boxes, image, size=48)
|
||||
if len(img_boxes) == 0:
|
||||
return [], []
|
||||
img_boxes = torch.from_numpy(img_boxes)
|
||||
img_boxes = img_boxes.to(self.device)
|
||||
output = self.onet(img_boxes)
|
||||
landmarks = output[0].to('cpu').data.numpy() # shape [n_boxes, 10]
|
||||
offsets = output[1].to('cpu').data.numpy() # shape [n_boxes, 4]
|
||||
probs = output[2].to('cpu').data.numpy() # shape [n_boxes, 2]
|
||||
|
||||
keep = np.where(probs[:, 1] > thresholds[2])[0]
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
bounding_boxes[:, 4] = probs[keep, 1].reshape((-1,))
|
||||
offsets = offsets[keep]
|
||||
landmarks = landmarks[keep]
|
||||
|
||||
# compute landmark points
|
||||
width = bounding_boxes[:, 2] - bounding_boxes[:, 0] + 1.0
|
||||
height = bounding_boxes[:, 3] - bounding_boxes[:, 1] + 1.0
|
||||
xmin, ymin = bounding_boxes[:, 0], bounding_boxes[:, 1]
|
||||
landmarks[:, 0:5] = np.expand_dims(xmin, 1) + np.expand_dims(width, 1) * landmarks[:, 0:5]
|
||||
landmarks[:, 5:10] = np.expand_dims(ymin, 1) + np.expand_dims(height, 1) * landmarks[:, 5:10]
|
||||
|
||||
bounding_boxes = calibrate_box(bounding_boxes, offsets)
|
||||
keep = nms(bounding_boxes, nms_thresholds[2], mode='min')
|
||||
bounding_boxes = bounding_boxes[keep]
|
||||
landmarks = landmarks[keep]
|
||||
|
||||
return bounding_boxes, landmarks
|
||||
|
||||
|
||||
def run_first_stage(image, net, scale, threshold, gpu_id=0):
|
||||
"""
|
||||
Run P-Net, generate bounding boxes, and do NMS.
|
||||
"""
|
||||
device = torch.device('cuda:{}'.format(gpu_id) if gpu_id is not None else 'cpu')
|
||||
(height, width, _) = image.shape
|
||||
sw, sh = math.ceil(width * scale), math.ceil(height * scale)
|
||||
img = cv2.resize(image, (sw, sh))
|
||||
# img = image.resize((sw, sh), Image.BILINEAR)
|
||||
img = np.asarray(img, 'float32')
|
||||
img = torch.from_numpy(_preprocess(img))
|
||||
img = img.to(device)
|
||||
|
||||
output = net(img)
|
||||
probs = output[1].to('cpu').data.numpy()[0, 1, :, :]
|
||||
offsets = output[0].to('cpu').data.numpy()
|
||||
|
||||
boxes = _generate_bboxes(probs, offsets, scale, threshold)
|
||||
if len(boxes) == 0:
|
||||
return None
|
||||
|
||||
keep = nms(boxes[:, 0:5], overlap_threshold=0.5)
|
||||
return boxes[keep]
|
||||
|
||||
|
||||
def _generate_bboxes(probs, offsets, scale, threshold):
|
||||
"""
|
||||
Generate bounding boxes at places where there is probably a face.
|
||||
"""
|
||||
stride = 2
|
||||
cell_size = 12
|
||||
|
||||
inds = np.where(probs > threshold)
|
||||
|
||||
if inds[0].size == 0:
|
||||
return np.array([])
|
||||
|
||||
tx1, ty1, tx2, ty2 = [offsets[0, i, inds[0], inds[1]] for i in range(4)]
|
||||
|
||||
offsets = np.array([tx1, ty1, tx2, ty2])
|
||||
score = probs[inds[0], inds[1]]
|
||||
|
||||
# P-Net is applied to scaled images, so we need to rescale bounding boxes back
|
||||
bounding_boxes = np.vstack([
|
||||
np.round((stride * inds[1] + 1.0) / scale),
|
||||
np.round((stride * inds[0] + 1.0) / scale),
|
||||
np.round((stride * inds[1] + 1.0 + cell_size) / scale),
|
||||
np.round((stride * inds[0] + 1.0 + cell_size) / scale),
|
||||
score, offsets
|
||||
])
|
||||
|
||||
return bounding_boxes.T
|
||||
Reference in New Issue
Block a user