def get_model():
pretrain_model = "/Users/yangjiang/temp/gpu1/FaceBoxes.pth"
num_classes = CONFIG["num_classes"]
params = torch.load(pretrain_model, map_location="cpu")
model = FaceBoxes(num_classes, "test")
model.load_state_dict(params)
model.eval()
return model
def __init__(self, **opt):
net = FaceBoxes(phase='test', size=None, num_classes=2) # initialize detector
net = load_model(net, opt.get("weights", 'weights/FaceBoxes.pth'), True)
net.eval()
self.net = net
self.top_k = opt.get('top_k', 5000)
self.confidence_threshold = opt.get('confidence_threshold', 0.05)
self.nms_threshold = opt.get('nms_threshold', 0.3)
self.keep_top_k = opt.get('keep_top_k', 750)
def load_face_model(self):
torch.set_grad_enabled(False)
# net and model
net = FaceBoxes(phase='test', size=None,
num_classes=2) # initialize detector
net = load_model(net, self.args.trained_model, self.args.cpu)
net.eval()
cudnn.benchmark = True
self.device = torch.device("cpu" if self.args.cpu else "cuda")
self.net = net.to(self.device)
class Detect(object):
def __init__(self, path, device):
self.net = FaceBoxes(phase='test', size=None, num_classes=2).to(device)
self.net = load_model(self.net, path, False)
self.net.eval()
self.device = device
def get_bbox(self, img_raw):
img = torch.FloatTensor(img_raw).to(self.device)
im_height, im_width, _ = img.size()
scale = torch.FloatTensor([im_width, im_height, im_width,
im_height]).to(self.device)
img -= torch.FloatTensor((104, 117, 123)).to(self.device)
img = img.permute(2, 0, 1).unsqueeze(0)
loc, conf = self.net(img) # forward pass
priorbox = PriorBox(cfg, image_size=(im_height, im_width))
priors = priorbox.forward()
priors = priors.to(self.device)
prior_data = priors.data
boxes = decode(loc.data.squeeze(0), prior_data, cfg['variance'])
boxes = boxes * scale
boxes = boxes.cpu().numpy()
scores = conf.squeeze(0).data.cpu().numpy()[:, 1]
# ignore low scores
inds = np.where(scores > 0.05)[0]
boxes = boxes[inds]
scores = scores[inds]
# keep top-K before NMS
order = scores.argsort()[::-1][:5000]
boxes = boxes[order]
scores = scores[order]
# do NMS
dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32,
copy=False)
# keep = py_cpu_nms(dets, args.nms_threshold)
keep = nms(dets, 0.3, force_cpu=False)
dets = dets[keep, :]
# keep top-K faster NMS
dets = dets[:750, :]
bboxes = []
for b in dets:
if b[4] < 0.65:
continue
b = list(map(int, b))
bboxes.append((b[0], b[1], b[2], b[3]))
return bboxes
def load_faceboxes():
"""
Load FaceBoxes model and weight in pytorch
"""
print('---------------------------------------------------')
pretrained_path = 'weights/FaceBoxes.pth'
net = FaceBoxes(phase='test', size=None,
num_classes=2) # initialize detector
net = load_model(net, pretrained_path)
net.eval()
print('Finished loading model')
print('---------------------------------------------------')
return net.cpu()
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 and model
net = FaceBoxes(phase='test', size=None,
num_classes=2) # initialize detector
net = load_model(net, args.trained_model, args.cpu)
net.eval()
print('Finished loading model!')
print(net)
cudnn.benchmark = True
device = torch.device("cpu" if args.cpu else "cuda")
net = net.to(device)
# save file
if not os.path.exists(args.save_folder):
os.makedirs(args.save_folder)
fw = open(os.path.join(args.save_folder, args.dataset + '_dets.txt'), 'w')
# testing dataset
testset_folder = os.path.join('data', args.dataset, 'images/')
testset_list = os.path.join('data', args.dataset, 'img_list.txt')
with open(testset_list, 'r') as fr:
chkpt = torch.load(ops.landmarks_model,
map_location=lambda storage, loc: storage)
landmarks_model.load_state_dict(chkpt)
landmarks_model.eval() # 设置为前向推断模式
print('load landmarks model : {}'.format(ops.landmarks_model))
print(
'\n/******************* landmarks model acc ******************/')
acc_model(ops, landmarks_model)
landmarks_model = landmarks_model.to(device)
#--------------------------------------------------------------------------- 构建人脸检测模型
# detect_model
detect_model = FaceBoxes(phase='test', size=None,
num_classes=2) # initialize detector
detect_model = load_model(detect_model, ops.detect_model, True)
detect_model.eval()
print('\n/******************* detect model acc ******************/')
acc_model(ops, detect_model)
detect_model = detect_model.to(device)
print('Finished loading model!')
# print(detect_model)
detect_model = detect_model.to(device)
video_capture = cv2.VideoCapture(ops.test_path)
resize = 1
with torch.no_grad():
idx = 0
while True:
class FaceDetector(object):
'''
Class for face detection
'''
def __init__(self,
trained_model,
cpu=True,
nms_threshold=0.5,
top_k=1000,
confidence_threshold=0.8,
keep_top_k=10):
super(FaceDetector, self).__init__()
self.trained_model = trained_model
self.net = FaceBoxes(phase='test', size=None, num_classes=2)
self.net = load_model(self.net, trained_model, cpu)
self.net.eval()
print('Finished loading model', trained_model)
self.device = torch.device("cpu" if cpu else "cuda")
self.net = self.net.to(self.device)
self.nms_threshold = nms_threshold
self.top_k = top_k
self.confidence_threshold = confidence_threshold
self.keep_top_k = keep_top_k
self.cpu = cpu
def predict(self, img_name):
img = np.float32(cv2.imread(img_name, cv2.IMREAD_COLOR))
resize = 1
if resize != 1:
img = cv2.resize(img,
None,
None,
fx=resize,
fy=resize,
interpolation=cv2.INTER_LINEAR)
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)
_t = {'forward_pass': Timer(), 'misc': Timer()}
_t['forward_pass'].tic()
loc, conf = self.net(img) # forward pass
_t['forward_pass'].toc()
_t['misc'].tic()
priorbox = PriorBox(cfg, image_size=(im_height, im_width))
priors = priorbox.forward()
priors = priors.to(self.device)
prior_data = priors.data
boxes = decode(loc.data.squeeze(0), prior_data, cfg['variance'])
boxes = boxes * scale / resize
boxes = boxes.cpu().numpy()
scores = conf.data.cpu().numpy()[:, 1]
# ignore low scores
inds = np.where(scores > self.confidence_threshold)[0]
boxes = boxes[inds]
scores = scores[inds]
# keep top-K before NMS
order = scores.argsort()[::-1][:self.top_k]
boxes = boxes[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)
keep = nms(dets, self.nms_threshold, force_cpu=self.cpu)
dets = dets[keep, :]
# keep top-K faster NMS
dets = dets[:self.keep_top_k, :]
_t['misc'].toc()
return dets
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 = 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
model = FaceBoxes(phase='test', size=None, num_classes=2)
device = torch.device("cuda")
model = load_model(model, trained_model_path)
model = model.to(device)
model.eval()
image_path = '/media/haoxue/WD/FaceBoxes.PyTorch/data/FDDB/images/2002/08/26/big/img_265.jpg'
img = np.float32(cv2.imread(image_path, cv2.IMREAD_COLOR))
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)
# if args.cuda:
img = img.cuda()
scale = scale.cuda()
out = model(img)
priorbox = PriorBox(cfg, out[2], (im_height, im_width), phase='test')
priors = priorbox.forward()
def faceboxes(img_raw, cur_frame_counter):
img = np.float32(img_raw)
torch.set_grad_enabled(False)
# net and model
net = FaceBoxes(phase='test', size=None,
num_classes=2) # initialize detector
net = load_model(net, args.trained_model, args.cpu)
net.eval()
# print('Finished loading model!')
# print(net)
cudnn.benchmark = True
device = torch.device("cpu" if args.cpu else "cuda")
net = net.to(device)
# testing scale
resize = 2
_t = {'forward_pass': Timer(), 'misc': Timer()}
# testing begin
if resize != 1:
img = cv2.resize(img,
None,
None,
fx=resize,
fy=resize,
interpolation=cv2.INTER_LINEAR)
im_height, im_width, _ = img.shape
scale = torch.Tensor(
[img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
# print(img)
img -= (104, 117, 123)
img = img.transpose(2, 0, 1)
img = torch.from_numpy(img).unsqueeze(0)
img = img.to(device)
scale = scale.to(device)
loc, conf = net(img) # forward pass
# print(loc.size(),conf.size())
priorbox = PriorBox(cfg, image_size=(im_height, im_width))
priors = priorbox.forward()
priors = priors.to(device)
prior_data = priors.data
boxes = decode(loc.data.squeeze(0), prior_data, cfg['variance'])
boxes = boxes * scale / resize
boxes = boxes.cpu().numpy()
scores = conf.squeeze(0).data.cpu().numpy()[:, 1]
# ignore low scores
inds = np.where(scores > args.confidence_threshold)[0]
boxes = boxes[inds]
scores = scores[inds]
# keep top-K before NMS
order = scores.argsort()[::-1][:args.top_k]
boxes = boxes[order]
scores = scores[order]
# do NMS
dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32,
copy=False)
# keep = py_cpu_nms(dets, args.nms_threshold)
keep = nms(dets, args.nms_threshold, force_cpu=args.cpu)
dets = dets[keep, :]
# keep top-K faster NMS
dets = dets[:args.keep_top_k, :]
_t['misc'].toc()
outputs_useful = []
for b in dets:
output_traffic = {}
if b[4] < args.vis_thres:
continue
b = list(map(int, b))
(left, right, top, bottom) = (b[0], b[2], b[1], b[3])
label_str = 'face'
output_traffic[label_str] = [left, right, top, bottom]
outputs_useful.append(output_traffic)
# show image
if args.show_image:
for b in dets:
if b[4] < args.vis_thres:
continue
text = "{:.4f}".format(b[4])
b = list(map(int, b))
cv2.rectangle(img_raw, (b[0], b[1]), (b[2], b[3]), (0, 0, 255), 2)
cx = b[0]
cy = b[1] + 12
# cv2.putText(img_raw, text, (cx, cy),
# cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255))
cv2.imwrite('boxes/' + str(cur_frame_counter) + '.jpg', img_raw)
return outputs_useful
class FaceBoxDetector:
def __init__(self, min_score=0.9, use_gpu=True):
# Minimum score to consider as a detection.
self.score_min = min_score
self.net = FaceBoxes(phase='test', size=None, num_classes=2)
self.use_gpu = use_gpu
self.logger = Logger()
def load_model(self, path_to_model):
self.logger.field('Loading pretrained model from', path_to_model)
device = torch.cuda.current_device()
pretrained_dict = torch.load(path_to_model, 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(self.net, pretrained_dict)
self.net.load_state_dict(pretrained_dict, strict=False)
self.net.eval()
if self.use_gpu:
self.net.cuda()
def detect(self, images) -> List[List[TrackingRegion]]:
frames = []
for img in images:
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)
if self.use_gpu:
img = img.cuda()
scale = scale.cuda()
out = self.net(img)
face_regions = self.nms_process(out, scale, im_height, im_width)
frames.append(face_regions)
return frames
def nms_process(self, network_output, scale, im_height, im_width) -> List[TrackingRegion]:
priorbox = PriorBox(cfg, network_output[2], (im_height, im_width), phase='test')
priors = priorbox.forward()
if self.use_gpu:
priors = priors.cuda()
loc, conf, _ = network_output
prior_data = priors.data
boxes = decode(loc.data.squeeze(0), prior_data, cfg['variance'])
boxes = boxes * scale
boxes = boxes.cpu().numpy()
scores = conf.data.cpu().numpy()[:, 1]
# ignore low scores
inds = np.where(scores > self.score_min)[0]
boxes = boxes[inds]
scores = scores[inds]
# keep top-K before NMS, top_k = 5
order = scores.argsort()[::-1][:5000]
boxes = boxes[order]
scores = scores[order]
# do NMS
dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
keep = nms(dets, 0.3, force_cpu=False)
dets = dets[keep, :]
# keep top-K faster NMS
dets = dets[:750, :]
regions = []
for i in range(dets.shape[0]):
face_region = TrackingRegion()
face_region.set_rect(left=dets[i, 0], top=dets[i, 1], right=dets[i, 2], bottom=dets[i, 3])
face_region.confidence = dets[i, 4]
face_region.data["class_id"] = "face"
regions.append(face_region)
return regions
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
self.logger.field('Missing keys', len(missing_keys))
self.logger.field('Unused checkpoint keys', len(unused_pretrained_keys))
self.logger.field('Used keys', len(used_pretrained_keys))
assert len(used_pretrained_keys) > 0, 'load NONE from pretrained checkpoint'
return True
@staticmethod
def remove_prefix(state_dict, prefix):
""" Old style model is stored with all names of parameters sharing common prefix 'module.' """
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()}
