def main():
num_classes = CONFIG["num_classes"]
max_epoch = CONFIG["max_epoch"]
save_directory = CONFIG["save_directory"]
phase = "train"
device = "cuda" if torch.cuda.is_available() else "cpu"
model = FaceBoxes(num_classes, phase)
model.to(device)
if not os.path.exists(save_directory):
os.makedirs(save_directory)
train_net(model, device, max_epoch, save_directory)
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)
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:
test_dataset = fr.read().split()
num_images = len(test_dataset)
# testing scale
if args.dataset == "FDDB":
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:
ret, img_raw = video_capture.read()
if ret:
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
print('Loading pretrained model from {}'.format(pretrained_path))
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')
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
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
net = load_model(net, model_path, False)
net.eval()
net.to("cuda")
model_name = model_path.split("/")[-1].split(".")[0]
onnx_model_path = f"models/onnx/base-model.onnx"
# export ONNX model
dummy_input = torch.randn(1, 3, input_dim, input_dim).to("cuda")
torch.onnx.export(net,
dummy_input,
onnx_model_path,
verbose=False,
input_names=['input'],
output_names=['output'])
"""
# try using pytorch2keras
keras_model = pytorch_to_keras(net, dummy_input, [(3, input_dim, input_dim)])
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, 'weights/Final_HandBoxes.pth', False)
net = torch.nn.DataParallel(net, [1,0]).cuda() # multiprocessing edler look here
net.to(f'cuda:{net.device_ids[0]}')
net.eval()
print('Finished loading model!')
#print(net)
cudnn.benchmark = True
device = torch.device(f'cuda:{net.device_ids[0]}')
#net = net.to(device)
# testing scale
resize = 2
_t = {'forward_pass': Timer(), 'misc': Timer()}
cap = cv2.VideoCapture(0)
# 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
net.load_state_dict(new_state_dict)
if args.ngpu > 1:
net = torch.nn.DataParallel(net, device_ids=list(range(args.ngpu)))
net.to(device)
cudnn.benchmark = True
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
criterion_m = CrossEntropyLoss2d()
criterion = MultiBoxLoss(num_classes, 0.35, True, 0, True, 7, 0.35, False)
priorbox = PriorBox(cfg)
with torch.no_grad():
priors = priorbox.forward()
priors = priors.to(device)