def train(train_loader, model, criterion, optimizer, idx, flip=True):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
# switch to train mode
model.train()
end = time.time()
gt_win, pred_win = None, None
bar = Bar('Processing', max=len(train_loader))
for i, (inputs, target, meta) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
input_var = torch.autograd.Variable(inputs.to(device=device))
target_var = torch.autograd.Variable(
target.to(device=device, non_blocking=True))
# compute output
output = model(input_var)
score_map = output[-1].data.cpu()
loss = criterion(output[0], target_var)
for j in range(1, len(output)):
loss += criterion(output[j], target_var)
acc = accuracy(score_map, target, idx, pck_threshold)
# measure accuracy and record loss
losses.update(loss.item(), inputs.size(0))
acces.update(acc[0], inputs.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = (
'({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | '
'Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | Acc: {acc: .4f}'
).format(batch=i + 1,
size=len(train_loader),
data=data_time.val,
bt=batch_time.val,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg)
bar.next()
bar.finish()
return losses.avg, acces.avg
def validate(val_loader, model, criterion):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
# predictions
predictions = torch.Tensor(val_loader.dataset.__len__(), 12)
# switch to evaluate mode
model.eval()
end = time.time()
bar = Bar('Processing', max=len(val_loader))
for i, (inputs, target) in enumerate(val_loader):
# measure data loading time
data_time.update(time.time() - end)
input_var = torch.autograd.Variable(inputs.cuda())
input_var = input_var.unsqueeze(0)
input_var = input_var.type(torch.FloatTensor)
target_var = torch.autograd.Variable(target)
target_var = target_var.type(torch.FloatTensor)
target_var = target_var.cuda()
# compute output
output = model(input_var)
output = output.unsqueeze(0)
loss = criterion(output, target_var)
acc = accuracy(output.detach(), target.cpu())
# generate predictions
predictions[i] = output.detach()
# measure accuracy and record loss
losses.update(loss.item(), inputs.size(0))
acces.update(acc, inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.6f} | Acc: {acc: .4f}'.format(
batch=i + 1,
size=len(val_loader),
data=data_time.val,
bt=batch_time.avg,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg)
bar.next()
bar.finish()
return losses.avg, acces.avg, predictions
def train(train_loader, model, criterion, optimizer):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
# switch to train mode
model.train()
end = time.time()
bar = Bar('Processing', max=len(train_loader))
for i, (inputs, target) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
# target = target.cuda(async=True)
input_var = torch.autograd.Variable(inputs.cuda())
input_var = input_var.unsqueeze(0)
input_var = input_var.type(torch.FloatTensor)
target_var = torch.autograd.Variable(target)
target_var = target_var.type(torch.FloatTensor)
target_var = target_var.cuda()
# compute output
output = model(input_var)
output = output.unsqueeze(0)
loss = criterion(output, target_var)
acc = accuracy(output.detach(), target.cpu())
# measure accuracy and record loss
losses.update(loss.item(), inputs.size(0))
acces.update(acc, inputs.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.6f} | Acc: {acc: .4f}'.format(
batch=i + 1,
size=len(train_loader),
data=data_time.val,
bt=batch_time.val,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg)
bar.next()
bar.finish()
return losses.avg, acces.avg
def train(train_loader, model, tmodel, criterion, optimizer, kdloss_alpha, debug=False, flip=True):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
kdlosses = AverageMeter()
unkdlosses = AverageMeter()
tslosses = AverageMeter()
gtlosses = AverageMeter()
acces = AverageMeter()
# switch to train mode
model.train()
end = time.time()
gt_win, pred_win = None, None
bar = Bar('Processing', max=len(train_loader))
for i, (inputs, target, meta) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
input_var = torch.autograd.Variable(inputs.cuda())
target_var = torch.autograd.Variable(target.cuda(async=True))
# compute output
output = model(input_var)
score_map = output[-1].data.cpu()
# compute teacher network output
toutput = tmodel(input_var)
toutput = toutput[-1].detach()
# lmse : student vs ground truth
# gtmask will filter out the samples without ground truth
gtloss = torch.tensor(0.0).cuda()
kdloss = torch.tensor(0.0).cuda()
kdloss_unlabeled = torch.tensor(0.0).cuda()
unkdloss_alpha = 1.0
gtmask = meta['gtmask']
train_batch = score_map.shape[0]
for j in range(0, len(output)):
_output = output[j]
for i in range(gtmask.shape[0]):
if gtmask[i] < 0.1:
# unlabeled data, gtmask=0.0, kdloss only
# need to dividen train_batch to keep number equal
kdloss_unlabeled += criterion(_output[i,:,:,:], toutput[i, :,:,:])/train_batch
else:
# labeled data: kdloss + gtloss
gtloss += criterion(_output[i,:,:,:], target_var[i, :,:,:])/train_batch
kdloss += criterion(_output[i,:,:,:], toutput[i,:,:,:])/train_batch
loss_labeled = kdloss_alpha * (kdloss) + (1 - kdloss_alpha)*gtloss
total_loss = loss_labeled + unkdloss_alpha * kdloss_unlabeled
acc = accuracy(score_map, target, idx)
if debug: # visualize groundtruth and predictions
gt_batch_img = batch_with_heatmap(inputs, target)
pred_batch_img = batch_with_heatmap(inputs, score_map)
teacher_batch_img = batch_with_heatmap(inputs, toutput)
if not gt_win or not pred_win or not pred_teacher:
ax1 = plt.subplot(131)
ax1.title.set_text('Groundtruth')
gt_win = plt.imshow(gt_batch_img)
ax2 = plt.subplot(132)
ax2.title.set_text('Prediction')
pred_win = plt.imshow(pred_batch_img)
ax2 = plt.subplot(133)
ax2.title.set_text('teacher')
pred_teacher = plt.imshow(teacher_batch_img)
else:
gt_win.set_data(gt_batch_img)
pred_win.set_data(pred_batch_img)
pred_teacher.set_data(teacher_batch_img)
plt.pause(.05)
plt.draw()
# measure accuracy and record loss
gtlosses.update(gtloss.item(), inputs.size(0))
kdlosses.update(kdloss.item(), inputs.size(0))
unkdlosses.update(kdloss_unlabeled.item(), inputs.size(0))
losses.update(total_loss.item(), inputs.size(0))
acces.update(acc[0], inputs.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} ' \
'| Loss: {loss:.6f} | KdLoss:{kdloss:.6f}| unKdLoss:{unkdloss:.6f}| GtLoss:{gtloss:.6f} | Acc: {acc: .4f}'.format(
batch=i + 1,
size=len(train_loader),
data=data_time.val,
bt=batch_time.val,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
kdloss=kdlosses.avg,
unkdloss=unkdlosses.avg,
tsloss=tslosses.avg,
gtloss=gtlosses.avg,
acc=acces.avg
)
bar.next()
bar.finish()
return losses.avg, acces.avg
def validate(val_loader,
model,
criterion,
num_classes,
debug=False,
flip=True):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
# predictions
predictions = torch.Tensor(val_loader.dataset.__len__(), num_classes, 2)
# switch to evaluate mode
model.eval()
gt_win, pred_win = None, None
end = time.time()
bar = Bar('Eval ', max=len(val_loader))
with torch.no_grad():
for i, (input, target, meta, img_path) in enumerate(val_loader):
# measure data loading time
data_time.update(time.time() - end)
indexes = []
input = input.to(device, non_blocking=True)
#print (input.shape)
#image = input.cpu().permute(0,2,3,1).numpy()
#image = np.squeeze(image)
path = str(img_path)
path = path[3:len(path) - 2]
image = cv2.imread(path)
# cv2.imshow("image", image)
# cv2.waitKey(10)
# time.sleep(1)
target = target.to(device, non_blocking=True)
target_weight = meta['target_weight'].to(device, non_blocking=True)
# compute output
#print (input.shape)
output = model(input)
score_map = output[-1].cpu() if type(
output) == list else output.cpu()
if flip:
flip_input = torch.from_numpy(fliplr(
input.clone().numpy())).float().to(device)
flip_output = model(flip_input)
flip_output = flip_output[-1].cpu() if type(
flip_output) == list else flip_output.cpu()
flip_output = flip_back(flip_output)
score_map += flip_output
if type(output) == list: # multiple output
loss = 0
for o in output:
loss += criterion(o, target, target_weight)
output = output[-1]
else: # single output
loss = criterion(output, target, target_weight)
#print (acc)
# generate predictions
preds, vals = final_preds(score_map, meta['center'], meta['scale'],
[64, 64])
# for z in range(target.shape[1]):
# for j in range(target.shape[2]):
# for k in range(target.shape[3]):
# if target[0,z,j,k]==1.0:
# indexes.append(z)
# coords = np.squeeze(preds)
# for m in range(0,len(coords)):
# val = vals[0][m].numpy()
# if val>0.6: #threshold for confidence score
# x,y = coords[m][0].cpu().numpy(), coords[m][1].cpu().numpy()
# cv2.circle(image, (x,y), 1, (0,0,255), -1)
# #indexes.append(m)
acc = accuracy(score_map, target.cpu(), indexes)
#print ((target.cpu()).shape[1])
for n in range(score_map.size(0)):
predictions[meta['index'][n], :, :] = preds[n, :, :]
#print ("scored", score_map.shape)
if debug:
gt_batch_img = batch_with_heatmap(input, target)
pred_batch_img = batch_with_heatmap(input, score_map)
if not gt_win or not pred_win:
plt.subplot(121)
gt_win = plt.imshow(gt_batch_img)
plt.subplot(122)
pred_win = plt.imshow(pred_batch_img)
else:
gt_win.set_data(gt_batch_img)
pred_win.set_data(pred_batch_img)
plt.pause(.05)
plt.draw()
cv2.imwrite(
'/home/shantam/Documents/Programs/pytorch-pose/example/predictions/pred'
+ str(i) + '.png', image)
#time.sleep(5)
# measure accuracy and record loss
losses.update(loss.item(), input.size(0))
acces.update(acc[0], input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | Acc: {acc: .4f}'.format(
batch=i + 1,
size=len(val_loader),
data=data_time.val,
bt=batch_time.avg,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg)
bar.next()
bar.finish()
return losses.avg, acces.avg, predictions
def validate(val_loader, model, criterion, debug=False, flip=True):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
# predictions
predictions = torch.Tensor(val_loader.dataset.__len__(), 17, 2)
# switch to evaluate mode
model.eval()
gt_win, pred_win = None, None
end = time.time()
bar = Bar('Processing', max=len(val_loader))
for i, (inputs, target, meta) in enumerate(val_loader):
# measure data loading time
data_time.update(time.time() - end)
target = target.cuda(async=True)
input_var = torch.autograd.Variable(inputs.cuda(), volatile=True)
target_var = torch.autograd.Variable(target, volatile=True)
# compute output
output = model(input_var)
score_map = output[-1].data.cpu()
if flip:
flip_input_var = torch.autograd.Variable(
torch.from_numpy(fliplr(inputs.clone().numpy())).float().cuda(),
volatile=True
)
flip_output_var = model(flip_input_var)
flip_output = flip_back(flip_output_var[-1].data.cpu())
score_map += flip_output
loss = 0
for o in output:
loss += criterion(o, target_var)
acc = accuracy(score_map, target.cpu(), idx)
# generate predictions
preds = final_preds(score_map, meta['center'], meta['scale'], [64, 64])
for n in range(score_map.size(0)):
predictions[meta['index'][n], :, :] = preds[n, :, :]
if debug:
gt_batch_img = batch_with_heatmap(inputs, target)
pred_batch_img = batch_with_heatmap(inputs, score_map)
if not gt_win or not pred_win:
plt.subplot(121)
gt_win = plt.imshow(gt_batch_img)
plt.subplot(122)
pred_win = plt.imshow(pred_batch_img)
else:
gt_win.set_data(gt_batch_img)
pred_win.set_data(pred_batch_img)
plt.pause(.5)
plt.draw()
# measure accuracy and record loss
losses.update(loss.data[0], inputs.size(0))
acces.update(acc[0], inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | Acc: {acc: .4f}'.format(
batch=i + 1,
size=len(val_loader),
data=data_time.val,
bt=batch_time.avg,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg
)
bar.next()
bar.finish()
return losses.avg, acces.avg, predictions
def validate(val_loader,
model,
criterion,
debug=False,
flip=True,
test_batch=6,
njoints=68):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
# predictions
predictions = torch.Tensor(val_loader.dataset.__len__(), njoints, 2)
# switch to evaluate mode
model.eval()
gt_win, pred_win = None, None
end = time.time()
bar = Bar('Eval ', max=len(val_loader))
interocular_dists = torch.zeros((njoints, val_loader.dataset.__len__()))
with torch.no_grad():
for i, (input, target, meta) in enumerate(val_loader):
# measure data loading time
data_time.update(time.time() - end)
input = input.to(device, non_blocking=True)
target = target.to(device, non_blocking=True)
target_weight = meta['target_weight'].to(device, non_blocking=True)
# compute output
output = model(input)
score_map = output[-1].cpu() if type(
output) == list else output.cpu()
if flip:
flip_input = torch.from_numpy(fliplr(
input.clone().numpy())).float().to(device)
flip_output = model(flip_input)
flip_output = flip_output[-1].cpu() if type(
flip_output) == list else flip_output.cpu()
flip_output = flip_back(flip_output)
score_map += flip_output
if type(output) == list: # multiple output
loss = 0
for o in output:
loss += criterion(o, target, target_weight, len(idx))
output = output[-1]
else: # single output
loss = criterion(output, target, target_weight, len(idx))
acc, batch_interocular_dists = accuracy(score_map, target.cpu(),
idx)
interocular_dists[:, i * test_batch:(i + 1) *
test_batch] = batch_interocular_dists
# generate predictions
preds = final_preds(score_map, meta['center'], meta['scale'],
[64, 64])
for n in range(score_map.size(0)):
predictions[meta['index'][n], :, :] = preds[n, :, :]
if debug:
gt_batch_img = batch_with_heatmap(input, target)
pred_batch_img = batch_with_heatmap(input, score_map)
if not gt_win or not pred_win:
plt.subplot(121)
gt_win = plt.imshow(gt_batch_img)
plt.subplot(122)
pred_win = plt.imshow(pred_batch_img)
else:
gt_win.set_data(gt_batch_img)
pred_win.set_data(pred_batch_img)
plt.pause(.05)
plt.draw()
# measure accuracy and record loss
losses.update(loss.item(), input.size(0))
acces.update(acc[0], input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.8f} | Acc: {acc: .8f}'.format(
batch=i + 1,
size=len(val_loader),
data=data_time.val,
bt=batch_time.avg,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg)
bar.next()
bar.finish()
idx_array = np.array(idx) - 1
interocular_dists_pickup = interocular_dists[idx_array, :]
mean_error = torch.mean(
interocular_dists_pickup[interocular_dists_pickup != -1])
auc = calc_metrics(interocular_dists,
idx) # this is auc of predicted maps and target.
#print("=> Mean Error: {:.8f}, [email protected]: {:.8f} based on maps".format(mean_error, auc))
return losses.avg, acces.avg, predictions, auc, mean_error
def train(train_loader,
model,
criterion,
optimizer,
debug=False,
flip=True,
train_batch=6,
epoch=0,
njoints=68):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
# switch to train mode
model.train()
end = time.time()
gt_win, pred_win = None, None
bar = Bar('Train', max=len(train_loader))
interocular_dists = torch.zeros((njoints, train_loader.dataset.__len__()))
for i, (input, target, meta) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
input, target = input.to(device), target.to(device, non_blocking=True)
target_weight = meta['target_weight'].to(device, non_blocking=True)
# compute output
output = model(input)
if type(output) == list: # multiple output
loss = 0
for o in output:
loss += criterion(o, target, target_weight, len(idx))
output = output[-1]
else: # single output
loss = criterion(output, target, target_weight, len(idx))
acc, batch_interocular_dists = accuracy(output, target, idx)
interocular_dists[:, i * train_batch:(i + 1) *
train_batch] = batch_interocular_dists
if debug: # visualize groundtruth and predictions
gt_batch_img = batch_with_heatmap(input, target)
pred_batch_img = batch_with_heatmap(input, output)
if not gt_win or not pred_win:
ax1 = plt.subplot(121)
ax1.title.set_text('Groundtruth')
gt_win = plt.imshow(gt_batch_img)
ax2 = plt.subplot(122)
ax2.title.set_text('Prediction')
pred_win = plt.imshow(pred_batch_img)
else:
gt_win.set_data(gt_batch_img)
pred_win.set_data(pred_batch_img)
plt.pause(.05)
plt.draw()
# measure accuracy and record loss
losses.update(loss.item(), input.size(0))
acces.update(acc[0], input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.8f} | Acc: {acc: .8f}'.format(
batch=i + 1,
size=len(train_loader),
data=data_time.val,
bt=batch_time.val,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg)
bar.next()
bar.finish()
idx_array = np.array(idx) - 1
interocular_dists_pickup = interocular_dists[idx_array, :]
mean_error = torch.mean(
interocular_dists_pickup[interocular_dists_pickup != -1])
auc = calc_metrics(interocular_dists,
idx) # this is auc of predicted maps and target.
#print("=> Mean Error: {:.8f}, [email protected]: {:.8f} based on maps".format(mean_error, auc))
return losses.avg, acces.avg
def train(train_loader, model, criterion, optimizer, debug=False, flip=True):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
# switch to train mode
model.train()
end = time.time()
gt_win, pred_win = None, None
bar = Bar('Train', max=len(train_loader))
for i, (input, target, meta) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
input, target = input.to(device), target.to(device, non_blocking=True)
target_weight = meta['target_weight'].to(device, non_blocking=True)
# compute output
output = model(input)
if type(output) == list: # multiple output
loss = 0
for o in output:
loss += criterion(o, target, target_weight)
output = output[-1]
else: # single output
loss = criterion(output, target, target_weight)
acc = accuracy(output, target, idx)
if debug: # visualize groundtruth and predictions
gt_batch_img = batch_with_heatmap(input, target)
pred_batch_img = batch_with_heatmap(input, output)
if not gt_win or not pred_win:
ax1 = plt.subplot(121)
ax1.title.set_text('Groundtruth')
gt_win = plt.imshow(gt_batch_img)
ax2 = plt.subplot(122)
ax2.title.set_text('Prediction')
pred_win = plt.imshow(pred_batch_img)
else:
gt_win.set_data(gt_batch_img)
pred_win.set_data(pred_batch_img)
plt.pause(.05)
plt.draw()
# measure accuracy and record loss
losses.update(loss.item(), input.size(0))
acces.update(acc[0], input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | Acc: {acc: .4f}'.format(
batch=i + 1,
size=len(train_loader),
data=data_time.val,
bt=batch_time.val,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg)
bar.next()
bar.finish()
return losses.avg, acces.avg
def train(train_loader,
model,
criterion,
optimizer,
debug=False,
flip=True,
train_iters=0):
print("Train iters: {}".format(train_iters))
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
# switch to train mode
model.train()
debug_count = 0
end = time.time()
gt_win, pred_win = None, None
bar_len = [train_iters if train_iters != 0 else len(train_loader)][0]
train_iters = [train_iters if train_iters != 0 else len(train_loader)][0]
bar = Bar('Train', max=bar_len)
curr_iter = 0
while curr_iter < train_iters:
for i, (input, target, meta) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
input, target = input.to(device), target.to(device,
non_blocking=True)
target_weight = meta['target_weight'].to(device, non_blocking=True)
# compute output
output = model(input)
if type(output) == list: # multiple output
loss = 0
for o in output:
loss += criterion(o, target, target_weight)
output = output[-1]
else: # single output
loss = criterion(output, target, target_weight)
acc = accuracy(output, target, idx)
if debug: # visualize groundtruth and predictions
gt_batch_img = batch_with_heatmap(input, target)
pred_batch_img = batch_with_heatmap(input, output)
fig = plt.figure()
ax1 = fig.add_subplot(121)
ax1.title.set_text('Groundtruth')
gt_win = plt.imshow(gt_batch_img)
ax2 = fig.add_subplot(122)
ax2.title.set_text('Prediction')
pred_win = plt.imshow(pred_batch_img)
plt.pause(.05)
plt.draw()
fig.savefig('debug/debug_{}.png'.format(str(debug_count)),
dpi=500)
debug_count += 1
# measure accuracy and record loss
losses.update(loss.item(), input.size(0))
acces.update(acc[0], input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | Acc: {acc: .4f}'.format(
batch=i + 1,
size=[len(train_loader) if train_iters == 0 else train_iters
][0],
data=data_time.val,
bt=batch_time.val,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg)
bar.next()
curr_iter += 1
if curr_iter >= train_iters - 1:
break
bar.finish()
return losses.avg, acces.avg
def main():
_t = {'iter time': Timer()}
# create directory
model_name = args.source + '_to_' + args.target
if not os.path.exists(args.checkpoint):
os.makedirs(args.checkpoint)
os.makedirs(os.path.join(args.checkpoint, 'logs'))
os.makedirs(os.path.join(args.checkpoint, 'ssl_labels'))
opt.print_options(args)
# load datasets
# load synthetic datasets
source_train_dataset = datasets.__dict__[args.source](is_train=True,
**vars(args))
source_train_loader = torch.utils.data.DataLoader(
source_train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
source_test_dataset = datasets.__dict__[args.source](is_train=False,
**vars(args))
source_test_loader = torch.utils.data.DataLoader(
source_test_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# load real training set images with pseudo-labels
target_train_dataset = datasets.__dict__[args.target_ssl](is_train=True,
is_aug=False,
**vars(args))
target_train_loader = torch.utils.data.DataLoader(
target_train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
# load original real test set with ground truth labels
target_test_dataset = datasets.__dict__[args.target](is_train=False,
is_aug=False,
**vars(args))
target_test_loader = torch.utils.data.DataLoader(
target_test_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# create model and optimizer
model, optimizer = CreateModel(args, models, datasets)
cudnn.enabled = True
cudnn.benchmark = True
model.train()
model = torch.nn.DataParallel(model).to(device)
if isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
best_acc = checkpoint['best_acc']
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
loss = ['loss_kpt_src', 'loss_kpt_trg']
_t['iter time'].tic()
trg_val_acc_best = 0
# CC-SSL training
for epoch in range(args.num_epochs):
if epoch == 0:
trg_val_loss, trg_val_acc = validate(target_test_loader, model,
criterion, args.flip,
args.batch_size, njoints)
# generate ssl labels every 10 epoch
if epoch % 10 == 0:
print("==> generating ssl labels")
target_train_dataset = datasets.__dict__[args.target](
is_train=True, is_aug=False, **vars(args))
target_train_loader = torch.utils.data.DataLoader(
target_train_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
model.eval()
# generate labels on target training set
ssl_kpts = {}
acces1 = AverageMeter()
previous_img = None
previous_kpts = None
for _, (trg_img, trg_lbl,
trg_meta) in enumerate(target_train_loader):
trg_img = trg_img.to(device)
trg_lbl = trg_lbl.to(device, non_blocking=True)
# generate labels for each image
for i in range(trg_img.size(0)):
score_map, generated_kpts = prediction_check(
previous_img, previous_kpts, trg_img[i], model,
target_train_dataset)
ssl_kpts[int(trg_meta['index'][i].cpu().numpy().astype(
np.int32))] = generated_kpts
if global_animal == 'tiger':
trg_lbl[i] = trg_lbl[
i,
np.array([
1, 2, 3, 4, 5, 6, 7, 8, 15, 16, 17, 18, 13, 14,
9, 10, 11, 12
]) - 1, :, :]
acc1, _ = accuracy(score_map,
trg_lbl[i].cpu().unsqueeze(0), idx)
acces1.update(acc1[0], 1)
previous_img = trg_img[i]
previous_kpts = generated_kpts
print('Acc on target training set (pseudo-labels):', acces1.avg)
# modify confidence score based on ranking
sorted_confidence = np.zeros(1)
for k in ssl_kpts:
sorted_confidence = np.concatenate(
(sorted_confidence, ssl_kpts[k][:, 2].reshape(-1)), axis=0)
sorted_confidence = np.sort(sorted_confidence)
np.save(
args.checkpoint +
'/ssl_labels/ssl_labels_train_confidence.npy',
sorted_confidence)
p = (1.0 - 0.02 * (epoch + 10))
if p > 0.2:
ccl_thresh = sorted_confidence[int(p *
sorted_confidence.shape[0])]
else:
p = 0.2
ccl_thresh = sorted_confidence[int(p *
sorted_confidence.shape[0])]
print("=====> ccl_thresh: ", ccl_thresh)
for k in ssl_kpts:
ssl_kpts[k][:, 2] = (ssl_kpts[k][:, 2] > ccl_thresh).astype(
np.float32)
np.save(args.checkpoint + 'ssl_labels/ssl_labels_train.npy',
ssl_kpts)
# generate labels on target test set for diagnosis
ssl_kpts = {}
acces1 = AverageMeter()
previous_img = None
previous_kpts = None
for jj, (trg_img, trg_lbl,
trg_meta) in enumerate(target_test_loader):
trg_img = trg_img.to(device)
trg_lbl = trg_lbl.to(device, non_blocking=True)
# generate labels for each image
for i in range(trg_img.size(0)):
score_map, generated_kpts = prediction_check(
previous_img, previous_kpts, trg_img[i], model,
target_test_dataset)
ssl_kpts[int(trg_meta['index'][i].cpu().numpy().astype(
np.int32))] = generated_kpts
if global_animal == 'tiger':
trg_lbl[i] = trg_lbl[
i,
np.array([
1, 2, 3, 4, 5, 6, 7, 8, 15, 16, 17, 18, 13, 14,
9, 10, 11, 12
]) - 1, :, :]
acc1, _ = accuracy(score_map,
trg_lbl[i].cpu().unsqueeze(0), idx)
acces1.update(acc1[0], 1)
previous_img = trg_img[i]
previous_kpts = generated_kpts
print('Acc on target testing set (pseudo-labels):', acces1.avg)
np.save(
args.checkpoint + 'ssl_labels/ssl_labels_valid' + str(epoch) +
'.npy', ssl_kpts)
# load real training set images with pseudo-labels
target_train_dataset = datasets.__dict__[args.target_ssl](
is_train=True, is_aug=True, **vars(args))
target_train_loader = torch.utils.data.DataLoader(
target_train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
print("======> start training")
loss_output_src = AverageMeter()
loss_output_trg = AverageMeter()
joint_loader = zip(source_train_loader, target_train_loader)
model.train()
# training with source images and target images jointly
for i, ((src_img, src_lbl, src_meta),
(trg_img, trg_lbl, trg_meta)) in enumerate(joint_loader):
optimizer.zero_grad()
# calculate loss on source dataset
src_img, src_lbl, src_weight = src_img.to(device), src_lbl.to(
device, non_blocking=True), src_meta['target_weight'].to(
device, non_blocking=True)
src_kpt_score = model(src_img)
if type(src_kpt_score) == list: # multiple output
loss_kpt_src = 0
for o in src_kpt_score:
loss_kpt_src += criterion(o, src_lbl, src_weight, len(idx))
src_kpt_score = src_kpt_score[-1]
else: # single output
loss_kpt_src = criterion(src_kpt_score, src_lbl, src_weight,
len(idx))
loss_output_src.update(loss_kpt_src.data.item(), src_img.size(0))
loss_kpt_src.backward()
# calculate loss on target dataset
trg_img, trg_lbl, trg_weight = trg_img.to(device), trg_lbl.to(
device, non_blocking=True), trg_meta['target_weight'].to(
device, non_blocking=True)
trg_kpt_score = model(trg_img)
if type(trg_kpt_score) == list: # multiple output
loss_kpt_trg = 0
for o in trg_kpt_score:
loss_kpt_trg += criterion(o, trg_lbl, trg_weight, len(idx))
trg_kpt_score = trg_kpt_score[-1]
else: # single output
loss_kpt_trg = criterion(trg_kpt_score, trg_lbl, trg_weight,
len(idx))
loss_kpt_trg *= args.gamma_
loss_output_trg.update(loss_kpt_trg.data.item(), src_img.size(0))
loss_kpt_trg.backward()
# update
optimizer.step()
# print logs
if (i + 1) % args.print_freq == 0:
_t['iter time'].toc(average=False)
print('[epoch %d][it %d][src kpt loss %.6f][trg kpt loss %.6f][lr %.6f][%.2fs]' % \
(epoch+1, i + 1, loss_output_src.avg, loss_output_trg.avg, optimizer.param_groups[0]['lr'], _t['iter time'].diff))
_t['iter time'].tic()
print('\nEvaluation')
src_val_loss, src_val_acc = validate(source_test_loader, model,
criterion, args.flip,
args.batch_size, njoints)
trg_val_loss, trg_val_acc = validate(target_test_loader, model,
criterion, args.flip,
args.batch_size, njoints)
# save best model
if trg_val_acc > trg_val_acc_best:
trg_val_acc_best = trg_val_acc
print('\ntaking snapshot ...')
state = {
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
}
torch.save(
state,
os.path.join(args.checkpoint,
'%s' % (args.source) + '.pth.tar'))
def validate(val_loader,
model,
criterion,
flip=True,
test_batch=6,
njoints=18):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
# predictions
predictions = torch.Tensor(val_loader.dataset.__len__(), njoints, 2)
# switch to evaluate mode
model.eval()
gt_win, pred_win = None, None
end = time.time()
bar = Bar('Eval ', max=len(val_loader))
with torch.no_grad():
for i, (input, target, meta) in enumerate(val_loader):
# measure data loading time
data_time.update(time.time() - end)
input = input.to(device, non_blocking=True)
if global_animal == 'horse':
target = target.to(device, non_blocking=True)
target_weight = meta['target_weight'].to(device,
non_blocking=True)
elif global_animal == 'tiger':
target = target.to(device, non_blocking=True)
target_weight = meta['target_weight'].to(device,
non_blocking=True)
target = target[:,
np.array([
1, 2, 3, 4, 5, 6, 7, 8, 15, 16, 17, 18, 13,
14, 9, 10, 11, 12
]) - 1, :, :]
target_weight = target_weight[:,
np.array([
1, 2, 3, 4, 5, 6, 7, 8, 15,
16, 17, 18, 13, 14, 9, 10,
11, 12
]) - 1, :]
else:
raise Exception('please add new animal category')
# compute output
output = model(input)
score_map = output[-1].cpu() if type(
output) == list else output.cpu()
if flip:
flip_input = torch.from_numpy(fliplr(
input.clone().numpy())).float().to(device)
flip_output = model(flip_input)
flip_output = flip_output[-1].cpu() if type(
flip_output) == list else flip_output.cpu()
flip_output = flip_back(flip_output)
score_map += flip_output
if type(output) == list: # multiple output
loss = 0
for o in output:
loss += criterion(o, target, target_weight, len(idx))
output = output[-1]
else: # single output
loss = criterion(output, target, target_weight, len(idx))
acc, _ = accuracy(score_map, target.cpu(), idx)
# generate predictions
preds = final_preds(score_map, meta['center'], meta['scale'],
[64, 64])
#for n in range(score_map.size(0)):
# predictions[meta['index'][n], :, :] = preds[n, :, :]
# measure accuracy and record loss
losses.update(loss.item(), input.size(0))
acces.update(acc[0], input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.8f} | Acc: {acc: .8f}'.format(
batch=i + 1,
size=len(val_loader),
data=data_time.val,
bt=batch_time.avg,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg)
bar.next()
bar.finish()
return losses.avg, acces.avg
def validate(val_loader,
model,
criterion,
criterion_seg,
debug=False,
flip=True,
test_batch=6,
njoints=68):
batch_time = AverageMeter()
data_time = AverageMeter()
losses_kpt = AverageMeter()
losses_seg = AverageMeter()
acces = AverageMeter()
inter_meter = AverageMeter()
union_meter = AverageMeter()
# predictions
predictions = torch.Tensor(val_loader.dataset.__len__(), njoints, 2)
# switch to evaluate mode
model.eval()
gt_win, pred_win = None, None
end = time.time()
bar = Bar('Eval ', max=len(val_loader))
interocular_dists = torch.zeros((njoints, val_loader.dataset.__len__()))
with torch.no_grad():
for i, (input, target, target_seg, meta) in enumerate(val_loader):
# measure data loading time
data_time.update(time.time() - end)
input, target, target_seg = input.to(device), target.to(
device, non_blocking=True), target_seg.to(device)
target_weight = meta['target_weight'].to(device, non_blocking=True)
# compute output
output_kpt, output_seg = model(input)
score_map = output_kpt[-1].cpu() if type(
output_kpt) == list else output_kpt.cpu()
if flip:
flip_input = torch.from_numpy(fliplr(
input.clone().numpy())).float().to(device)
flip_output = model(flip_input)
flip_output = flip_output[-1].cpu() if type(
flip_output) == list else flip_output.cpu()
flip_output = flip_back(flip_output)
score_map += flip_output
if type(output_kpt) == list: # multiple output
loss_kpt = 0
loss_seg = 0
for (o, o_seg) in zip(output_kpt, output_seg):
loss_kpt += criterion(o, target, target_weight, len(idx))
loss_seg += criterion_seg(o_seg, target_seg)
output = output_kpt[-1]
output_seg = output_seg[-1]
else: # single output
loss_kpt = criterion(output_kpt, target, target_weight,
len(idx))
loss_seg = criterion(output_seg, target_seg)
acc, batch_interocular_dists = accuracy(score_map, target.cpu(),
idx)
_, pred_seg = torch.max(output_seg, 1)
# generate predictions
preds = final_preds(score_map, meta['center'], meta['scale'],
[64, 64])
for n in range(score_map.size(0)):
predictions[meta['index'][n], :, :] = preds[n, :, :]
if debug:
gt_batch_img = batch_with_heatmap(input, target)
pred_batch_img = batch_with_heatmap(input, score_map)
if not gt_win or not pred_win:
plt.subplot(121)
gt_win = plt.imshow(gt_batch_img)
plt.subplot(122)
pred_win = plt.imshow(pred_batch_img)
else:
gt_win.set_data(gt_batch_img)
pred_win.set_data(pred_batch_img)
plt.pause(.05)
plt.draw()
# measure accuracy and record loss
losses_kpt.update(loss_kpt.item(), input.size(0))
losses_seg.update(loss_seg.item(), input.size(0))
acces.update(acc[0], input.size(0))
inter, union = inter_and_union(
pred_seg.data.cpu().numpy().astype(np.uint8),
target_seg.data.cpu().numpy().astype(np.uint8))
inter_meter.update(inter)
union_meter.update(union)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
iou = inter_meter.sum / (union_meter.sum + 1e-10)
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss_kpt: {loss_kpt:.8f} | Loss_seg: {loss_seg:.8f} | Acc: {acc: .8f} | IOU: {iou:.2f}'.format(
batch=i + 1,
size=len(val_loader),
data=data_time.val,
bt=batch_time.avg,
total=bar.elapsed_td,
eta=bar.eta_td,
loss_kpt=losses_kpt.avg,
loss_seg=losses_seg.avg,
acc=acces.avg,
iou=iou.mean() * 100)
bar.next()
bar.finish()
print(iou)
return losses_kpt.avg, acces.avg, predictions, iou.mean() * 100
def train(train_loader,
model,
criterion,
criterion_seg,
optimizer,
debug=False,
flip=True,
train_batch=6,
epoch=0,
njoints=68):
batch_time = AverageMeter()
data_time = AverageMeter()
losses_kpt = AverageMeter()
losses_seg = AverageMeter()
acces = AverageMeter()
inter_meter = AverageMeter()
union_meter = AverageMeter()
# switch to train mode
model.train()
end = time.time()
gt_win, pred_win = None, None
bar = Bar('Train', max=len(train_loader))
for i, (input, target, target_seg, meta) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
input, target, target_seg = input.to(device), target.to(
device, non_blocking=True), target_seg.to(device)
target_weight = meta['target_weight'].to(device, non_blocking=True)
# compute output
output_kpt, output_seg = model(input)
if type(output_kpt) == list: # multiple output
loss_kpt = 0
loss_seg = 0
for (o, o_seg) in zip(output_kpt, output_seg):
loss_kpt += criterion(o, target, target_weight, len(idx))
loss_seg += criterion_seg(o_seg, target_seg)
output_kpt = output_kpt[-1]
output_seg = output_seg[-1]
else: # single output
loss_kpt = criterion(output_kpt, target, target_weight, len(idx))
loss_seg = criterion_seg(output_seg, target_seg)
acc, batch_interocular_dists = accuracy(output_kpt, target, idx)
_, pred_seg = torch.max(output_seg, 1)
if debug: # visualize groundtruth and predictions
gt_batch_img = batch_with_heatmap(input, target)
pred_batch_img = batch_with_heatmap(input, output)
if not gt_win or not pred_win:
ax1 = plt.subplot(121)
ax1.title.set_text('Groundtruth')
gt_win = plt.imshow(gt_batch_img)
ax2 = plt.subplot(122)
ax2.title.set_text('Prediction')
pred_win = plt.imshow(pred_batch_img)
else:
gt_win.set_data(gt_batch_img)
pred_win.set_data(pred_batch_img)
plt.pause(.05)
plt.draw()
loss = loss_kpt + (0.01 / (epoch + 1)) * loss_seg
# measure accuracy and record loss
losses_kpt.update(loss_kpt.item(), input.size(0))
losses_seg.update(loss_seg.item(), input.size(0))
acces.update(acc[0], input.size(0))
inter, union = inter_and_union(
pred_seg.data.cpu().numpy().astype(np.uint8),
target_seg.data.cpu().numpy().astype(np.uint8))
inter_meter.update(inter)
union_meter.update(union)
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
iou = inter_meter.sum / (union_meter.sum + 1e-10)
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss_kpt: {loss_kpt:.8f} | Loss_seg: {loss_seg:.8f} | Acc: {acc: .4f} | IOU: {iou: .2f}'.format(
batch=i + 1,
size=len(train_loader),
data=data_time.val,
bt=batch_time.val,
total=bar.elapsed_td,
eta=bar.eta_td,
loss_kpt=losses_kpt.avg,
loss_seg=losses_seg.avg,
acc=acces.avg,
iou=iou.mean() * 100)
bar.next()
bar.finish()
print(iou)
return losses_kpt.avg, acces.avg
def validate(val_loader,
model,
criterion,
num_classes,
idx,
save_result_dir,
meta_dir,
anno_type,
flip=True,
evaluate=False,
scales=[0.7, 0.8, 0.9, 1, 1.2, 1.4, 1.6],
multi_scale=False,
save_heatmap=False):
anno_type = anno_type[0].lower()
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
num_scales = len(scales)
# switch to evaluate mode
model.eval()
meanstd_file = '../datasets/arm/mean.pth.tar'
meanstd = torch.load(meanstd_file)
mean = meanstd['mean']
gt_win, pred_win = None, None
end = time.time()
bar = Bar('Processing', max=len(val_loader))
for i, (inputs, target, meta) in enumerate(val_loader):
#print(inputs.shape)
# measure data loading time
data_time.update(time.time() - end)
if anno_type != 'none':
target = target.cuda(async=True)
target_var = torch.autograd.Variable(target)
input_var = torch.autograd.Variable(inputs.cuda())
with torch.no_grad():
# compute output
output = model(input_var)
score_map = output[-1].data.cpu()
if flip:
flip_input_var = torch.autograd.Variable(
torch.from_numpy(fliplr(
inputs.clone().numpy())).float().cuda(), )
flip_output_var = model(flip_input_var)
flip_output = flip_back(flip_output_var[-1].data.cpu(),
meta_dir=meta_dir[0])
score_map += flip_output
score_map /= 2
if anno_type != 'none':
loss = 0
for o in output:
loss += criterion(o, target_var)
acc = accuracy(score_map, target.cpu(), idx, pck_threshold)
if multi_scale:
new_scales = []
new_res = []
new_score_map = []
new_inp = []
new_meta = []
img_name = []
confidence = []
new_center = []
num_imgs = score_map.size(0) // num_scales
for n in range(num_imgs):
score_map_merged, res, conf = multi_scale_merge(
score_map[num_scales * n:num_scales * (n + 1)].numpy(),
meta['scale'][num_scales * n:num_scales * (n + 1)])
inp_merged, _, _ = multi_scale_merge(
inputs[num_scales * n:num_scales * (n + 1)].numpy(),
meta['scale'][num_scales * n:num_scales * (n + 1)])
new_score_map.append(score_map_merged)
new_scales.append(meta['scale'][num_scales * (n + 1) - 1])
new_center.append(meta['center'][num_scales * n])
new_res.append(res)
new_inp.append(inp_merged)
img_name.append(meta['img_name'][num_scales * n])
confidence.append(conf)
if len(new_score_map) > 1:
score_map = torch.tensor(
np.stack(new_score_map)) #stack back to 4-dim
inputs = torch.tensor(np.stack(new_inp))
else:
score_map = torch.tensor(
np.expand_dims(new_score_map[0], axis=0))
inputs = torch.tensor(np.expand_dims(new_inp[0], axis=0))
else:
img_name = []
confidence = []
for n in range(score_map.size(0)):
img_name.append(meta['img_name'][n])
confidence.append(
np.amax(score_map[n].numpy(), axis=(1, 2)).tolist())
# generate predictions
if multi_scale:
preds = final_preds(score_map, new_center, new_scales, new_res[0])
else:
preds = final_preds(score_map, meta['center'], meta['scale'],
[64, 64])
for n in range(score_map.size(0)):
if evaluate:
with open(
os.path.join(save_result_dir, 'preds',
img_name[n] + '.json'), 'w') as f:
obj = {
'd2_key': preds[n].numpy().tolist(),
'score': confidence[n]
}
json.dump(obj, f)
if evaluate:
for n in range(score_map.size(0)):
inp = inputs[n]
pred = score_map[n]
for t, m in zip(inp, mean):
t.add_(m)
scipy.misc.imsave(
os.path.join(save_result_dir, 'visualization',
'{}.jpg'.format(img_name[n])),
sample_with_heatmap(inp, pred))
if save_heatmap:
score_map_original_size = align_back(
score_map[n], meta['center'][n],
meta['scale'][len(scales) * n - 1],
meta['original_size'][n])
np.save(
os.path.join(save_result_dir, 'heatmaps',
'{}.npy'.format(img_name[n])),
score_map_original_size)
if anno_type != 'none':
# measure accuracy and record loss
losses.update(loss.item(), inputs.size(0))
acces.update(acc[0], inputs.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | Acc: {acc: .4f}'.format(
batch=i + 1,
size=len(val_loader),
data=data_time.val,
bt=batch_time.avg,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg)
bar.next()
bar.finish()
if anno_type != 'none':
return losses.avg, acces.avg
else:
return 0, 0
def validate(unit_path, model, criterion, num_classes, debug=False, flip=True):
losses = AverageMeter()
acces = AverageMeter()
dimensions = [0, 0]
scale_factor = 0
with Image.open(unit_path) as img:
dimensions = map(lambda x: x / 2, img.size)
scale_factor = (img.size[1]) / 200.0
print(dimensions)
print(scale_factor)
# Custom edits [for testing]
# dimensions = [459.0, 335.0]
scale_factor = 6.0969 + 3
# predictions
predictions = torch.Tensor(1, num_classes, 2)
# switch to evaluate mode
model.eval()
end = time.time()
with torch.no_grad():
input_temp, target_temp, meta = customMpiiObject.get_image_data(
'./data/custom/images/015620151.jpg', dimensions, scale_factor)
input = torch.zeros([1, 3, 256, 256])
target = torch.zeros([1, 16, 64, 64])
target_weight = torch.zeros([1, 16, 1])
input[0, :, :, :] = input_temp.to(device, non_blocking=True)
target[0, :, :, :] = target.to(device, non_blocking=True)
target_weight[0, :, :] = meta['target_weight'].to(device,
non_blocking=True)
# compute output
output = model(input)
score_map = output[-1].cpu() if type(output) == list else output.cpu()
# if type(output) == list: # multiple output
# loss = 0
# for o in output:
# loss += criterion(o, target, target_weight)
# output = output[-1]
# else: # single output
# loss = criterion(output, target, target_weight)
acc = accuracy(score_map, target.cpu(), idx)
meta_center_temp = torch.Tensor(meta['center'])
meta['center'] = torch.zeros([1, 2])
meta['center'][0, :] = meta_center_temp
# generate predictions
preds = final_preds(score_map, meta['center'], meta['scale'], [64, 64])
# for n in range(score_map.size(0)):
# predictions[meta['index'][n], :, :] = preds[n, :, :]
predictions = preds
acces.update(acc[0], input.size(0))
# measure accuracy and record loss
return losses.avg, acces.avg, predictions
def train(train_loader, model, criterion, optimizer, debug=False, flip=True):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
# switch to train mode
model.train()
end = time.time()
gt_win, pred_win = None, None
bar = Bar('Processing', max=len(train_loader))
for i, (inputs, target, meta) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
input_var = torch.autograd.Variable(inputs.cuda())
target_var = torch.autograd.Variable(target.cuda(async=True))
# compute output
output = model(input_var)
score_map = output[-1].data.cpu()
loss = criterion(output[0], target_var)
for j in range(1, len(output)):
loss += criterion(output[j], target_var)
acc = accuracy(score_map, target, idx)
if debug: # visualize groundtruth and predictions
gt_batch_img = batch_with_heatmap(inputs, target)
pred_batch_img = batch_with_heatmap(inputs, score_map)
if not gt_win or not pred_win:
ax1 = plt.subplot(121)
ax1.title.set_text('Groundtruth')
gt_win = plt.imshow(gt_batch_img)
ax2 = plt.subplot(122)
ax2.title.set_text('Prediction')
pred_win = plt.imshow(pred_batch_img)
else:
gt_win.set_data(gt_batch_img)
pred_win.set_data(pred_batch_img)
plt.pause(.05)
plt.draw()
# measure accuracy and record loss
losses.update(loss.data[0], inputs.size(0))
acces.update(acc[0], inputs.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | Acc: {acc: .4f}'.format(
batch=i + 1,
size=len(train_loader),
data=data_time.val,
bt=batch_time.val,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg
)
bar.next()
bar.finish()
return losses.avg, acces.avg
def validate(self):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
predictions = torch.Tensor(self.val_loader.dataset.__len__(),
self.num_classes, 2)
self.netG.eval()
gt_win, pred_win = None, None
end = time.time()
bar = Bar('Eval ', max=len(self.val_loader))
with torch.no_grad():
for i, (input, target, meta, mpii) in enumerate(self.val_loader):
if mpii == False:
continue
data_time.update(time.time() - end)
input = input.to(self.device, non_blocking=True)
target = target.to(self.device, non_blocking=True)
target_weight = meta['target_weight'].to(self.device,
non_blocking=True)
output = self.netG(input)
score_map = output[-1].cpu() if type(
output) == list else output.cpu()
if self.flip:
flip_input = torch.from_numpy
flip_output = self.netG(flip_input)
flip_output = flip_output[-1].cpu() if type(
flip_output) == list else flip_output.cpu()
flip_output = flip_back(flip_output)
score_map += flip_output
if type(output) == list:
loss = 0
for o in output:
loss += self.criterion(o, target, target_weight)
output = output[-1]
else:
loss = self.criterion(output, target, target_weight)
acc = accuracy(score_map, target.cpu(), self.idx)
preds = final_preds(score_map, meta['center'], meta['scale'],
[64, 64])
for n in range(score_map.size(0)):
predictions[meta['index'][n], :, :] = preds[n, :, :]
if self.debug:
gt_batch_img = batch_with_heatmap(input, target)
pred_batch_img = batch_with_heatmap(input, score_map)
if not gt_win or not pred_win:
plt.subplot(121)
gt_win = plt.imshow(gt_batch_img)
plt.subplot(122)
pred_win = plt.imshow(pred_batch_img)
else:
gt_win.set_data(gt_batch_img)
pred_win.set_data(pred_batch_img)
plt.pause(.05)
plt.draw()
losses.update(loss.item, input.size(0))
acces.update(acc[0], input.size(0))
batch_time.update(time.time() - end)
end = time.time()
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | Acc: {acc: .4f}'.format(
batch=i + 1,
size=len(self.val_loader),
data=data_time.val,
bt=batch_time.avg,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg)
bar.next()
bar.finish()
return losses.avg, acces.avg, predictions
def validate(val_loader,
model,
criterion,
num_classes,
debug=False,
flip=True):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
# predictions
predictions = torch.Tensor(val_loader.dataset.__len__(), num_classes, 2)
# switch to evaluate mode
model.eval()
gt_win, pred_win = None, None
end = time.time()
bar = Bar('Eval ', max=len(val_loader))
with torch.no_grad():
for i, (input, target, meta) in enumerate(val_loader):
# measure data loading time
data_time.update(time.time() - end)
input = input.to(device, non_blocking=True)
target = target.to(device, non_blocking=True)
target_weight = meta['target_weight'].to(device, non_blocking=True)
# compute output
output = model(input)
score_map = output[-1].cpu() if type(
output) == list else output.cpu()
if flip:
flip_input = torch.from_numpy(fliplr(
input.clone().numpy())).float().to(device)
flip_output = model(flip_input)
flip_output = flip_output[-1].cpu() if type(
flip_output) == list else flip_output.cpu()
flip_output = flip_back(flip_output)
score_map += flip_output
if type(output) == list: # multiple output
loss = 0
for o in output:
loss += criterion(o, target, target_weight)
output = output[-1]
else: # single output
loss = criterion(output, target, target_weight)
acc = accuracy(score_map, target.cpu(), idx)
# generate predictions
preds = final_preds(score_map, meta['center'], meta['scale'],
[64, 64])
for n in range(score_map.size(0)):
predictions[meta['index'][n], :, :] = preds[n, :, :]
if debug:
gt_batch_img = batch_with_heatmap(input, target)
pred_batch_img = batch_with_heatmap(input, score_map)
if not gt_win or not pred_win:
plt.subplot(121)
gt_win = plt.imshow(gt_batch_img)
plt.subplot(122)
pred_win = plt.imshow(pred_batch_img)
else:
gt_win.set_data(gt_batch_img)
pred_win.set_data(pred_batch_img)
plt.pause(.05)
plt.draw()
# measure accuracy and record loss
losses.update(loss.item(), input.size(0))
acces.update(acc[0], input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.6f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | Acc: {acc: .4f}'.format(
batch=i + 1,
size=len(val_loader),
data=data_time.val,
bt=batch_time.avg,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc=acces.avg)
bar.next()
bar.finish()
return losses.avg, acces.avg, predictions
def validate(val_loader,
model,
criterion,
debug=False,
flip=True,
test_batch=6,
njoints=68):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces_re = AverageMeter()
# switch to evaluate mode
model.eval()
gt_win, pred_win = None, None
end = time.time()
bar = Bar('Eval ', max=len(val_loader))
with torch.no_grad():
for i, (input, target, meta) in enumerate(val_loader):
# measure data loading time
data_time.update(time.time() - end)
input = input.to(device, non_blocking=True)
target = target.to(device, non_blocking=True)
target_weight = meta['target_weight'].to(device, non_blocking=True)
# compute output
output, output_refine = model(input)
score_map = output[-1].cpu() if type(
output) == list else output.cpu()
score_map_refine = output_refine[-1].cpu() if type(
output_refine) == list else output_refine.cpu()
if flip:
flip_input = torch.from_numpy(fliplr(
input.clone().numpy())).float().to(device)
flip_output, flip_output_re = model(flip_input)
flip_output = flip_output[-1].cpu() if type(
flip_output) == list else flip_output.cpu()
flip_output_re = flip_output_re[-1].cpu() if type(
flip_output_re) == list else flip_output_re.cpu()
flip_output = flip_back(flip_output, 'real_animal')
flip_output_re = flip_back(flip_output_re, 'real_animal')
score_map += flip_output
score_map_refine += flip_output_re
if type(output) == list: # multiple output
loss = 0
for (o, o_re) in (output, output_refine):
loss = loss + criterion(
o, target, target_weight, len(idx)) + criterion(
o_re, target, target_weight, len(idx))
else: # single output
loss = criterion(
output, target, target_weight, len(idx)) + criterion(
output_refine, target, target_weight, len(idx))
acc_re, _ = accuracy(score_map_refine, target.cpu(), idx)
if debug:
gt_batch_img = batch_with_heatmap(input, target)
pred_batch_img = batch_with_heatmap(input, score_map)
if not gt_win or not pred_win:
plt.subplot(121)
gt_win = plt.imshow(gt_batch_img)
plt.subplot(122)
pred_win = plt.imshow(pred_batch_img)
else:
gt_win.set_data(gt_batch_img)
pred_win.set_data(pred_batch_img)
plt.pause(.05)
plt.draw()
# measure accuracy and record loss
losses.update(loss.item(), input.size(0))
acces_re.update(acc_re[0], input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# plot progress
bar.suffix = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} ' \
'| Loss: {loss:.8f} | Acc_re: {acc_re: .8f}'.format(
batch=i + 1,
size=len(val_loader),
data=data_time.val,
bt=batch_time.avg,
total=bar.elapsed_td,
eta=bar.eta_td,
loss=losses.avg,
acc_re=acces_re.avg
)
bar.next()
bar.finish()
return losses.avg, acces_re.avg
