def detect_face(net, img, resize):
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 = img[[2, 1, 0], :, :]
img = torch.from_numpy(img).unsqueeze(0)
if args.cuda:
img = img.cuda()
scale = scale.cuda()
out = net(img) # forward pass
priorbox = PriorBox(cfg, out[2], (im_height, im_width), phase='test')
priors = priorbox.forward()
if args.cuda:
priors = priors.cuda()
loc, conf, _ = out
print(loc.size(), conf.size())
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 > 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]
#print(boxes)
# do NMS
dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32,
copy=False)
keep = nms(dets, args.nms_threshold, force_cpu=args.cpu)
dets = dets[keep, :]
#print(dets)
# keep top-K faster NMS
dets = dets[:args.keep_top_k, :]
return dets
net = torch.nn.DataParallel(net, device_ids=list(range(args.ngpu)))
if args.cuda:
net.cuda()
cudnn.benchmark = True
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
criterion = MultiBoxLoss(num_classes, (0.5, 0.35, 0.1), True, 0, True, 3, 0.35,
False)
priorbox = PriorBox(cfg)
with torch.no_grad():
priors = priorbox.forward()
if args.cuda:
priors = priors.cuda()
def train():
net.train()
epoch = 0 + args.resume_epoch
print('Loading Dataset...')
dataset = VOCDetection(
args.training_dataset,
preproc_s3fd(img_dim, rgb_means, cfg['max_expand_ratio']),
AnnotationTransform())
epoch_size = math.ceil(len(dataset) / args.batch_size)