def generate(self):
model_path = os.path.expanduser(self.model_path)
assert model_path.endswith('.h5'), 'Keras model or weights must be a .h5 file.'
# 载入模型
self.retinaface = RetinaFace(self.cfg, self.backbone)
self.retinaface.load_weights(self.model_path,by_name=True)
def generate(self):
model_path = os.path.expanduser(self.model_path)
assert model_path.endswith(
'.h5'), 'Keras model or weights must be a .h5 file.'
self.retinaface = RetinaFace(self.cfg, self.backbone)
self.retinaface.load_weights(self.model_path, by_name=True)
print('{} model, anchors, and classes loaded.'.format(model_path))
def generate(self):
model_path = os.path.expanduser(self.model_path)
assert model_path.endswith('.h5'), 'tensorflow.keras model or weights must be a .h5 file.'
# 加快模型训练的效率
print('Loading weights into state dict...')
# 载入模型
self.retinaface = RetinaFace(self.cfg, self.backbone)
self.retinaface.load_weights(self.model_path)
print('{} model, anchors loaded.'.format(self.model_path))
def generate(self):
model_path = os.path.expanduser(self.model_path)
assert model_path.endswith(
'.h5'), 'tensorflow.keras model or weights must be a .h5 file.'
#-------------------------------#
# 载入模型与权值
#-------------------------------#
self.retinaface = RetinaFace(self.cfg, self.backbone)
self.retinaface.load_weights(self.model_path)
print('{} model, anchors loaded.'.format(self.model_path))
def generate(self):
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
self.net = RetinaFace(cfg=self.cfg, phase='eval').eval()
# 加快模型训练的效率
print('Loading weights into state dict...')
state_dict = torch.load(self.model_path)
self.net.load_state_dict(state_dict)
if self.cuda:
self.net = nn.DataParallel(self.net)
self.net = self.net.cuda()
print('Finished!')
class Retinaface(object):
_defaults = {
"model_path":
'model_data/Epoch114-Total_Loss0.3237-Conf_Loss.pth0.0999-Regression_Loss0.2241.pth',
"backbone": 'resnet50',
"confidence": 0.9,
"nms_iou": 0.45,
"cuda": True,
#----------------------------------------------------------------------#
# 是否需要进行图像大小限制。
# 开启后，会将输入图像的大小限制为input_shape。否则使用原图进行预测。
# 可根据输入图像的大小自行调整input_shape，注意为32的倍数，如[640, 640, 3]
#----------------------------------------------------------------------#
"input_shape": [768, 768, 3],
"letterbox_image": True
}
@classmethod
def get_defaults(cls, n):
if n in cls._defaults:
return cls._defaults[n]
else:
return "Unrecognized attribute name '" + n + "'"
#---------------------------------------------------#
# 初始化Retinaface
#---------------------------------------------------#
def __init__(self, **kwargs):
self.confidence = None
self.__dict__.update(self._defaults)
if self.backbone == "mobilenet":
self.cfg = cfg_mnet
else:
self.cfg = cfg_re50
self.generate()
if self.letterbox_image:
self.anchors = Anchors(
self.cfg,
image_size=[self.input_shape[0],
self.input_shape[1]]).get_anchors()
#---------------------------------------------------#
# 载入模型
#---------------------------------------------------#
def generate(self):
self.net = RetinaFace(cfg=self.cfg, mode='eval').eval()
#-------------------------------#
# 载入模型与权值
#-------------------------------#
print('Loading weights into state dict...')
state_dict = torch.load(self.model_path)
self.net.load_state_dict(state_dict)
if self.cuda:
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
self.net = nn.DataParallel(self.net)
self.net = self.net.cuda()
print('Finished!')
#---------------------------------------------------#
# 检测图片
#---------------------------------------------------#
def detect_image(self, image):
#---------------------------------------------------#
# 对输入图像进行一个备份，后面用于绘图
#---------------------------------------------------#
old_image = image.copy()
image = np.array(image, np.float32)
#---------------------------------------------------#
# 计算scale，用于将获得的预测框转换成原图的高宽
#---------------------------------------------------#
scale = [
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0]
]
scale_for_landmarks = [
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0]
]
im_height, im_width, _ = np.shape(image)
#---------------------------------------------------------#
# letterbox_image可以给图像增加灰条，实现不失真的resize
#---------------------------------------------------------#
if self.letterbox_image:
image = np.array(
letterbox_image(image,
[self.input_shape[1], self.input_shape[0]]),
np.float32)
else:
self.anchors = Anchors(self.cfg,
image_size=(im_height,
im_width)).get_anchors()
with torch.no_grad():
#-----------------------------------------------------------#
# 图片预处理，归一化。
#-----------------------------------------------------------#
image = torch.from_numpy(
preprocess_input(image).transpose(2, 0, 1)).unsqueeze(0)
if self.cuda:
self.anchors = self.anchors.cuda()
image = image.cuda()
loc, conf = self.net(image)
#-----------------------------------------------------------#
# 将预测结果进行解码
#-----------------------------------------------------------#
boxes = decode(loc.data.squeeze(0), self.anchors,
self.cfg['variance'])
boxes = boxes.cpu().numpy()
conf = conf.data.squeeze(0)[:, 1:2].cpu().numpy()
boxes_conf_landms = np.concatenate([boxes, conf], -1)
print(boxes_conf_landms.shape)
# boxes_conf_landms = non_max_suppression(boxes_conf_landms, 0.9)
mask = boxes_conf_landms[:, 4] >= 0.15
boxes_conf_landms = boxes_conf_landms[mask]
print(boxes_conf_landms.shape)
if len(boxes_conf_landms) <= 0:
return old_image
#---------------------------------------------------------#
# 如果使用了letterbox_image的话，要把灰条的部分去除掉。
#---------------------------------------------------------#
if self.letterbox_image:
boxes_conf_landms = retinaface_correct_boxes(boxes_conf_landms, \
np.array([self.input_shape[0], self.input_shape[1]]), np.array([im_height, im_width]))
boxes_conf_landms[:, :4] = boxes_conf_landms[:, :4] * scale
for b in boxes_conf_landms:
text = "{:.4f}".format(b[4])
b = list(map(int, b))
# b[0]-b[3]为人脸框的坐标，b[4]为得分
cv2.rectangle(old_image, (b[0], b[1]), (b[2], b[3]), (0, 0, 255),
2)
cx = b[0]
cy = b[1] + 12
cv2.putText(old_image, text, (cx, cy), cv2.FONT_HERSHEY_DUPLEX,
0.5, (255, 255, 255))
print(b[0], b[1], b[2], b[3], b[4])
# b[5]-b[14]为人脸关键点的坐标
return old_image
class Retinaface(object):
#-------------------------------#
# 请注意主干网络
# 与预训练权重的对应
# 即注意修改model_path
# 和backbone
#-------------------------------#
_defaults = {
"model_path": 'model_data/retinaface_mobilenet025.h5',
"backbone": "mobilenet",
"confidence": 0.5,
}
@classmethod
def get_defaults(cls, n):
if n in cls._defaults:
return cls._defaults[n]
else:
return "Unrecognized attribute name '" + n + "'"
#---------------------------------------------------#
# 初始化Retinaface
#---------------------------------------------------#
def __init__(self, **kwargs):
self.__dict__.update(self._defaults)
if self.backbone == "mobilenet":
self.cfg = cfg_mnet
else:
self.cfg = cfg_re50
self.bbox_util = BBoxUtility()
self.generate()
#---------------------------------------------------#
# 获得所有的分类
#---------------------------------------------------#
def generate(self):
model_path = os.path.expanduser(self.model_path)
assert model_path.endswith(
'.h5'), 'Keras model or weights must be a .h5 file.'
# 载入模型
self.retinaface = RetinaFace(self.cfg, self.backbone)
self.retinaface.load_weights(self.model_path, by_name=True)
#---------------------------------------------------#
# 检测图片
#---------------------------------------------------#
def detect_image(self, image):
old_image = image.copy()
image = np.array(image, np.float32)
im_height, im_width, _ = np.shape(image)
scale = [
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0]
]
scale_for_landmarks = [
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0]
]
# 图片预处理，归一化
photo = np.expand_dims(preprocess_input(image), 0)
anchors = Anchors(self.cfg,
image_size=(im_height, im_width)).get_anchors()
preds = self.retinaface.predict(photo)
# 将预测结果进行解码和非极大抑制
results = self.bbox_util.detection_out(
preds, anchors, confidence_threshold=self.confidence)
if len(results) <= 0:
return old_image
results = np.array(results)
results[:, :4] = results[:, :4] * scale
results[:, 5:] = results[:, 5:] * scale_for_landmarks
for b in results:
text = "{:.4f}".format(b[4])
b = list(map(int, b))
cv2.rectangle(old_image, (b[0], b[1]), (b[2], b[3]), (0, 0, 255),
2)
cx = b[0]
cy = b[1] + 12
cv2.putText(old_image, text, (cx, cy), cv2.FONT_HERSHEY_DUPLEX,
0.5, (255, 255, 255))
# landms
cv2.circle(old_image, (b[5], b[6]), 1, (0, 0, 255), 4)
cv2.circle(old_image, (b[7], b[8]), 1, (0, 255, 255), 4)
cv2.circle(old_image, (b[9], b[10]), 1, (255, 0, 255), 4)
cv2.circle(old_image, (b[11], b[12]), 1, (0, 255, 0), 4)
cv2.circle(old_image, (b[13], b[14]), 1, (255, 0, 0), 4)
return old_image
elif backbone == "resnet50":
cfg = cfg_re50
else:
raise ValueError(
'Unsupported backbone - `{}`, Use mobilenet, resnet50.'.format(
backbone))
img_dim = cfg['train_image_size']
# -------------------------------#
# 获得先验框anchors
# -------------------------------#
anchors = Anchors(cfg, image_size=(img_dim, img_dim)).get_anchors()
if Cuda:
anchors = anchors.cuda()
model = RetinaFace(cfg=cfg, pretrained=pretrained).train()
# -------------------------------------------#
# 权值文件的下载请看README
# 权值和主干特征提取网络一定要对应
# -------------------------------------------#
# model_path = "model_data/Retinaface_mobilenet0.25.pth"
# print('Loading weights into state dict...')
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# model_dict = model.state_dict()
# pretrained_dict = torch.load(model_path, map_location=device)
# pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
# model_dict.update(pretrained_dict)
# model.load_state_dict(model_dict)
# print('Finished!')
net = model
class Retinaface(object):
_defaults = {
"model_path": 'model_data/retinaface_mobilenet025.h5',
"backbone": "mobilenet",
"confidence": 0.5,
}
@classmethod
def get_defaults(cls, n):
if n in cls._defaults:
return cls._defaults[n]
else:
return "Unrecognized attribute name '" + n + "'"
#---------------------------------------------------#
# 初始化Retinaface
#---------------------------------------------------#
def __init__(self, **kwargs):
self.__dict__.update(self._defaults)
if self.backbone == "mobilenet":
self.cfg = cfg_mnet
else:
self.cfg = cfg_re50
self.bbox_util = BBoxUtility()
self.generate()
#---------------------------------------------------#
# 获得所有的分类
#---------------------------------------------------#
def generate(self):
model_path = os.path.expanduser(self.model_path)
assert model_path.endswith(
'.h5'), 'tensorflow.keras model or weights must be a .h5 file.'
# 加快模型训练的效率
print('Loading weights into state dict...')
# 载入模型
self.retinaface = RetinaFace(self.cfg, self.backbone)
self.retinaface.load_weights(self.model_path)
print('{} model, anchors loaded.'.format(self.model_path))
@tf.function
def get_pred(self, photo):
preds = self.retinaface(photo, training=False)
return preds
#---------------------------------------------------#
# 检测图片
#---------------------------------------------------#
def detect_image(self, image):
old_image = image.copy()
image = np.array(image, np.float32)
im_height, im_width, _ = np.shape(image)
scale = [im_width, im_height, im_width, im_height]
scale_for_landmarks = [
im_width, im_height, im_width, im_height, im_width, im_height,
im_width, im_height, im_width, im_height
]
# 图片预处理，归一化
photo = np.expand_dims(preprocess_input(image), 0)
anchors = Anchors(self.cfg,
image_size=(im_height, im_width)).get_anchors()
preds = self.get_pred(photo)
preds = [pred.numpy() for pred in preds]
# 将预测结果进行解码和非极大抑制
results = self.bbox_util.detection_out(
preds, anchors, confidence_threshold=self.confidence)
if len(results) <= 0:
return old_image, 0, 0
results = np.array(results)
results[:, :4] = results[:, :4] * scale
results[:, 5:] = results[:, 5:] * scale_for_landmarks
for b in results:
text = "{:.4f}".format(b[4])
b = list(map(int, b))
cv2.rectangle(old_image, (b[0], b[1]), (b[2], b[3]), (0, 0, 255),
2)
#####################
global cnt, t0, t1
t1 = time.time()
image_clip = old_image
# if t1 - t0 > 1:
# t0 = t1
# image_clip = old_image[b[1]-20:b[3]+20, b[0]-20:b[2]+20]
image_clip = old_image[b[1]:b[3], b[0]:b[2]]
image_clip = cv2.cvtColor(image_clip, cv2.COLOR_RGB2BGR)
# 保存剪切的图片
# cv2.imshow("clip", image_clip)
# cv2.imwrite("savedImg/wang/" + str(t1) + ".png", image_clip)
# cnt += 1
# print(cnt)
#####################
cx = b[0]
cy = b[1] + 12
cv2.putText(old_image, text, (cx, cy), cv2.FONT_HERSHEY_DUPLEX,
0.5, (255, 255, 255))
# landms
cv2.circle(old_image, (b[5], b[6]), 1, (0, 0, 255), 4)
cv2.circle(old_image, (b[7], b[8]), 1, (0, 255, 255), 4)
cv2.circle(old_image, (b[9], b[10]), 1, (255, 0, 255), 4)
cv2.circle(old_image, (b[11], b[12]), 1, (0, 255, 0), 4)
cv2.circle(old_image, (b[13], b[14]), 1, (255, 0, 0), 4)
return old_image, image_clip, len(results)
cfg = cfg_mnet
freeze_layers = 81
elif backbone == "resnet50":
cfg = cfg_re50
freeze_layers = 173
else:
raise ValueError(
'Unsupported backbone - `{}`, Use mobilenet, resnet50.'.format(
backbone))
img_dim = cfg['train_image_size']
#--------------------------------------#
# 载入模型与权值
# 请注意主干网络与预训练权重的对应
#--------------------------------------#
model = RetinaFace(cfg, backbone=backbone)
model_path = "model_data/retinaface_mobilenet025.h5"
model.load_weights(model_path, by_name=True, skip_mismatch=True)
#-------------------------------#
# 获得先验框和工具箱
#-------------------------------#
anchors = Anchors(cfg, image_size=(img_dim, img_dim)).get_anchors()
bbox_util = BBoxUtility(anchors)
#-------------------------------------------------------------------------------#
# 训练参数的设置
# logging表示tensorboard的保存地址
# checkpoint用于设置权值保存的细节，period用于修改多少epoch保存一次
# reduce_lr用于设置学习率下降的方式
# early_stopping用于设定早停，val_loss多次不下降自动结束训练，表示模型基本收敛
class Retinaface(object):
_defaults = {
"model_path": 'model_data/retinaface_mobilenet025.h5',
"backbone": 'mobilenet',
"confidence": 0.6,
"nms_iou": 0.4,
#----------------------------------------------------------------------#
# 是否需要进行图像大小限制。
# 开启后，会将输入图像的大小限制为input_shape。否则使用原图进行预测。
# keras代码中主干为mobilenet时存在小bug，当输入图像的宽高不为32的倍数
# 会导致检测结果偏差，主干为resnet50不存在此问题。
# 可根据输入图像的大小自行调整input_shape，注意为32的倍数，如[640, 640, 3]
#----------------------------------------------------------------------#
"input_shape": [1280, 1280, 3],
"letterbox_image": False
}
@classmethod
def get_defaults(cls, n):
if n in cls._defaults:
return cls._defaults[n]
else:
return "Unrecognized attribute name '" + n + "'"
#---------------------------------------------------#
# 初始化Retinaface
#---------------------------------------------------#
def __init__(self, **kwargs):
self.__dict__.update(self._defaults)
if self.backbone == "mobilenet":
self.cfg = cfg_mnet
else:
self.cfg = cfg_re50
self.bbox_util = BBoxUtility(nms_thresh=self.nms_iou)
self.generate()
self.anchors = Anchors(self.cfg,
image_size=(self.input_shape[0],
self.input_shape[1])).get_anchors()
def generate(self):
model_path = os.path.expanduser(self.model_path)
assert model_path.endswith(
'.h5'), 'Keras model or weights must be a .h5 file.'
self.retinaface = RetinaFace(self.cfg, self.backbone)
self.retinaface.load_weights(self.model_path, by_name=True)
print('{} model, anchors, and classes loaded.'.format(model_path))
def detect_image(self, image):
old_image = image.copy()
image = np.array(image, np.float32)
im_height, im_width, _ = np.shape(image)
scale = [
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0]
]
scale_for_landmarks = [
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0]
]
if self.letterbox_image:
image = letterbox_image(image,
[self.input_shape[1], self.input_shape[0]])
else:
self.anchors = Anchors(self.cfg,
image_size=(im_height,
im_width)).get_anchors()
photo = np.expand_dims(preprocess_input(image), 0)
preds = self.retinaface.predict(photo)
results = self.bbox_util.detection_out(
preds, self.anchors, confidence_threshold=self.confidence)
if len(results) <= 0:
return old_image, []
results = np.array(results)
if self.letterbox_image:
results = retinaface_correct_boxes(
results, np.array([self.input_shape[0], self.input_shape[1]]),
np.array([im_height, im_width]))
results[:, :4] = results[:, :4] * scale
results[:, 5:] = results[:, 5:] * scale_for_landmarks
ans = []
for b in results:
confidence = b[4].astype(float)
each_ans = {'box': [0, 0, 0, 0], 'confidence': 0, 'landmarks': []}
text = "{:.4f}".format(b[4])
b = list(map(int, b))
each_ans['box'][0] = b[0]
each_ans['box'][1] = b[1]
each_ans['box'][2] = b[2]
each_ans['box'][3] = b[3]
each_ans['confidence'] = confidence
cv2.rectangle(old_image, (b[0], b[1]), (b[2], b[3]), (0, 0, 255),
2)
cx = b[0]
cy = b[1] + 12
cv2.putText(old_image, text, (cx, cy), cv2.FONT_HERSHEY_DUPLEX,
0.5, (255, 255, 255))
print(b[0], b[1], b[2], b[3], b[4])
cv2.circle(old_image, (b[5], b[6]), 1, (0, 0, 255), 4)
cv2.circle(old_image, (b[7], b[8]), 1, (0, 255, 255), 4)
cv2.circle(old_image, (b[9], b[10]), 1, (255, 0, 255), 4)
cv2.circle(old_image, (b[11], b[12]), 1, (0, 255, 0), 4)
cv2.circle(old_image, (b[13], b[14]), 1, (255, 0, 0), 4)
landmarks = [
(b[5], b[6]),
(b[7], b[8]),
(b[9], b[10]),
(b[11], b[12]),
(b[13], b[14]),
]
each_ans['landmarks'] = landmarks
ans.append(each_ans)
return old_image, ans
from nets.retinaface import RetinaFace
from utils.config import cfg_mnet, cfg_re50
if __name__ == '__main__':
model = RetinaFace(cfg_mnet,backbone="mobilenet")
model.summary()
for i,layer in enumerate(model.layers):
print(i,layer.name)
#--------------------------------------------#
# 该部分代码只用于看网络结构，并非测试代码
#--------------------------------------------#
import torch
from torchsummary import summary
from nets.retinaface import RetinaFace
from utils.config import cfg_mnet
if __name__ == '__main__':
# 需要使用device来指定网络在GPU还是CPU运行
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = RetinaFace(cfg_mnet).to(device)
print('# generator parameters:',
sum(param.numel() for param in model.parameters()))
summary(model, input_size=(3, 2150, 2150))
class Retinaface(object):
_defaults = {
"model_path" : 'model_data/Retinaface_mobilenet0.25.pth',
"backbone" : 'mobilenet',
"confidence" : 0.5,
"nms_iou" : 0.45,
"cuda" : True,
#----------------------------------------------------------------------#
# 是否需要进行图像大小限制。
# 开启后，会将输入图像的大小限制为input_shape。否则使用原图进行预测。
# 可根据输入图像的大小自行调整input_shape，注意为32的倍数，如[640, 640, 3]
#----------------------------------------------------------------------#
"input_shape" : [1280, 1280, 3],
"letterbox_image" : True
}
@classmethod
def get_defaults(cls, n):
if n in cls._defaults:
return cls._defaults[n]
else:
return "Unrecognized attribute name '" + n + "'"
#---------------------------------------------------#
# 初始化Retinaface
#---------------------------------------------------#
def __init__(self, **kwargs):
self.__dict__.update(self._defaults)
if self.backbone == "mobilenet":
self.cfg = cfg_mnet
else:
self.cfg = cfg_re50
self.generate()
if self.letterbox_image:
self.anchors = Anchors(self.cfg, image_size=[self.input_shape[0], self.input_shape[1]]).get_anchors()
#---------------------------------------------------#
# 载入模型
#---------------------------------------------------#
def generate(self):
self.net = RetinaFace(cfg=self.cfg, mode='eval').eval()
#-------------------------------#
# 载入模型与权值
#-------------------------------#
print('Loading weights into state dict...')
state_dict = torch.load(self.model_path)
self.net.load_state_dict(state_dict)
if self.cuda:
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
self.net = nn.DataParallel(self.net)
self.net = self.net.cuda()
print('Finished!')
#---------------------------------------------------#
# 检测图片
#---------------------------------------------------#
def detect_image(self, image, picpath, frameNo):
#---------------------------------------------------#
# 对输入图像进行一个备份，后面用于绘图
#---------------------------------------------------#
old_image = image.copy()
image = np.array(image,np.float32)
#---------------------------------------------------#
# 计算scale，用于将获得的预测框转换成原图的高宽
#---------------------------------------------------#
scale = [np.shape(image)[1], np.shape(image)[0], np.shape(image)[1], np.shape(image)[0]]
scale_for_landmarks = [np.shape(image)[1], np.shape(image)[0], np.shape(image)[1], np.shape(image)[0],
np.shape(image)[1], np.shape(image)[0], np.shape(image)[1], np.shape(image)[0],
np.shape(image)[1], np.shape(image)[0]]
im_height, im_width, _ = np.shape(image)
#---------------------------------------------------------#
# letterbox_image可以给图像增加灰条，实现不失真的resize
#---------------------------------------------------------#
if self.letterbox_image:
image = np.array(letterbox_image(image, [self.input_shape[1], self.input_shape[0]]), np.float32)
else:
self.anchors = Anchors(self.cfg, image_size=(im_height, im_width)).get_anchors()
with torch.no_grad():
#-----------------------------------------------------------#
# 图片预处理，归一化。
#-----------------------------------------------------------#
image = torch.from_numpy(preprocess_input(image).transpose(2, 0, 1)).unsqueeze(0)
if self.cuda:
self.anchors = self.anchors.cuda()
image = image.cuda()
loc, conf, landms = self.net(image)
#-----------------------------------------------------------#
# 将预测结果进行解码
#-----------------------------------------------------------#
boxes = decode(loc.data.squeeze(0), self.anchors, self.cfg['variance'])
boxes = boxes.cpu().numpy()
conf = conf.data.squeeze(0)[:,1:2].cpu().numpy()
landms = decode_landm(landms.data.squeeze(0), self.anchors, self.cfg['variance'])
landms = landms.cpu().numpy()
boxes_conf_landms = np.concatenate([boxes, conf, landms],-1)
#到这里为止我们已经利用Retinaface_pytorch的预训练模型检测完了人脸，并获得了人脸框和人脸五个特征点的坐标信息，全保存在dets中，接下来为人脸剪切部分
# 用来储存生成的单张人脸的路径
path_save = "./curve/faces/" #你可以将这里的路径换成你自己的路径
'''
#剪切图片
#if args.show_cutting_image：
for num, b in enumerate(boxes_conf_landms): # dets中包含了人脸框和五个特征点的坐标
#if b[4] < 0.6:
# continue
b = list(map(int, b))
# landms，在人脸上画出特征点，要是你想保存不显示特征点的人脸图，你可以把这里注释掉
cv2.circle(old_image, (b[5], b[6]), 1, (0, 0, 255), 4)
cv2.circle(old_image, (b[7], b[8]), 1, (0, 255, 255), 4)
cv2.circle(old_image, (b[9], b[10]), 1, (255, 0, 255), 4)
cv2.circle(old_image, (b[11], b[12]), 1, (0, 255, 0), 4)
cv2.circle(old_image, (b[13], b[14]), 1, (255, 0, 0), 4)
#计算人脸框矩形大小
Height = b[3] - b[1]
Width = b[2] - b[0]
# 显示人脸矩阵大小
print("人脸数 / faces in all:", str(num+1), "\n")
print("窗口大小 / The size of window:"
, '\n', "高度 / height:", Height
, '\n', "宽度 / width: ", Width)
#根据人脸框大小，生成空白的图片
img_blank = np.zeros((Height, Width, 3), np.uint8)
# 将人脸填充到空白图片
for h in range(Height):
for w in range(Width):
img_blank[h][w] = old_image[b[1] + h][b[0] + w]
cv2.namedWindow("img_faces") # , 2)
#cv2.imshow("img_faces", img_blank) #显示图片
cv2.imwrite(path_save + "img_face_4" + str(num + 1) + ".jpg", img_blank) #将图片保存至你指定的文件夹
print("Save into:", path_save + "img_face_4" + str(num + 1) + ".jpg")
cv2.waitKey(0)
'''
boxes_conf_landms = non_max_suppression(boxes_conf_landms, self.confidence)
if len(boxes_conf_landms)<=0:
return False, old_image
#---------------------------------------------------------#
# 如果使用了letterbox_image的话，要把灰条的部分去除掉。
#---------------------------------------------------------#
if self.letterbox_image:
boxes_conf_landms = retinaface_correct_boxes(boxes_conf_landms, \
np.array([self.input_shape[0], self.input_shape[1]]), np.array([im_height, im_width]))
boxes_conf_landms[:,:4] = boxes_conf_landms[:,:4]*scale
boxes_conf_landms[:,5:] = boxes_conf_landms[:,5:]*scale_for_landmarks
num = 0
for b in boxes_conf_landms:
text = "{:.4f}".format(b[4])
b = list(map(int, b))
# b[0]-b[3]为人脸框的坐标，b[4]为得分
cv2.rectangle(old_image, (b[0], b[1]), (b[2], b[3]), (0, 0, 255), 2)
cx = b[0]
cy = b[1] + 12
cv2.putText(old_image, text, (cx, cy),
cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255))
print(b[0], b[1], b[2], b[3], b[4])
c = [-50, -50,50, 50, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
#b += c
c = np.array(b) + np.array(c)
if (c > 0).all() :
b = c
else:
for i, ind in enumerate(b[:4]):
if ind < 0:
print(i, b[i])
b[i] = 0
print(b[0], b[1], b[2], b[3], b[4])
#计算人脸框矩形大小
Height = b[3] - b[1]
Width = b[2] - b[0]
print("窗口大小 / The size of window:"
, '\n', "高度 / height:", Height
, '\n', "宽度 / width: ", Width)
img_blank = old_image[int(b[1]):int(b[3]), int(b[0]):int(b[2])] # height, width
savepath = picpath + "/" + str(frameNo) + "_" + str(num)+".jpg"
#cv2.namedWindow("img_faces") # , 2)
#cv2.imshow("img_faces", img_blank) #显示图片
#img_blank = cv2.cvtColor(img_blank,cv2.COLOR_RGB2BGR)
cv2.imwrite(savepath, img_blank) #将图片保存至你指定的文件夹
print("Save into:", savepath)
#cv2.waitKey(0)
'''
# b[5]-b[14]为人脸关键点的坐标
cv2.circle(old_image, (b[5], b[6]), 1, (0, 0, 255), 4)
cv2.circle(old_image, (b[7], b[8]), 1, (0, 255, 255), 4)
cv2.circle(old_image, (b[9], b[10]), 1, (255, 0, 255), 4)
cv2.circle(old_image, (b[11], b[12]), 1, (0, 255, 0), 4)
cv2.circle(old_image, (b[13], b[14]), 1, (255, 0, 0), 4)
'''
num += 1
return True, old_image
pretrained_path,
map_location=lambda storage, loc: storage.cuda(device))
if "state_dict" in pretrained_dict.keys():
pretrained_dict = remove_prefix(pretrained_dict['state_dict'],
'module.')
else:
pretrained_dict = remove_prefix(pretrained_dict, 'module.')
check_keys(model, pretrained_dict)
model.load_state_dict(pretrained_dict, strict=False)
return model
if __name__ == '__main__':
torch.set_grad_enabled(False)
net = RetinaFace(cfg=cfg, phase='test')
net = load_model(net, args.trained_model, args.cpu)
net.eval()
print('Finished loading model!')
cudnn.benchmark = True
device = torch.device("cpu" if args.cpu else "cuda")
net = net.to(device)
# testing dataset
dataset_folder = args.dataset_folder
anno_file = "label.txt"
# read images from annotation file
image_list = []
with open(os.path.join(dataset_folder, anno_file), 'r') as f:
type=int,
help='resume iter for retraining')
parser.add_argument('--save_folder',
default='./weights/',
help='Location to save checkpoint models')
args = parser.parse_args()
if not os.path.exists(args.save_folder):
os.mkdir(args.save_folder)
training_dataset = args.dataset_root
anno_file = args.anno_file
save_folder = args.save_folder
net = RetinaFace(cfg=cfg)
print(net)
if args.resume_net is not None:
print('Loading resume network...')
state_dict = torch.load(args.resume_net)
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
head = k[:7]
if head == 'module.':
name = k[7:] # remove `module.`
else:
name = k
new_state_dict[name] = v
class Retinaface(object):
_defaults = {
"model_path": 'model_data/Retinaface_mobilenet0.25.pth',
"confidence": 0.5,
"backbone": "mobilenet",
"cuda": True
}
@classmethod
def get_defaults(cls, n):
if n in cls._defaults:
return cls._defaults[n]
else:
return "Unrecognized attribute name '" + n + "'"
#---------------------------------------------------#
# 初始化Retinaface
#---------------------------------------------------#
def __init__(self, **kwargs):
self.__dict__.update(self._defaults)
if self.backbone == "mobilenet":
self.cfg = cfg_mnet
else:
self.cfg = cfg_re50
self.generate()
#---------------------------------------------------#
# 获得所有的分类
#---------------------------------------------------#
def generate(self):
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
self.net = RetinaFace(cfg=self.cfg, phase='eval').eval()
# 加快模型训练的效率
print('Loading weights into state dict...')
state_dict = torch.load(self.model_path)
self.net.load_state_dict(state_dict)
if self.cuda:
self.net = nn.DataParallel(self.net)
self.net = self.net.cuda()
print('Finished!')
#---------------------------------------------------#
# 检测图片
#---------------------------------------------------#
def detect_image(self, image):
# 绘制人脸框
old_image = image.copy()
image = np.array(image,np.float32)
im_height, im_width, _ = np.shape(image)
# 它的作用是将归一化后的框坐标转换成原图的大小
scale = torch.Tensor([np.shape(image)[1], np.shape(image)[0], np.shape(image)[1], np.shape(image)[0]])
scale_for_landmarks = torch.Tensor([np.shape(image)[1], np.shape(image)[0], np.shape(image)[1], np.shape(image)[0],
np.shape(image)[1], np.shape(image)[0], np.shape(image)[1], np.shape(image)[0],
np.shape(image)[1], np.shape(image)[0]])
# pytorch
image = preprocess_input(image).transpose(2, 0, 1)
# 增加batch_size维度
image = torch.from_numpy(image).unsqueeze(0)
# 计算先验框
anchors = Anchors(self.cfg, image_size=(im_height, im_width)).get_anchors()
with torch.no_grad():
if self.cuda:
scale = scale.cuda()
scale_for_landmarks = scale_for_landmarks.cuda()
image = image.cuda()
anchors = anchors.cuda()
loc, conf, landms = self.net(image) # forward pass
boxes = decode(loc.data.squeeze(0), anchors, self.cfg['variance'])
boxes = boxes * scale
boxes = boxes.cpu().numpy()
conf = conf.data.squeeze(0)[:,1:2].cpu().numpy()
landms = decode_landm(landms.data.squeeze(0), anchors, self.cfg['variance'])
landms = landms * scale_for_landmarks
landms = landms.cpu().numpy()
boxes_conf_landms = np.concatenate([boxes,conf,landms],-1)
boxes_conf_landms = non_max_suppression(boxes_conf_landms, self.confidence)
for b in boxes_conf_landms:
text = "{:.4f}".format(b[4])
b = list(map(int, b))
cv2.rectangle(old_image, (b[0], b[1]), (b[2], b[3]), (0, 0, 255), 2)
cx = b[0]
cy = b[1] + 12
cv2.putText(old_image, text, (cx, cy),
cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255))
# landms
cv2.circle(old_image, (b[5], b[6]), 1, (0, 0, 255), 4)
cv2.circle(old_image, (b[7], b[8]), 1, (0, 255, 255), 4)
cv2.circle(old_image, (b[9], b[10]), 1, (255, 0, 255), 4)
cv2.circle(old_image, (b[11], b[12]), 1, (0, 255, 0), 4)
cv2.circle(old_image, (b[13], b[14]), 1, (255, 0, 0), 4)
return old_image
if backbone == "mobilenet":
cfg = cfg_mnet
freeze_layers = 81
elif backbone == "resnet50":
cfg = cfg_re50
freeze_layers = 173
else:
raise ValueError('Unsupported backbone - `{}`, Use mobilenet, resnet50.'.format(backbone))
img_dim = cfg['image_size']
#-------------------------------#
# 创立模型
#-------------------------------#
model = RetinaFace(cfg, backbone=backbone)
model_path = "model_data/retinaface_mobilenet025.h5"
model.load_weights(model_path,by_name=True,skip_mismatch=True)
#-------------------------------#
# 获得先验框和工具箱
#-------------------------------#
anchors = Anchors(cfg, image_size=(img_dim, img_dim)).get_anchors()
bbox_util = BBoxUtility(anchors)
# 训练参数设置
logging = TensorBoard(log_dir="logs")
checkpoint = ModelCheckpoint('logs/ep{epoch:03d}-loss{loss:.3f}.h5',
monitor='loss', save_weights_only=True, save_best_only=False, period=1)
reduce_lr = ReduceLROnPlateau(monitor='loss', factor=0.5, patience=2, verbose=1)
early_stopping = EarlyStopping(monitor='loss', min_delta=0, patience=6, verbose=1)
class Retinaface(object):
_defaults = {
"model_path": 'model_data/retinaface_mobilenet025.h5',
"backbone": 'mobilenet',
"confidence": 0.5,
"nms_iou": 0.45,
#----------------------------------------------------------------------#
# 是否需要进行图像大小限制。
# 开启后，会将输入图像的大小限制为input_shape。否则使用原图进行预测。
# tf2代码中主干为mobilenet时存在小bug，当输入图像的宽高不为32的倍数
# 会导致检测结果偏差，主干为resnet50不存在此问题。
# 可根据输入图像的大小自行调整input_shape，注意为32的倍数，如[640, 640, 3]
#----------------------------------------------------------------------#
"input_shape": [1280, 1280, 3],
"letterbox_image": True
}
@classmethod
def get_defaults(cls, n):
if n in cls._defaults:
return cls._defaults[n]
else:
return "Unrecognized attribute name '" + n + "'"
#---------------------------------------------------#
# 初始化Retinaface
#---------------------------------------------------#
def __init__(self, **kwargs):
self.__dict__.update(self._defaults)
if self.backbone == "mobilenet":
self.cfg = cfg_mnet
else:
self.cfg = cfg_re50
self.bbox_util = BBoxUtility(nms_thresh=self.nms_iou)
self.generate()
self.anchors = Anchors(self.cfg,
image_size=(self.input_shape[0],
self.input_shape[1])).get_anchors()
#---------------------------------------------------#
# 载入模型
#---------------------------------------------------#
def generate(self):
model_path = os.path.expanduser(self.model_path)
assert model_path.endswith(
'.h5'), 'tensorflow.keras model or weights must be a .h5 file.'
#-------------------------------#
# 载入模型与权值
#-------------------------------#
self.retinaface = RetinaFace(self.cfg, self.backbone)
self.retinaface.load_weights(self.model_path)
print('{} model, anchors loaded.'.format(self.model_path))
@tf.function
def get_pred(self, photo):
preds = self.retinaface(photo, training=False)
return preds
#---------------------------------------------------#
# 检测图片
#---------------------------------------------------#
def detect_image(self, image):
#---------------------------------------------------#
# 对输入图像进行一个备份，后面用于绘图
#---------------------------------------------------#
old_image = image.copy()
image = np.array(image, np.float32)
im_height, im_width, _ = np.shape(image)
#---------------------------------------------------#
# 计算scale，用于将获得的预测框转换成原图的高宽
#---------------------------------------------------#
scale = [
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0]
]
scale_for_landmarks = [
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0],
np.shape(image)[1],
np.shape(image)[0]
]
#---------------------------------------------------------#
# letterbox_image可以给图像增加灰条，实现不失真的resize
#---------------------------------------------------------#
if self.letterbox_image:
image = letterbox_image(image,
[self.input_shape[1], self.input_shape[0]])
else:
self.anchors = Anchors(self.cfg,
image_size=(im_height,
im_width)).get_anchors()
#-----------------------------------------------------------#
# 图片预处理，归一化。
#-----------------------------------------------------------#
photo = np.expand_dims(preprocess_input(image), 0)
preds = self.get_pred(photo)
preds = [pred.numpy() for pred in preds]
#-----------------------------------------------------------#
# 将预测结果进行解码
#-----------------------------------------------------------#
results = self.bbox_util.detection_out(
preds, self.anchors, confidence_threshold=self.confidence)
#--------------------------------------#
# 如果没有检测到物体，则返回原图
#--------------------------------------#
if len(results) <= 0:
return old_image
results = np.array(results)
#---------------------------------------------------------#
# 如果使用了letterbox_image的话，要把灰条的部分去除掉。
#---------------------------------------------------------#
if self.letterbox_image:
results = retinaface_correct_boxes(
results, np.array([self.input_shape[0], self.input_shape[1]]),
np.array([im_height, im_width]))
results[:, :4] = results[:, :4] * scale
results[:, 5:] = results[:, 5:] * scale_for_landmarks
for b in results:
text = "{:.4f}".format(b[4])
b = list(map(int, b))
# b[0]-b[3]为人脸框的坐标，b[4]为得分
cv2.rectangle(old_image, (b[0], b[1]), (b[2], b[3]), (0, 0, 255),
2)
cx = b[0]
cy = b[1] + 12
cv2.putText(old_image, text, (cx, cy), cv2.FONT_HERSHEY_DUPLEX,
0.5, (255, 255, 255))
print(b[0], b[1], b[2], b[3], b[4])
# b[5]-b[14]为人脸关键点的坐标
cv2.circle(old_image, (b[5], b[6]), 1, (0, 0, 255), 4)
cv2.circle(old_image, (b[7], b[8]), 1, (0, 255, 255), 4)
cv2.circle(old_image, (b[9], b[10]), 1, (255, 0, 255), 4)
cv2.circle(old_image, (b[11], b[12]), 1, (0, 255, 0), 4)
cv2.circle(old_image, (b[13], b[14]), 1, (255, 0, 0), 4)
return old_image
img_dim = cfg['image_size']
#-------------------------------#
# Cuda的使用
#-------------------------------#
Cuda = True
#-------------------------------#
# Dataloder的使用
#-------------------------------#
Use_Data_Loader = True
#-------------------------------------------#
# 权值文件的下载请看README
# 权值和主干特征提取网络一定要对应
#-------------------------------------------#
model = RetinaFace(cfg=cfg, pretrained = pretrained).train()
model_path = "model_data/Retinaface_mobilenet0.25.pth"
# 加快模型训练的效率
print('Loading weights into state dict...')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_dict = model.state_dict()
pretrained_dict = torch.load(model_path, map_location=device)
pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
print('Finished!')
net = model
if Cuda:
net = torch.nn.DataParallel(model)
cudnn.benchmark = True
cfg = cfg_mnet
freeze_layers = 81
elif backbone == "resnet50":
cfg = cfg_re50
freeze_layers = 173
else:
raise ValueError(
'Unsupported backbone - `{}`, Use mobilenet, resnet50.'.format(
backbone))
img_dim = cfg['image_size']
#-------------------------------#
# 创立模型
#-------------------------------#
model = RetinaFace(cfg, backbone=backbone)
model_path = "model_data/retinaface_mobilenet025.h5"
model.load_weights(model_path, by_name=True, skip_mismatch=True)
#-------------------------------#
# Dataloder的使用
#-------------------------------#
Use_Data_Loader = True
#-------------------------------#
# 获得先验框和工具箱
#-------------------------------#
anchors = Anchors(cfg, image_size=(img_dim, img_dim)).get_anchors()
bbox_util = BBoxUtility(anchors)
#------------------------------------------------------#
# 主干特征提取网络特征通用，冻结训练可以加快训练速度
