import argparse
import os
import sys
import shutil
import time
import numpy as np
from PIL import Image
import json
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models
from nets.vgg_based_network import HPE_with_PIL_VGG_MSRAInit
from nets.hourglass_based_network import HPE_with_PIL_HG_MSRAInit
from utils.data_loader import LIPDataset
from utils.calc_pckh import calc_pck_lip_dataset
import utils.eval_util as eval_util
parser = argparse.ArgumentParser(description='PyTorch Human Pose Estimation with Parsing Induced Learner on LIP dataset')
parser.add_argument('--train-data', default='dataset/lip/train_images/', metavar='DIR', help='path to training dataset')
parser.add_argument('--train-pose-anno', default='dataset/lip/jsons/LIP_SP_TRAIN_annotations.json', type=str, metavar='PATH', help='path to training pose annotations')
parser.add_argument('--train-parsing-anno', default='dataset/lip/train_segmentations', metavar='DIR', help='path to training parsing annotations')
parser.add_argument('--eval-data', default='dataset/lip/val_images', metavar='DIR', help='path to eval dataset')
parser.add_argument('--eval-pose-anno', default='dataset/lip/jsons/LIP_SP_VAL_annotations.json', type=str, metavar='PATH', help='path to eval pose annotations')
parser.add_argument('--eval-parsing-anno', default='dataset/lip/val_segmentations', metavar='DIR', help='path to eval parsing annotations')
parser.add_argument('--arch', default='HG', type=str, metavar='PATH', help='Network architecture (VGG or HG (Hourglass), default: HG)')
parser.add_argument('-b', '--batch_size', default=10, type=int, metavar='N', help='mini-batch size (default: 10)')
parser.add_argument('--lr', '--learning-rate', default=0.0015, type=float, metavar='LR', help='initial learning rate')
parser.add_argument('--epochs', default=250, type=int, metavar='N', help='number of total epochs to run (default: 250)')
parser.add_argument('--snapshot-fname-prefix', default='exps/snapshots/pil_lip', type=str, metavar='PATH', help='path to snapshot')
parser.add_argument('--resume', default='', type=str, metavar='PATH', help='path to latest checkpoint (default: none)')
parser.add_argument('-j', '--workers', default=8, type=int, metavar='N', help='number of data loading workers (default: 8)')
parser.add_argument('--print-freq', '-p', default=10, type=int, metavar='N', help='print frequency (default: 10)')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N', help='manual epoch number (useful on restarts)')
parser.add_argument('--evaluate', default=False, type=bool, metavar='BOOL', help='evaluate or train')
parser.add_argument('--calc-pck', default=False, type=bool, metavar='BOOL', help='caculate PCK or not')
parser.add_argument('--pred-path', default='exps/preds/csv_results/pred_keypoints_lip.csv', type=str, metavar='PATH', help='path to save the prediction results in .csv format')
parser.add_argument('--visualization', default=False, type=bool, metavar='BOOL', help='visualizae prediction or not')
parser.add_argument('--vis-dir', default='exps/preds/vis_results', metavar='DIR', help='path to save visualization results')
best_pck = 0
pck_avg_list = []
pck_all_list = []
def main():
# Global variables
global args, best_pck, pck_avg_list, pck_all_list
args = parser.parse_args()
# Welcome msg
phase_str = '[Train and Val Phase]'
if args.evaluate:
phase_str = '[Testing Phase]'
print('Human Pose Estimation with Parsing Induced Learner: {0}'.format(phase_str))
# Create network
if args.arch == 'VGG':
hpe_with_pil_net = HPE_with_PIL_VGG_MSRAInit()
pose_net_stride = 8
elif args.arch == 'HG':
hpe_with_pil_net = HPE_with_PIL_HG_MSRAInit()
pose_net_stride = 4
else:
raise RuntimeError('Unknown network architecture!')
# Multi-GPU setting
hpe_with_pil_net = nn.DataParallel(hpe_with_pil_net).cuda()
# CUDNN setting
cudnn.benchmark = True
cudnn.enabled = True
# Optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print('=> Loading checkpoints {0}'.format(args.resume))
checkpoint = torch.load(args.resume)
hpe_with_pil_net.load_state_dict(checkpoint['state_dict'])
args.start_epoch = checkpoint['epoch']
best_pck = checkpoint['best_pck']
pck_avg_list = checkpoint['pck_avg_list']
pck_all_list = checkpoint['pck_all_list']
hpe_with_pil_net_params = hpe_with_pil_net.parameters()
else:
print('=> No checkpoint found at {0}'.format(args.resume))
hpe_with_pil_net_params = hpe_with_pil_net.parameters()
# Snapshot file names
snapshot_fname = '{0}.pth.tar'.format(args.snapshot_fname_prefix)
snapshot_best_fname = '{0}_best.pth.tar'.format(args.snapshot_fname_prefix)
# Image normalization
normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[1, 1, 1])
# Data transform
data_transform = transforms.Compose([transforms.ToTensor(), normalize,])
# LIP dataset
lip_ds = LIPDataset(args.train_data, \
args.train_pose_anno, \
args.train_parsing_anno, \
transform=data_transform, \
pose_net_stride=pose_net_stride, \
parsing_net_stride=1, \
crop_size=256, \
target_dist=1.171, scale_min=0.8, scale_max=1.5, \
max_rotate_degree=40, \
max_center_trans=40, \
flip_prob=0.5, \
is_visualization=False)
# Load training data
train_loader = torch.utils.data.DataLoader(lip_ds, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True)
# Load validation data
print('Loading evaluation json file: {0}...'.format(args.eval_pose_anno))
eval_list = []
with open(args.eval_pose_anno) as data_file:
data_this = json.load(data_file)
data_this = data_this['root']
eval_list = eval_list + data_this
eval_im_name_list = []
for ii in range(0, len(eval_list)):
eval_item = eval_list[ii]
eval_im_name_list.append(eval_item['im_name'])
print('Finished loading evaluation json file')
# MSE Loss function for pose estimation and CrossEntropy Loss function for parsing estimation
pose_criterion = nn.MSELoss().cuda()
parsing_criterion = nn.NLLLoss2d().cuda()
# RMSProp as the optimizer
optimizer = torch.optim.RMSprop(hpe_with_pil_net_params, args.lr)
# Testing
if args.evaluate == True:
evaluate(hpe_with_pil_net, \
args.eval_data, \
eval_im_name_list, \
transform=data_transform, \
stride=pose_net_stride, \
crop_size=256, \
scale_multiplier=[1], \
visualization=args.visualization, \
vis_result_dir=args.vis_dir, \
pred_path=args.pred_path, \
is_calc_pck=args.calc_pck)
return
for epoch in range(args.start_epoch, args.epochs):
# Training
train(train_loader, hpe_with_pil_net, pose_criterion, parsing_criterion, optimizer, epoch)
# Save snapshot
torch.save({
'epoch': epoch + 1,
'state_dict': hpe_with_pil_net.state_dict(),
'best_pck': best_pck,
'pck_avg_list': pck_avg_list,
'pck_all_list': pck_all_list,
}, snapshot_fname)
# Validation
if epoch < 100:
val_freq = 10
elif epoch < 150:
val_freq = 2
else:
val_freq = 1
if (epoch + 1) % val_freq == 0:
pck_avg = evaluate(hpe_with_pil_net, \
args.eval_data, \
eval_im_name_list, \
transform=data_transform, \
stride=pose_net_stride, \
crop_size=256, \
scale_multiplier=[1], \
visualization=args.visualization, \
vis_result_dir=args.vis_dir, \
pred_path=args.pred_path, \
is_calc_pck=True)
is_best = pck_avg > best_pck
best_pck = max(pck_avg, best_pck)
torch.save({
'epoch': epoch + 1,
'state_dict': hpe_with_pil_net.state_dict(),
'best_pck': best_pck,
'pck_avg_list': pck_avg_list,
'pck_all_list': pck_all_list,
}, snapshot_fname)
if is_best:
shutil.copyfile(snapshot_fname,snapshot_best_fname)
def train(train_loader, model, pose_criterion, parsing_criterion, optimizer, epoch):
cur_lr = adjust_learning_rate(optimizer, epoch)
losses = AverageMeter()
cost_time = AverageMeter()
train_acc = AverageMeter()
model.train()
iter_start_time = time.time()
for i, (im, pose_target, parsing_target) in enumerate(train_loader):
# Prepare input and target variables
im = im.cuda(async=True)
pose_target = pose_target.float().cuda(async=True)
parsing_target = parsing_target.long().cuda(async=True)
input_var = torch.autograd.Variable(im)
pose_target_var = torch.autograd.Variable(pose_target)
parsing_target_var = torch.autograd.Variable(parsing_target)
# Network forward
pose_output, parsing_output = model(input_var)
# Calculate parsing loss
total_loss = 0.01 * parsing_criterion(parsing_output, parsing_target_var)
# Calculate pose loss
# Case 1: pose output is a list from Hourglass network
# Case 2: pose output is a tensor from VGG network
if isinstance(pose_output, list):
avg_acc = cal_train_acc(pose_output[-1].data, pose_target)
for s in range(0, len(pose_output)):
pose_loss = pose_criterion(pose_output[s], pose_target_var)
total_loss += pose_loss
else:
avg_acc = cal_train_acc(pose_output.data, pose_target)
pose_loss = pose_criterion(pose_output, pose_target_var)
total_loss += pose_loss
train_acc.update(avg_acc, 1)
losses.update(total_loss.data[0], im.size(0))
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
cost_time.update(time.time() - iter_start_time)
iter_start_time = time.time()
if i == 0 or (i + 1) % args.print_freq == 0:
print('Epoch: [{0}][{1}/{2}] \t'
'CurLR: {3} \t'
'Loss {loss.val:.6f} ({loss.avg:.6f}) \t'
'Acc {accuracy.val:.3f} ({accuracy.avg:.3f}) \t'
'BatchTime {cost_time.val:.3f} ({cost_time.avg:.3f}) \t'.format(
epoch + 1, i + 1, len(train_loader),
cur_lr,
loss=losses,
accuracy=train_acc,
cost_time=cost_time))
def evaluate(model, \
eval_im_root_dir, \
eval_im_name_list, \
transform=None, \
stride=4, \
crop_size=256, \
scale_multiplier=[1], \
num_of_joints=16, \
visualization=False, \
vis_result_dir='exps/preds/vis_results', \
gt_path='dataset/lip/val_gt/lip_val_groundtruth.csv', \
pred_path='exps/preds/csv_results/pred_keypoints_lip.csv', \
is_calc_pck=True):
model.eval()
pose_list = eval_util.multi_image_testing_on_lip_dataset(model, \
eval_im_root_dir, \
eval_im_name_list, \
transform=transform, \
stride=stride, \
crop_size=crop_size, \
scale_multiplier=scale_multiplier, \
num_of_joints=num_of_joints, \
visualization=visualization, \
vis_result_dir=vis_result_dir)
eval_util.save_hpe_results_to_lip_format(eval_im_name_list, pose_list, save_path=pred_path)
pck_avg = 0.0
if is_calc_pck:
pck_all = calc_pck_lip_dataset(gt_path, pred_path, method_name='hpe_with_pil', eval_num=len(eval_im_name_list))
pck_avg = pck_all[-1][-1]
pck_all_list.append(pck_all)
pck_avg_list.append(pck_avg)
return pck_avg
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
def adjust_learning_rate(optimizer, epoch):
decay = 0
if epoch + 1 >= 230:
decay = 0.05
elif epoch + 1 >= 200:
decay = 0.1
elif epoch + 1 >= 170:
decay = 0.25
elif epoch + 1 >= 150:
decay = 0.5
else:
decay = 1
lr = args.lr * decay
for param_group in optimizer.param_groups:
param_group['lr'] = lr
return lr
# Get predictions
def get_preds(heatmaps):
if heatmaps.dim() != 4:
raise ValueError('Input must be 4-D tensor')
max_val, max_idx = torch.max(heatmaps.view(heatmaps.size(0), heatmaps.size(1), heatmaps.size(2) * heatmaps.size(3)), 2)
preds = torch.Tensor(max_idx.size(0), max_idx.size(1), 2)
preds[:, :, 0] = max_idx[:, :] % heatmaps.size(3)
preds[:, :, 1] = max_idx[:, :] / heatmaps.size(3)
preds[:, :, 1] = preds[:, :, 1].floor()
return preds
def calc_dists(preds, labels, normalize):
dists = torch.Tensor(preds.size(1), preds.size(0))
for i in range(preds.size(0)):
for j in range(preds.size(1)):
if labels[i, j, 0] == 0 and labels[i, j, 1] == 0:
dists[j, i] = -1
else:
dists[j, i] = torch.dist(labels[i, j, :], preds[i, j, :]) / normalize
return dists
def dist_accuracy(dists, th=0.5):
if torch.ne(dists, -1).sum() > 0:
return (dists.le(th).eq(dists.ne(-1)).sum()) * 1.0 / dists.ne(-1).sum()
else:
return -1
def cal_train_acc(output, target):
num_of_joints = target.size(1) - 1
preds = get_preds(output)
gt = get_preds(target)
dists = calc_dists(preds, gt, output.size(3) / 10.0)
avg_acc = 0.0
bad_idx_count = 0
for ji in range(num_of_joints):
acc = dist_accuracy(dists[ji, :])
if acc > 0:
avg_acc += acc
else:
bad_idx_count += 1
if bad_idx_count != num_of_joints:
avg_acc = avg_acc / (num_of_joints - bad_idx_count)
return avg_acc
if __name__ == '__main__':
main()
