xsl
443cfa298f
包含： - 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密钥已脱敏为环境变量，原文件备份在本地
134 lines
3.7 KiB
Python
134 lines
3.7 KiB
Python
import numpy as np
|
|
import cv2
|
|
|
|
def nms(boxes, overlap_threshold=0.5, mode='union'):
|
|
""" Pure Python NMS baseline. """
|
|
x1 = boxes[:, 0]
|
|
y1 = boxes[:, 1]
|
|
x2 = boxes[:, 2]
|
|
y2 = boxes[:, 3]
|
|
scores = boxes[:, 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
|
|
|
|
if mode == 'min':
|
|
ovr = inter / np.minimum(areas[i], areas[order[1:]])
|
|
else:
|
|
ovr = inter / (areas[i] + areas[order[1:]] - inter)
|
|
|
|
inds = np.where(ovr <= overlap_threshold)[0]
|
|
order = order[inds + 1]
|
|
|
|
return keep
|
|
|
|
|
|
def convert_to_square(bboxes):
|
|
"""
|
|
Convert bounding boxes to a square form.
|
|
"""
|
|
square_bboxes = np.zeros_like(bboxes)
|
|
x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
|
|
h = y2 - y1 + 1.0
|
|
w = x2 - x1 + 1.0
|
|
max_side = np.maximum(h, w)
|
|
square_bboxes[:, 0] = x1 + w*0.5 - max_side*0.5
|
|
square_bboxes[:, 1] = y1 + h*0.5 - max_side*0.5
|
|
square_bboxes[:, 2] = square_bboxes[:, 0] + max_side - 1.0
|
|
square_bboxes[:, 3] = square_bboxes[:, 1] + max_side - 1.0
|
|
return square_bboxes
|
|
|
|
|
|
def calibrate_box(bboxes, offsets):
|
|
"""Transform bounding boxes to be more like true bounding boxes.
|
|
'offsets' is one of the outputs of the nets.
|
|
"""
|
|
x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
|
|
w = x2 - x1 + 1.0
|
|
h = y2 - y1 + 1.0
|
|
w = np.expand_dims(w, 1)
|
|
h = np.expand_dims(h, 1)
|
|
|
|
translation = np.hstack([w, h, w, h])*offsets
|
|
bboxes[:, 0:4] = bboxes[:, 0:4] + translation
|
|
return bboxes
|
|
|
|
|
|
def get_image_boxes(bounding_boxes, img, size=24):
|
|
"""Cut out boxes from the image.
|
|
"""
|
|
num_boxes = len(bounding_boxes)
|
|
(height, width, _) = img.shape
|
|
|
|
[dy, edy, dx, edx, y, ey, x, ex, w, h] = correct_bboxes(bounding_boxes, width, height)
|
|
img_boxes = np.zeros((num_boxes, 3, size, size), 'float32')
|
|
|
|
for i in range(num_boxes):
|
|
img_box = np.zeros((h[i], w[i], 3), 'uint8')
|
|
|
|
img_array = np.asarray(img, 'uint8')
|
|
img_box[dy[i]:(edy[i] + 1), dx[i]:(edx[i] + 1), :] =\
|
|
img_array[y[i]:(ey[i] + 1), x[i]:(ex[i] + 1), :]
|
|
|
|
img_box = cv2.resize(img_box, (size, size))
|
|
img_box = np.asarray(img_box, 'float32')
|
|
|
|
img_boxes[i, :, :, :] = _preprocess(img_box)
|
|
|
|
return img_boxes
|
|
|
|
|
|
def correct_bboxes(bboxes, width, height):
|
|
"""Crop boxes that are too big and get coordinates
|
|
with respect to cutouts.
|
|
"""
|
|
x1, y1, x2, y2 = [bboxes[:, i] for i in range(4)]
|
|
w, h = x2 - x1 + 1.0, y2 - y1 + 1.0
|
|
num_boxes = bboxes.shape[0]
|
|
|
|
x, y, ex, ey = x1, y1, x2, y2
|
|
dx, dy = np.zeros((num_boxes,)), np.zeros((num_boxes,))
|
|
edx, edy = w.copy() - 1.0, h.copy() - 1.0
|
|
|
|
ind = np.where(ex > width - 1.0)[0]
|
|
edx[ind] = w[ind] + width - 2.0 - ex[ind]
|
|
ex[ind] = width - 1.0
|
|
|
|
ind = np.where(ey > height - 1.0)[0]
|
|
edy[ind] = h[ind] + height - 2.0 - ey[ind]
|
|
ey[ind] = height - 1.0
|
|
|
|
ind = np.where(x < 0.0)[0]
|
|
dx[ind] = 0.0 - x[ind]
|
|
x[ind] = 0.0
|
|
|
|
ind = np.where(y < 0.0)[0]
|
|
dy[ind] = 0.0 - y[ind]
|
|
y[ind] = 0.0
|
|
return_list = [dy, edy, dx, edx, y, ey, x, ex, w, h]
|
|
return_list = [i.astype('int32') for i in return_list]
|
|
|
|
return return_list
|
|
|
|
|
|
def _preprocess(img):
|
|
"""Preprocessing step before feeding the network.
|
|
"""
|
|
img = img.transpose((2, 0, 1))
|
|
img = np.expand_dims(img, 0)
|
|
img = (img - 127.5)*0.0078125
|
|
return img
|