bboxes = []
for i in range(detections.size(1)):
j = 0
while detections[0, i, j, 0] >= thresh:
box = []
score = detections[0, i, j, 0]
pt = (detections[0, i, j, 1:] * scale).cpu().numpy().astype(np.int)
j += 1
box += [pt[0], pt[1], pt[2] - pt[0], pt[3] - pt[1], score]
bboxes += [box]
return bboxes
if __name__ == '__main__':
net = build_net('test', cfg.NUM_CLASSES)
net.load_state_dict(torch.load(args.model))
net.eval()
if use_cuda:
net.cuda()
cudnn.benckmark = True
#transform = S3FDBasicTransform(cfg.INPUT_SIZE, cfg.MEANS)
counter = 0
for i in range(10):
txt_in = os.path.join(FDDB_FOLD_DIR, 'FDDB-fold-%02d.txt' % (i + 1))
txt_out = os.path.join(FDDB_RESULT_DIR, 'fold-%02d-out.txt' % (i + 1))
answer_in = os.path.join(FDDB_FOLD_DIR,
def train():
iteration = 0
start_epoch = 0
step_index = 0
per_epoch_size = len(train_dataset) // args.batch_size
pyramidbox = build_net('train', cfg.NUM_CLASSES)
net = pyramidbox
if args.resume:
print(('Resuming training, loading {}...'.format(args.resume)))
state_dict = torch.load(args.resume)
net.load_state_dict(state_dict)
#iteration = int(args.resume[args.resume.rfind('pyramidbox_')+len('pyramidbox_'):-4])
#start_epoch = iteration // per_epoch_size
#print('Start iteration: %d (epoch %d)'%(iteration, start_epoch))
else:
vgg_weights = torch.load(args.save_folder + args.basenet)
print('Load base network....')
net.vgg.load_state_dict(vgg_weights)
if args.cuda:
torch.backends.cudnn.benchmark = True
if args.multigpu:
net = torch.nn.DataParallel(pyramidbox)
net = net.cuda()
if not args.resume:
print('Initializing weights...')
pyramidbox.extras.apply(pyramidbox.weights_init)
pyramidbox.lfpn_topdown.apply(pyramidbox.weights_init)
pyramidbox.lfpn_later.apply(pyramidbox.weights_init)
pyramidbox.cpm.apply(pyramidbox.weights_init)
pyramidbox.loc_layers.apply(pyramidbox.weights_init)
pyramidbox.conf_layers.apply(pyramidbox.weights_init)
optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum,
weight_decay=args.weight_decay)
criterion1 = MultiBoxLoss(cfg, args.cuda)
criterion2 = MultiBoxLoss(cfg, args.cuda, use_head_loss=True)
print('Loading wider dataset...')
print('Using the specified args:')
print(args)
for step in cfg.LR_STEPS:
if iteration > step:
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
net.train()
for epoch in range(start_epoch, cfg.EPOCHES):
losses = 0
for batch_idx, (images, face_targets, head_targets) in enumerate(train_loader):
with torch.no_grad():
if args.cuda:
images = Variable(images.cuda())
face_targets = [Variable(ann.cuda())
for ann in face_targets]
head_targets = [Variable(ann.cuda())
for ann in head_targets]
else:
images = Variable(images)
face_targets = [Variable(ann)
for ann in face_targets]
head_targets = [Variable(ann)
for ann in head_targets]
if iteration in cfg.LR_STEPS:
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
t0 = time.time()
out = net(images)
# backprop
optimizer.zero_grad()
face_loss_l, face_loss_c = criterion1(out, face_targets)
head_loss_l, head_loss_c = criterion2(out, head_targets)
loss = face_loss_l + face_loss_c + head_loss_l + head_loss_c
losses += loss.item()
loss.backward()
optimizer.step()
t1 = time.time()
face_loss = (face_loss_l + face_loss_c).item()
head_loss = (head_loss_l + head_loss_c).item()
if iteration % 10 == 0:
loss_ = losses / (batch_idx + 1)
print(('Timer: {:.4f} sec.'.format(t1 - t0)))
print(('epoch ' + repr(epoch) + ' iter ' +
repr(iteration) + ' || Loss:%.4f' % (loss_)))
print(('->> face Loss: {:.4f} || head loss : {:.4f}'.format(
face_loss, head_loss)))
print(('->> lr: {}'.format(optimizer.param_groups[0]['lr'])))
if iteration != 0 and iteration % 5000 == 0:
print(('Saving state, iter:', iteration))
file = 'pyramidbox_' + repr(iteration) + '.pth'
torch.save(pyramidbox.state_dict(),
os.path.join(args.save_folder, file))
iteration += 1
logger.add_scalar('data/loss', loss, epoch)
logger.add_scalar('data/face_loss', face_loss, epoch)
logger.add_scalar('data/head_loss', head_loss, epoch)
val(epoch, net, pyramidbox, criterion1, criterion2)
