def load_model(arch, stacks, blocks, num_classes, mobile, checkpoint_resume):
# create model
model = models.__dict__[arch](num_stacks=stacks,
num_blocks=blocks,
num_classes=num_classes,
mobile=mobile)
# optionally resume from a checkpoint
if isfile(checkpoint_resume):
print("=> loading checkpoint '{}'".format(checkpoint_resume))
checkpoint = torch.load(checkpoint_resume,
map_location=lambda storage, loc: storage)
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in checkpoint['state_dict'].items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
# load params
model.load_state_dict(new_state_dict)
print("=> loaded checkpoint '{}' (epoch {})".format(
checkpoint_resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(checkpoint_resume))
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
model.eval()
return model
def main(args):
global best_acc
global idx1
global idx2
# idx is the index of joints used to compute accuracy for dataset2
idx1 = range(1, 19)
idx2 = range(1, 19) # horse
# create model
njoints = datasets.__dict__[args.dataset].njoints
print("==> creating model '{}'".format(args.arch))
model = models.__dict__[args.arch](num_classes=njoints,
resnet_layers=args.resnet_layers,
pretrained=None,
dual_branch=True)
model = torch.nn.DataParallel(model).to(device)
# define loss function (criterion) and optimizer
criterion = losses.JointsMSELoss().to(device)
# optionally resume from a checkpoint
if args.resume:
if isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_acc = checkpoint['best_acc']
model.load_state_dict(checkpoint['state_dict_ema'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
else:
raise Exception('please provide a checkpoint')
val_dataset = datasets.__dict__[args.dataset](is_train=False,
is_aug=False,
**vars(args))
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
_, acc, predictions = validate(val_loader, model, criterion, njoints, args,
args.flip, args.test_batch)
return
def main(args):
# create checkpoint dir
if not isdir(args.checkpoint):
mkdir_p(args.checkpoint)
# create model
print("==> creating model '{}', stacks={}, blocks={}".format(
args.arch, args.stacks, args.blocks))
model = models.__dict__[args.arch](num_stacks=args.stacks,
num_blocks=args.blocks,
num_classes=args.num_classes,
mobile=args.mobile)
model.eval()
# optionally resume from a checkpoint
title = 'mpii-' + args.arch
if args.checkpoint:
if isfile(args.checkpoint):
print("=> loading checkpoint '{}'".format(args.checkpoint))
checkpoint = torch.load(args.checkpoint)
args.start_epoch = checkpoint['epoch']
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in checkpoint['state_dict'].items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
# load params
model.load_state_dict(new_state_dict)
print("=> loaded checkpoint '{}' (epoch {})".format(
args.checkpoint, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.checkpoint))
else:
logger = Logger(join(args.checkpoint, 'log.txt'), title=title)
logger.set_names(
['Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train Acc', 'Val Acc'])
cudnn.benchmark = True
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
dummy_input = torch.randn(1, 3, args.in_res, args.in_res)
torch.onnx.export(model, dummy_input, args.out_onnx)
def main(args):
global best_acc
global idx
# idx is the index of joints used to compute accuracy
if args.dataset in ['mpii', 'lsp']:
idx = [1, 2, 3, 4, 5, 6, 11, 12, 15, 16]
elif args.dataset == 'coco':
idx = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
else:
print("Unknown dataset: {}".format(args.dataset))
assert False
# create checkpoint dir
if not isdir(args.checkpoint):
mkdir_p(args.checkpoint)
# create model
njoints = datasets.__dict__[args.dataset].njoints
print("==> creating model '{}', stacks={}, blocks={}".format(
args.arch, args.stacks, args.blocks))
model = models.__dict__[args.arch](num_stacks=args.stacks,
num_blocks=args.blocks,
num_classes=njoints,
resnet_layers=args.resnet_layers)
model = torch.nn.DataParallel(model).to(device)
# define loss function (criterion) and optimizer
criterion = losses.JointsMSELoss().to(device)
if args.solver == 'rms':
optimizer = torch.optim.RMSprop(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.solver == 'adam':
optimizer = torch.optim.Adam(
model.parameters(),
lr=args.lr,
)
else:
print('Unknown solver: {}'.format(args.solver))
assert False
# optionally resume from a checkpoint
title = args.dataset + ' ' + args.arch
if args.resume:
if isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_acc = checkpoint['best_acc']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
logger = Logger(join(args.checkpoint, 'log.txt'),
title=title,
resume=True)
else:
print("=> no checkpoint found at '{}'".format(args.resume))
else:
logger = Logger(join(args.checkpoint, 'log.txt'), title=title)
logger.set_names(
['Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train Acc', 'Val Acc'])
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
# create data loader
train_dataset = datasets.__dict__[args.dataset](is_train=True,
**vars(args))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_dataset = datasets.__dict__[args.dataset](is_train=False, **vars(args))
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# evaluation only
if args.evaluate:
print('\nEvaluation only')
loss, acc, predictions = validate(val_loader, model, criterion,
njoints, args.debug, args.flip)
save_pred(predictions, checkpoint=args.checkpoint)
return
# train and eval
lr = args.lr
for epoch in range(args.start_epoch, args.epochs):
lr = adjust_learning_rate(optimizer, epoch, lr, args.schedule,
args.gamma)
print('\nEpoch: %d | LR: %.8f' % (epoch + 1, lr))
# decay sigma
if args.sigma_decay > 0:
train_loader.dataset.sigma *= args.sigma_decay
val_loader.dataset.sigma *= args.sigma_decay
# train for one epoch
train_loss, train_acc = train(train_loader, model, criterion,
optimizer, args.debug, args.flip)
# evaluate on validation set
valid_loss, valid_acc, predictions = validate(val_loader, model,
criterion, njoints,
args.debug, args.flip)
# append logger file
logger.append(
[epoch + 1, lr, train_loss, valid_loss, train_acc, valid_acc])
# remember best acc and save checkpoint
is_best = valid_acc > best_acc
best_acc = max(valid_acc, best_acc)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
},
predictions,
is_best,
checkpoint=args.checkpoint,
snapshot=args.snapshot)
logger.close()
logger.plot(['Train Acc', 'Val Acc'])
savefig(os.path.join(args.checkpoint, 'log.eps'))
def main(args):
global best_acc
# create checkpoint dir
if not isdir(args.checkpoint):
mkdir_p(args.checkpoint)
# create model
print("==> creating model '{}', stacks={}, blocks={}".format(args.arch, args.stacks, args.blocks))
model = models.__dict__[args.arch](num_stacks=args.stacks, num_blocks=args.blocks, num_classes=args.num_classes)
model = torch.nn.DataParallel(model).cuda()
# define loss function (criterion) and optimizer
criterion = torch.nn.MSELoss(size_average=True).cuda()
optimizer = torch.optim.RMSprop(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
title = 'mpii-' + args.arch
if args.resume:
if isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_acc = checkpoint['best_acc']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
logger = Logger(join(args.checkpoint, 'log.txt'), title=title, resume=True)
else:
print("=> no checkpoint found at '{}'".format(args.resume))
else:
logger = Logger(join(args.checkpoint, 'log.txt'), title=title)
logger.set_names(['Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train Acc', 'Val Acc'])
cudnn.benchmark = True
print(' Total params: %.2fM' % (sum(p.numel() for p in model.parameters())/1000000.0))
# Data loading code
train_loader = torch.utils.data.DataLoader(
datasets.Mpii('data/mpii/mpii_annotations.json', 'data/mpii/images',
sigma=args.sigma, label_type=args.label_type),
batch_size=args.train_batch, shuffle=True,
num_workers=args.workers, pin_memory=True)
val_loader = torch.utils.data.DataLoader(
datasets.Mpii('data/mpii/mpii_annotations.json', 'data/mpii/images',
sigma=args.sigma, label_type=args.label_type, train=False),
batch_size=args.test_batch, shuffle=False,
num_workers=args.workers, pin_memory=True)
if args.evaluate:
print('\nEvaluation only')
loss, acc, predictions = validate(val_loader, model, criterion, args.num_classes, args.debug, args.flip)
save_pred(predictions, checkpoint=args.checkpoint)
return
lr = args.lr
for epoch in range(args.start_epoch, args.epochs):
from time import sleep
sleep(2)
lr = adjust_learning_rate(optimizer, epoch, lr, args.schedule, args.gamma)
print('\nEpoch: %d | LR: %.8f' % (epoch + 1, lr))
# decay sigma
if args.sigma_decay > 0:
train_loader.dataset.sigma *= args.sigma_decay
val_loader.dataset.sigma *= args.sigma_decay
# train for one epoch
train_loss, train_acc = train(train_loader, model, criterion, optimizer, args.debug, args.flip)
# evaluate on validation set
valid_loss, valid_acc, predictions = validate(val_loader, model, criterion, args.num_classes,
args.debug, args.flip)
# append logger file
logger.append([epoch + 1, lr, train_loss, valid_loss, train_acc, valid_acc])
# remember best acc and save checkpoint
is_best = valid_acc > best_acc
best_acc = max(valid_acc, best_acc)
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc': best_acc,
'optimizer' : optimizer.state_dict(),
}, predictions, is_best, checkpoint=args.checkpoint)
logger.close()
logger.plot(['Train Acc', 'Val Acc'])
savefig(os.path.join(args.checkpoint, 'log.eps'))
def main(args):
num_datasets = len(args.data_dir) #number of datasets
for item in [
args.training_set_percentage, args.meta_dir, args.anno_type,
args.ratio
]:
if len(item) == 1:
for i in range(num_datasets - 1):
item.append(item[0])
assert len(item) == num_datasets
scales = [0.7, 0.85, 1, 1.3, 1.6]
if args.meta_dir == '':
args.meta_dir = args.data_dir #if not specified, assume meta info is stored in data dir.
# create checkpoint dir
if not isdir(args.checkpoint):
mkdir_p(args.checkpoint)
#create the log file not exist
file = open(join(args.checkpoint, 'log.txt'), 'w+')
file.close()
if args.evaluate: #creatng path for evaluation
if not isdir(args.save_result_dir):
mkdir_p(args.save_result_dir)
folders_to_create = ['preds', 'visualization']
if args.save_heatmap:
folders_to_create.append('heatmaps')
for folder_name in folders_to_create:
if not os.path.isdir(
os.path.join(args.save_result_dir, folder_name)):
print('creating path: ' +
os.path.join(args.save_result_dir, folder_name))
os.mkdir(os.path.join(args.save_result_dir, folder_name))
idx = range(args.num_classes)
global best_acc
cams = ['FusionCameraActor3_2']
# create model
print("==> creating model '{}', stacks={}, blocks={}".format(
args.arch, args.stacks, args.blocks))
model = models.__dict__[args.arch](num_stacks=args.stacks,
num_blocks=args.blocks,
num_classes=args.num_classes)
model = torch.nn.DataParallel(model).cuda()
# define loss function (criterion) and optimizer
criterion = torch.nn.MSELoss(size_average=True).cuda()
optimizer = torch.optim.RMSprop(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
title = 'arm-' + args.arch
if args.resume:
if isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_acc = checkpoint['best_acc']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
logger = Logger(join(args.checkpoint, 'log.txt'),
title=title,
resume=True)
else:
print("=> no checkpoint found at '{}'".format(args.resume))
else:
logger = Logger(join(args.checkpoint, 'log.txt'), title=title)
logger.set_names(
['Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train Acc', 'Val Acc'])
cudnn.benchmark = True
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
train_set_list = []
val_set_list = []
for i in range(num_datasets):
train_set_list.append(
datasets.Arm(
args.data_dir[i],
args.meta_dir[i],
args.random_bg_dir,
cams[0],
args.anno_type[i],
train=True,
training_set_percentage=args.training_set_percentage[i],
replace_bg=args.replace_bg))
val_set_list.append(
datasets.Arm(
args.data_dir[i],
args.meta_dir[i],
args.random_bg_dir,
cams[0],
args.anno_type[i],
train=False,
training_set_percentage=args.training_set_percentage[i],
scales=scales,
multi_scale=args.multi_scale,
ignore_invis_pts=args.ignore_invis_pts))
# Data loading code
if not args.evaluate:
train_loader = torch.utils.data.DataLoader(datasets.Concat(
datasets=train_set_list, ratio=args.ratio),
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
print("No. minibatches in training set:{}".format(len(train_loader)))
if args.multi_scale: #multi scale testing
args.test_batch = args.test_batch * len(scales)
val_loader = torch.utils.data.DataLoader(datasets.Concat(
datasets=val_set_list, ratio=None),
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
print("No. minibatches in validation set:{}".format(len(val_loader)))
if args.evaluate:
print('\nEvaluation only')
# if not args.compute_3d:
loss, acc = validate(val_loader, model, criterion, args.num_classes,
idx, args.save_result_dir, args.meta_dir,
args.anno_type, args.flip, args.evaluate, scales,
args.multi_scale, args.save_heatmap)
if args.compute_3d:
preds = []
gts = []
hit, d3_pred, file_name_list = d2tod3(
data_dir=args.save_result_dir,
meta_dir=args.meta_dir[0],
cam_type=args.camera_type,
pred_from_heatmap=False,
em_test=False)
# validate the 3d reconstruction accuracy
with open(os.path.join(args.save_result_dir, 'd3_pred.json'),
'r') as f:
obj = json.load(f)
hit, d3_pred, file_name_list = obj['hit'], obj['d3_pred'], obj[
'file_name_list']
for file_name in file_name_list:
preds.append(d3_pred[file_name]['preds']) #predicted x
with open(os.path.join(args.data_dir[0], 'angles', file_name),
'r') as f:
gts.append(json.load(f))
print('average error in angle: [base, elbow, ankle, wrist]:{}'.
format(d3_acc(preds, gts)))
return
lr = args.lr
for epoch in range(args.start_epoch, args.epochs):
lr = adjust_learning_rate(optimizer, epoch, lr, args.schedule,
args.gamma)
print('\nEpoch: %d | LR: %.8f' % (epoch + 1, lr))
# decay sigma
if args.sigma_decay > 0:
train_loader.dataset.sigma *= args.sigma_decay
val_loader.dataset.sigma *= args.sigma_decay
# train for one epoch
train_loss, train_acc = train(train_loader, model, criterion,
optimizer, idx, args.flip)
# evaluate on validation set
valid_loss, valid_acc = validate(val_loader, model, criterion,
args.num_classes, idx,
args.save_result_dir, args.meta_dir,
args.anno_type, args.flip,
args.evaluate)
#If concatenated dataset is used, re-random after each epoch
train_loader.dataset.reset(), val_loader.dataset.reset()
# append logger file
logger.append(
[epoch + 1, lr, train_loss, valid_loss, train_acc, valid_acc])
# remember best acc and save checkpoint
is_best = valid_acc > best_acc
best_acc = max(valid_acc, best_acc)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
},
is_best,
checkpoint=args.checkpoint)
logger.close()
logger.plot(['Train Acc', 'Val Acc'])
savefig(os.path.join(args.checkpoint, 'log.eps'))
def main(args):
global best_acc
global idx1
global idx2
# idx is the index of joints used to compute accuracy for dataset2
if args.dataset1 == 'real_animal' or args.dataset1 == 'real_animal_sp' or args.dataset1 == 'synthetic_animal' or args.dataset1 == 'synthetic_animal_sp':
idx1 = range(1, 19)
else:
print("Unknown dataset: {}".format(args.dataset1))
assert False
if args.dataset2 == 'real_animal' or args.dataset2 == 'real_animal_sp':
if args.animal == 'horse':
idx2 = range(1, 19) # horse
elif args.animal == 'tiger':
idx2 = [
1, 2, 3, 4, 5, 6, 7, 8, 15, 16, 17, 18, 13, 14, 9, 10, 11, 12
] # tiger
else:
print("Unknown dataset: {}".format(args.dataset2))
assert False
# create checkpoint dir
if not isdir(args.checkpoint):
mkdir_p(args.checkpoint)
# create model
njoints = datasets.__dict__[args.dataset1].njoints
print(njoints)
print("==> creating model '{}', stacks={}, blocks={}".format(
args.arch, args.stacks, args.blocks))
model = models.__dict__[args.arch](num_stacks=args.stacks,
num_blocks=args.blocks,
num_classes=njoints,
resnet_layers=args.resnet_layers)
model = torch.nn.DataParallel(model).to(device)
# define loss function (criterion) and optimizer
criterion = losses.JointsMSELoss().to(device)
if args.solver == 'rms':
optimizer = torch.optim.RMSprop(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.solver == 'adam':
optimizer = torch.optim.Adam(
model.parameters(),
lr=args.lr,
)
else:
print('Unknown solver: {}'.format(args.solver))
assert False
# optionally resume from a checkpoint
title = args.dataset1 + ' ' + args.arch
if args.resume:
if isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
#model.load_state_dict(checkpoint)
args.start_epoch = checkpoint['epoch']
best_acc = checkpoint['best_acc']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
else:
raise Exception('please provide a checkpoint')
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
# create data loader
train_dataset = datasets.__dict__[args.dataset1](is_train=True,
**vars(args))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_dataset = datasets.__dict__[args.dataset2](is_train=False,
**vars(args))
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# evaluation only
if args.evaluate:
print('\nEvaluation only')
_, acc, predictions = validate(val_loader, model, criterion, njoints,
args, args.flip, args.test_batch)
#save_pred(predictions, checkpoint=args.checkpoint)
return
def main(args):
global best_acc
# create checkpoint dir
if not isdir(args.checkpoint):
mkdir_p(args.checkpoint)
_logger = log.get_logger(__name__, args)
_logger.info(print_args(args))
# create model
print("==> creating model '{}', stacks={}, blocks={}".format(
args.arch, args.stacks, args.blocks))
model = models.__dict__[args.arch](num_stacks=args.stacks,
num_blocks=args.blocks,
num_classes=len(args.index_classes))
model = torch.nn.DataParallel(model).cuda()
# define loss function (criterion) and optimizer
criterion = models.loss.UniLoss(a_points=args.a_points)
optimizer = torch.optim.RMSprop(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
title = 'mpii-' + args.arch
if args.resume:
if isfile(args.resume):
_logger.info("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_acc = checkpoint['best_acc']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
for param_group in optimizer.param_groups:
param_group['lr'] = args.lr
print(param_group['lr'])
_logger.info("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
logger = Logger(join(args.checkpoint, 'log.txt'),
title=title,
resume=False)
logger.set_names([
'Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train Acc', 'Val Acc'
])
else:
_logger.info("=> no checkpoint found at '{}'".format(args.resume))
else:
logger = Logger(join(args.checkpoint, 'log.txt'), title=title)
logger.set_names(
['Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train Acc', 'Val Acc'])
cudnn.benchmark = True
_logger.info(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
# Data loading code
train_loader = torch.utils.data.DataLoader(
datasets.Mpii('data/mpii/mpii_annotations.json',
'data/mpii/images',
sigma=args.sigma,
label_type=args.label_type,
_idx=args.index_classes,
direct=True,
n_points=args.n_points),
batch_size=args.train_batch,
shuffle=True,
collate_fn=datasets.mpii.mycollate,
num_workers=args.workers,
pin_memory=False)
val_loader = torch.utils.data.DataLoader(
datasets.Mpii('data/mpii/mpii_annotations.json',
'data/mpii/images',
sigma=args.sigma,
label_type=args.label_type,
_idx=args.index_classes,
train=False,
direct=True),
batch_size=args.test_batch,
shuffle=False,
collate_fn=datasets.mpii.mycollate,
num_workers=args.workers,
pin_memory=False)
if args.evaluate:
_logger.warning('\nEvaluation only')
loss, acc, predictions = validate(val_loader,
model,
criterion,
len(args.index_classes),
False,
args.flip,
_logger,
evaluate_only=True)
save_pred(predictions, checkpoint=args.checkpoint)
return
# multi-thread
inqueues = []
outqueues = []
valid_accs = []
lr = args.lr
for epoch in range(args.start_epoch, args.epochs):
lr = adjust_learning_rate(optimizer, epoch, lr, args.schedule,
args.gamma)
_logger.warning('\nEpoch: %d | LR: %.8f' % (epoch + 1, lr))
# decay sigma
if args.sigma_decay > 0:
train_loader.dataset.sigma *= args.sigma_decay
val_loader.dataset.sigma *= args.sigma_decay
# train for one epoch
train_loss, train_acc = train(inqueues, outqueues, train_loader, model,
criterion, optimizer, args.debug,
args.flip, args.clip, _logger)
# evaluate on validation set
with torch.no_grad():
valid_loss, valid_acc, predictions = validate(
val_loader, model, criterion, len(args.index_classes),
args.debug, args.flip, _logger)
# append logger file
logger.append(
[epoch + 1, lr, train_loss, valid_loss, train_acc, valid_acc])
valid_accs.append(valid_acc)
if args.schedule[0] == -1:
if len(valid_accs) > 8:
if sum(valid_accs[-4:]) / 4 * 0.99 < sum(
valid_accs[-8:-4]) / 4:
args.schedule.append(epoch + 1)
valid_accs = []
# remember best acc and save checkpoint
is_best = valid_acc > best_acc
best_acc = max(valid_acc, best_acc)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
},
predictions,
is_best,
checkpoint=args.checkpoint,
snapshot=1)
logger.close()
logger.plot(['Train Acc', 'Val Acc'])
savefig(os.path.join(args.checkpoint, 'log.eps'))
def main(args):
global best_acc
global idx
global customMpiiObject
# idx is the index of joints used to compute accuracy
if args.dataset in ['mpii', 'lsp']:
idx = [1, 2, 3, 4, 5, 6, 11, 12, 15, 16]
elif args.dataset == 'coco':
idx = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
else:
print("Unknown dataset: {}".format(args.dataset))
assert False
# create checkpoint dir
if not isdir(args.checkpoint):
mkdir_p(args.checkpoint)
# create model
njoints = CustomMpii.njoints
print("==> creating model '{}', stacks={}, blocks={}".format(
args.arch, args.stacks, args.blocks))
model = models.__dict__[args.arch](num_stacks=args.stacks,
num_blocks=args.blocks,
num_classes=njoints,
resnet_layers=args.resnet_layers)
model = torch.nn.DataParallel(model).to(device)
# define loss function (criterion) and optimizer
criterion = losses.JointsMSELoss().to(device)
if args.solver == 'rms':
optimizer = torch.optim.RMSprop(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.solver == 'adam':
optimizer = torch.optim.Adam(
model.parameters(),
lr=args.lr,
)
else:
print('Unknown solver: {}'.format(args.solver))
assert False
# optionally resume from a checkpoint
title = args.dataset + ' ' + args.arch
if args.resume:
if isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume, map_location='cpu')
args.start_epoch = checkpoint['epoch']
best_acc = checkpoint['best_acc']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
logger = Logger(join(args.checkpoint, 'log.txt'),
title=title,
resume=True)
else:
print("=> no checkpoint found at '{}'".format(args.resume))
else:
logger = Logger(join(args.checkpoint, 'log.txt'), title=title)
logger.set_names(
['Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train Acc', 'Val Acc'])
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
print("Loading validation dataset...")
customMpiiObject = CustomMpii(is_train=False, **vars(args))
loss, acc, predictions = validate(args.unit_path, model, criterion,
njoints, args.debug, args.flip)
print(predictions)
if args.debug:
paint.main(args.unit_path, predictions.squeeze().tolist())
return
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 main(args):
global best_acc
global idx
# idx is the index of joints used to compute accuracy
if args.dataset in ['mpii', 'lsp']:
idx = [1, 2, 3, 4, 5, 6, 11, 12, 15, 16]
elif args.dataset == 'coco':
idx = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
elif args.dataset == 'homes':
idx = list(range(1, 71))
else:
print("Unknown dataset: {}".format(args.dataset))
assert False
# create checkpoint dir
if not isdir(args.checkpoint):
mkdir_p(args.checkpoint)
# create model
njoints = datasets.__dict__[args.dataset].njoints
print("==> creating model '{}', stacks={}, blocks={}".format(
args.arch, args.stacks, args.blocks))
model = models.__dict__[args.arch](
num_stacks=args.stacks,
num_blocks=args.blocks,
num_classes=16, # read as 16, the nmodified to 70
resnet_layers=args.resnet_layers)
model = torch.nn.DataParallel(model).to(device)
#print(list(model.children())); exit(1)
# define loss function (criterion) and optimizer
#criterion = losses.JointsMSELoss().to(device)
#criterion = torch.nn.BCEWithLogitsLoss().to(device)
criterion = torch.nn.MSELoss().to(device)
print('==> loss function: %s' % criterion.__str__())
if args.solver == 'rms':
optimizer = torch.optim.RMSprop(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.solver == 'adam':
optimizer = torch.optim.Adam(
model.parameters(),
lr=args.lr,
)
else:
print('Unknown solver: {}'.format(args.solver))
assert False
# optionally resume from a checkpoint
title = args.dataset + ' ' + args.arch
if args.resume:
if isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_acc = checkpoint['best_acc']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
logger = Logger(join(args.checkpoint, 'log.txt'),
title=title,
resume=True)
else:
print("=> no checkpoint found at '{}'".format(args.resume))
else:
logger = Logger(join(args.checkpoint, 'log.txt'), title=title)
logger.set_names(
['Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train Acc', 'Val Acc'])
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
# update the final score layer and fine tuning => update number of output channels to 70
#print(list(model.children())[-1].score[0]); exit(1)
model_net = list(model.children())
model_net[-1].score[0] = nn.Conv2d(256,
70,
kernel_size=(1, 1),
stride=(1, 1)).to(device)
model = nn.Sequential(*model_net)
# create data loader
train_dataset = datasets.__dict__[args.dataset](**vars(args))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_dataset = datasets.__dict__[args.dataset](is_train=False,
is_valid=True,
**vars(args))
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# evaluation only
if args.evaluate:
print('\nEvaluation only')
loss, acc, predictions = validate(val_loader, model, criterion,
njoints, args.debug, args.flip)
save_pred(predictions, checkpoint=args.checkpoint)
return
# train and eval
df_loss = pd.DataFrame()
train_epo_loss, val_epo_loss = [], []
lr = args.lr
best_vel = np.float('inf')
for epoch in range(args.start_epoch, args.epochs + args.start_epoch):
lr = adjust_learning_rate(optimizer, epoch, lr, args.schedule,
args.gamma)
print('\nEpoch: %d | LR: %.8f' % (epoch + 1, lr))
# decay sigma
if args.sigma_decay > 0:
train_loader.dataset.sigma *= args.sigma_decay
val_loader.dataset.sigma *= args.sigma_decay
# train for one epoch
train_loss, train_acc, tel = train(train_loader, model, criterion,
optimizer, args.debug, args.flip)
# evaluate on validation set
valid_loss, valid_acc, predictions, vel = validate(
val_loader, model, criterion, njoints, args.debug, args.flip)
# save epoch loss to csv
train_epo_loss += [tel]
val_epo_loss += [vel]
df_loss.assign(train=train_epo_loss,
val=val_epo_loss).to_csv('./loss.csv')
# append logger file
logger.append([
epoch + 1, lr, train_loss, valid_loss,
train_acc.item(),
valid_acc.item()
])
# remember best acc and save checkpoint
#is_best = valid_acc > best_acc
#best_acc = max(valid_acc, best_acc)
is_best = vel < best_vel
best_vel = min(best_vel, vel)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
},
predictions,
is_best,
checkpoint=args.checkpoint,
snapshot=args.snapshot)
logger.close()
def main(args):
global best_acc
global idx
# idx is the index of joints used to compute accuracy
if args.dataset in ['mpii', 'lsp']:
idx = [1, 2, 3, 4, 5, 6, 11, 12, 15, 16]
elif args.dataset == 'mscoco':
idx = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
else:
print("Unknown dataset: {}".format(args.dataset))
assert False
# create checkpoint dir
if not isdir(args.checkpoint):
mkdir_p(args.checkpoint)
# create model
njoints = datasets.__dict__[args.dataset].njoints
print("==> creating model '{}', stacks={}, blocks={}".format(
args.arch, args.stacks, args.blocks))
model = models.__dict__[args.arch](num_stacks=args.stacks,
num_blocks=args.blocks,
num_classes=njoints)
model = torch.nn.DataParallel(model).to(device)
# define loss function (criterion) and optimizer
criterion = torch.nn.MSELoss(reduction='mean').to(device)
optimizer = torch.optim.RMSprop(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
title = args.dataset + ' ' + args.arch
if args.resume:
if isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_acc = checkpoint['best_acc']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
logger = Logger(join(args.checkpoint, 'log.txt'),
title=title,
resume=True)
else:
print("=> no checkpoint found at '{}'".format(args.resume))
else:
logger = Logger(join(args.checkpoint, 'log.txt'), title=title)
logger.set_names(
['Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train Acc', 'Val Acc'])
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
# save a cleaned version of model without dict and DataParallel
clean_model = checkpoint['state_dict']
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
clean_model = OrderedDict()
if any(key.startswith('module') for key in checkpoint['state_dict']):
for k, v in checkpoint['state_dict'].items():
name = k[7:] # remove `module.`
clean_model[name] = v
# import pdb; pdb.set_trace()
# import pdb; pdb.set_trace()
# while any(key.startswith('module') for key in clean_model):
# clean_model = clean_model.module
clean_model_path = join(args.checkpoint, 'clean_model.pth.tar')
torch.save(clean_model, clean_model_path)
print('clean model saved: {}'.format(clean_model_path))
def myvalidate( model, criterion, num_classes, debug=False, flip=True):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acces = AverageMeter()
img_folder = '/data3/wzwu/dataset/my'
img_num = 1
r = 0
center1 = torch.Tensor([1281,2169])
center2 = torch.Tensor([[1281,2169]])
scale = torch.Tensor([10.0])
inp_res = 256
meanstd_file = './data/mpii/mean.pth.tar'
if isfile(meanstd_file):
meanstd = torch.load(meanstd_file)
mean = meanstd['mean']
std = meanstd['std']
input_list = []
for i in range(img_num):
img_name = str(i)+'.jpg'
img_path = os.path.join(img_folder,img_name)
print('img_path')
print(img_path)
set_trace()
img = load_image(img_path)
inp = crop(img, center1, scale, [inp_res, inp_res], rot=r)
inp = color_normalize(inp, mean, std)
input_list.append(inp)
# predictions
predictions = torch.Tensor(img_num, num_classes, 2)
# switch to evaluate mode
model.eval()
gt_win, pred_win = None, None
end = time.time()
bar = Bar('Eval ', max=img_num)
with torch.no_grad():
for i, input in enumerate(input_list):
# measure data loading time
s0, s1, s2 = input.size()
input = input.view(1, s0, s1, s2)
data_time.update(time.time() - end)
input = input.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
# generate predictions
set_trace()
preds = final_preds(score_map, center2, scale, [64, 64])
set_trace()
print('preds')
print(preds)
print('predictions')
print(predictions)
for n in range(score_map.size(0)):
predictions[i, :, :] = preds[n, :, :]
if debug:
pred_batch_img = batch_with_heatmap(input, score_map)
if not gt_win or not pred_win:
#plt.subplot(121)
#plt.subplot(122)
pred_win = plt.imshow(pred_batch_img)
else:
pred_win.set_data(pred_batch_img)
plt.pause(.05)
plt.draw()
plt.savefig('/data3/wzwu/test/'+str(i)+'.png')
# 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=img_num,
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 predictions
