def test_do_forward(device, example_input):
model = hg2(pretrained=True)
predictor = HumanPosePredictor(model, device=device)
output = predictor.do_forward(example_input.to(device))
assert len(output) == 2 # Expect one set of heatmaps per stack.
heatmaps = output[-1]
assert heatmaps.shape == (1, 16, 64, 64)
def test_estimate_joints_with_flip(device, man_running_image,
man_running_pose):
model = hg2(pretrained=True)
predictor = HumanPosePredictor(model, device=device)
joints = predictor.estimate_joints(man_running_image, flip=True)
assert joints.shape == (16, 2)
assert_allclose(joints, man_running_pose, rtol=0, atol=20)
def test_estimate_joints_h36m(device, h36m_image, h36m_pose):
device = torch.device(device)
model = hg2(pretrained=True)
predictor = HumanPosePredictor(model, device=device)
joints = predictor.estimate_joints(h36m_image)
assert joints.shape == (16, 2)
assert_allclose(joints, h36m_pose, rtol=0, atol=15)
def test_do_validation_step(device):
model = hg2(pretrained=False)
model = model.to(device)
model.eval()
inp = torch.randn((1, 3, 256, 256), device=device)
target = torch.randn((1, 16, 64, 64), device=device)
output, loss = do_validation_step(model, inp, target, Mpii.DATA_INFO)
assert output.shape == (1, 16, 64, 64)
assert loss > 0
def test_prepare_image(device, man_running_image):
model = hg2(pretrained=True)
predictor = HumanPosePredictor(model, device=device)
orig_image = man_running_image.clone()
image = predictor.prepare_image(orig_image)
assert_allclose(orig_image,
man_running_image) # Input image should be unchanged.
assert image.shape == (3, 256, 256)
assert image.device.type == 'cpu'
def main(args):
# Select the hardware device to use for inference.
if torch.cuda.is_available():
device = torch.device('cuda', torch.cuda.current_device())
torch.backends.cudnn.benchmark = True
else:
device = torch.device('cpu')
# Disable gradient calculations.
torch.set_grad_enabled(False)
pretrained = not args.model_file
if pretrained:
print(
'No model weights file specified, using pretrained weights instead.'
)
# Create the model, downloading pretrained weights if necessary.
if args.arch == 'hg1':
model = hg1(pretrained=pretrained)
elif args.arch == 'hg2':
model = hg2(pretrained=pretrained)
elif args.arch == 'hg8':
model = hg8(pretrained=pretrained)
else:
raise Exception('unrecognised model architecture: ' + args.model)
model = model.to(device)
if not pretrained:
assert os.path.isfile(args.model_file)
print('Loading model weights from file: {}'.format(args.model_file))
checkpoint = torch.load(args.model_file)
state_dict = checkpoint['state_dict']
if sorted(state_dict.keys())[0].startswith('module.'):
model = DataParallel(model)
model.load_state_dict(state_dict)
# Initialise the MPII validation set dataloader.
# val_dataset = Mpii(args.image_path, is_train=False)
# val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False,
# num_workers=args.workers, pin_memory=True)
# Generate predictions for the validation set.
# _, _, predictions = do_validation_epoch(val_loader, model, device, Mpii.DATA_INFO, args.flip)
model = hg1(pretrained=True)
predictor = HumanPosePredictor(model, device='cpu')
# my_image = image_loader("../inference-img/1.jpg")
# joints = image_inference(predictor, image_path=None, my_image=my_image)
# imshow(my_image, joints=joints)
if args.camera == False:
inference_video(predictor, "../inference-video/R6llTwEh07w.mp4")
elif args.camera:
inference_video(predictor, 0)
def test_estimate_joints_tensor_batch(device, h36m_image, h36m_pose):
model = hg2(pretrained=True)
predictor = HumanPosePredictor(model, device=device)
batch_size = 4
joints = predictor.estimate_joints(h36m_image.repeat(batch_size, 1, 1, 1))
assert joints.shape == (batch_size, 16, 2)
assert_allclose(joints,
h36m_pose.repeat(batch_size, 1, 1),
rtol=0,
atol=15)
def test_prepare_image_aspect_ratio(device, dummy_data_info):
orig_image = torch.ones((3, 256, 512), dtype=torch.float32, device=device)
model = hg2(pretrained=True)
predictor = HumanPosePredictor(model,
device=device,
data_info=dummy_data_info)
image = predictor.prepare_image(orig_image)
expected = torch.zeros((3, 256, 256), dtype=torch.float32, device=device)
expected[:, 64:192] = 1.0
assert_allclose(image, expected)
def test_asymmetric_input(device, man_running_image):
model = hg2(pretrained=True)
predictor = HumanPosePredictor(model, device=device, input_shape=(512, 64))
orig_image = man_running_image.clone()
image = predictor.prepare_image(orig_image)
assert image.shape == (3, 512, 64)
heatmaps = predictor.estimate_heatmaps(image)
assert heatmaps.shape == (16, 128, 16)
joints = predictor.estimate_joints(image)
assert all(joints[:, 0] < 64)
assert all(joints[:, 1] < 512)
def test_do_training_step(device):
device = torch.device(device)
model = hg2(pretrained=False)
model = model.to(device)
model.train()
optimiser = Adam(model.parameters())
inp = torch.randn((1, 3, 256, 256), device=device)
target = torch.randn((1, 16, 64, 64), device=device)
output, loss = do_training_step(model, optimiser, inp, target)
assert output.shape == (1, 16, 64, 64)
assert loss > 0
def main(args):
# Select the hardware device to use for inference.
if torch.cuda.is_available():
device = torch.device('cuda', torch.cuda.current_device())
torch.backends.cudnn.benchmark = True
else:
device = torch.device('cpu')
# Disable gradient calculations.
torch.set_grad_enabled(False)
pretrained = not args.model_file
if pretrained:
print(
'No model weights file specified, using pretrained weights instead.'
)
# Create the model, downloading pretrained weights if necessary.
if args.arch == 'hg1':
model = hg1(pretrained=pretrained)
elif args.arch == 'hg2':
model = hg2(pretrained=pretrained)
elif args.arch == 'hg8':
model = hg8(pretrained=pretrained)
else:
raise Exception('unrecognised model architecture: ' + args.model)
model = model.to(device)
if not pretrained:
assert os.path.isfile(args.model_file)
print('Loading model weights from file: {}'.format(args.model_file))
checkpoint = torch.load(args.model_file)
state_dict = checkpoint['state_dict']
if sorted(state_dict.keys())[0].startswith('module.'):
model = DataParallel(model)
model.load_state_dict(state_dict)
# Initialise the MPII validation set dataloader.
val_dataset = Mpii(args.image_path, is_train=False)
val_loader = DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# Generate predictions for the validation set.
_, _, predictions = do_validation_epoch(val_loader, model, device,
Mpii.DATA_INFO, args.flip)
# Report PCKh for the predictions.
print('\nFinal validation PCKh scores:\n')
print_mpii_validation_accuracy(predictions)
def test_reorder_hg_outputs(device):
in_data = torch.randn((4, 3, 256, 256), dtype=torch.float32, device=device)
model = hg2(pretrained=True).to(device)
orig_output = model(in_data)
output_indices = list(reversed(range(16)))
change_hg_outputs(model, output_indices)
new_output = model(in_data)
for orig_stage_output, new_stage_output in zip(orig_output, new_output):
assert_allclose(new_stage_output, orig_stage_output.flip(1))
assert_allclose(new_stage_output, orig_stage_output.flip(1))
def test_prepare_image_mostly_ready(device):
# This test is for preparing an image which already has the correct dtype and size.
image_float32 = torch.empty((3, 256, 256),
device=device,
dtype=torch.float32).uniform_()
model = hg2(pretrained=True)
predictor = HumanPosePredictor(model, device=device)
orig_image = image_float32.clone()
image = predictor.prepare_image(orig_image)
assert_allclose(image_float32,
orig_image) # Input image should be unchanged.
assert image.shape == (3, 256, 256)
assert image.device == device
def test_estimate_joints_fit_contain(device, man_running_image,
man_running_pose):
# Crop the example so that it is no longer square.
narrow_width = 256
image = fit(man_running_image, (512, narrow_width), fit_mode='cover')
gt_joints = man_running_pose.clone()
gt_joints[..., 0] -= (512 - narrow_width) / 2
# Run inference, enforcing square input.
model = hg2(pretrained=True)
predictor = HumanPosePredictor(model,
device=device,
input_shape=(256, 256))
joints = predictor.estimate_joints(image)
# Check that the results are as expected.
assert joints.shape == (16, 2)
assert_allclose(joints, gt_joints, rtol=0, atol=20)
def test_change_hg_outputs(device):
in_data = torch.randn((4, 3, 256, 256), dtype=torch.float32, device=device)
model = hg2(pretrained=True).to(device)
orig_output = model(in_data)
# New output, left ankle, right ankle, new output, new output.
output_indices = [None, 5, 0, None, None]
change_hg_outputs(model, output_indices)
new_output = model(in_data)
orig_first_stage_output = orig_output[0]
new_first_stage_output = new_output[0]
assert_allclose(new_first_stage_output[:, 1], orig_first_stage_output[:,
5])
assert_allclose(new_first_stage_output[:, 2], orig_first_stage_output[:,
0])
def main(args):
"""Train/ Cross validate for data source = YogiDB."""
# Create data loader
"""Generic(data.Dataset)(image_set, annotations,
is_train=True, inp_res=256, out_res=64, sigma=1,
scale_factor=0, rot_factor=0, label_type='Gaussian',
rgb_mean=RGB_MEAN, rgb_stddev=RGB_STDDEV)."""
annotations_source = 'basic-thresholder'
# Get the data from yogi
db_obj = YogiDB(config.db_url)
imageset = db_obj.get_filtered(ImageSet,
name=args.image_set_name)
annotations = db_obj.get_annotations(image_set_name=args.image_set_name,
annotation_source=annotations_source)
pts = torch.Tensor(annotations[0]['joint_self'])
num_classes = pts.size(0)
crop_size = 512
if args.crop:
crop_size = args.crop
crop = True
else:
crop = False
# Using the default RGB mean and std dev as 0
RGB_MEAN = torch.as_tensor([0.0, 0.0, 0.0])
RGB_STDDEV = torch.as_tensor([0.0, 0.0, 0.0])
dataset = Generic(image_set=imageset,
inp_res=args.inp_res,
out_res=args.out_res,
annotations=annotations,
mode=args.mode,
crop=crop, crop_size=crop_size,
rgb_mean=RGB_MEAN, rgb_stddev=RGB_STDDEV)
train_dataset = dataset
train_dataset.is_train = True
train_loader = DataLoader(train_dataset,
batch_size=args.train_batch, shuffle=True,
num_workers=args.workers, pin_memory=True)
val_dataset = dataset
val_dataset.is_train = False
val_loader = DataLoader(val_dataset,
batch_size=args.test_batch, shuffle=False,
num_workers=args.workers, pin_memory=True)
# Select the hardware device to use for inference.
if torch.cuda.is_available():
device = torch.device('cuda', torch.cuda.current_device())
torch.backends.cudnn.benchmark = True
else:
device = torch.device('cpu')
# Disable gradient calculations by default.
torch.set_grad_enabled(False)
# create checkpoint dir
os.makedirs(args.checkpoint, exist_ok=True)
if args.arch == 'hg1':
model = hg1(pretrained=False, num_classes=num_classes)
elif args.arch == 'hg2':
model = hg2(pretrained=False, num_classes=num_classes)
elif args.arch == 'hg8':
model = hg8(pretrained=False, num_classes=num_classes)
else:
raise Exception('unrecognised model architecture: ' + args.model)
model = DataParallel(model).to(device)
if args.optimizer == "Adam":
optimizer = Adam(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
else:
optimizer = RMSprop(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
best_acc = 0
# optionally resume from a checkpoint
title = args.data_identifier + ' ' + args.arch
if args.resume:
assert os.path.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(os.path.join(args.checkpoint, 'log.txt'), title=title, resume=True)
else:
logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title)
logger.set_names(['Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train Acc', 'Val Acc'])
# 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))
# train for one epoch
train_loss, train_acc = do_training_epoch(train_loader, model, device, optimizer)
# evaluate on validation set
if args.debug == 1:
valid_loss, valid_acc, predictions, validation_log = do_validation_epoch(val_loader, model, device, False, True, os.path.join(args.checkpoint, 'debug.csv'), epoch + 1)
else:
valid_loss, valid_acc, predictions, _ = do_validation_epoch(val_loader, model, device, False)
# 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):
# Select the hardware device to use for inference.
if torch.cuda.is_available():
device = torch.device('cuda', torch.cuda.current_device())
torch.backends.cudnn.benchmark = True
else:
device = torch.device('cpu')
# Disable gradient calculations by default.
torch.set_grad_enabled(False)
# create checkpoint dir
os.makedirs(args.checkpoint, exist_ok=True)
if args.arch == 'hg1':
model = hg1(pretrained=False)
elif args.arch == 'hg2':
model = hg2(pretrained=False)
elif args.arch == 'hg8':
model = hg8(pretrained=False)
else:
raise Exception('unrecognised model architecture: ' + args.arch)
model = DataParallel(model).to(device)
optimizer = RMSprop(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
best_acc = 0
# optionally resume from a checkpoint
if args.resume:
assert os.path.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(os.path.join(args.checkpoint, 'log.txt'), resume=True)
else:
logger = Logger(os.path.join(args.checkpoint, 'log.txt'))
logger.set_names(
['Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train Acc', 'Val Acc'])
# create data loader
train_dataset = Mpii(args.image_path, is_train=True)
train_loader = DataLoader(train_dataset,
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_dataset = Mpii(args.image_path, is_train=False)
val_loader = DataLoader(val_dataset,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# train and eval
lr = args.lr
for epoch in trange(args.start_epoch,
args.epochs,
desc='Overall',
ascii=True):
lr = adjust_learning_rate(optimizer, epoch, lr, args.schedule,
args.gamma)
# train for one epoch
train_loss, train_acc = do_training_epoch(train_loader,
model,
device,
Mpii.DATA_INFO,
optimizer,
acc_joints=Mpii.ACC_JOINTS)
# evaluate on validation set
valid_loss, valid_acc, predictions = do_validation_epoch(
val_loader,
model,
device,
Mpii.DATA_INFO,
False,
acc_joints=Mpii.ACC_JOINTS)
# print metrics
tqdm.write(
f'[{epoch + 1:3d}/{args.epochs:3d}] lr={lr:0.2e} '
f'train_loss={train_loss:0.4f} train_acc={100 * train_acc:0.2f} '
f'valid_loss={valid_loss:0.4f} valid_acc={100 * valid_acc:0.2f}')
# append logger file
logger.append(
[epoch + 1, lr, train_loss, valid_loss, train_acc, valid_acc])
logger.plot_to_file(os.path.join(args.checkpoint, 'log.svg'),
['Train Acc', 'Val 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()
def main(args):
"""Predict for data source = YogiDB."""
# Select the hardware device to use for inference.
if torch.cuda.is_available():
device = torch.device('cuda', torch.cuda.current_device())
torch.backends.cudnn.benchmark = True
else:
device = torch.device('cpu')
# Disable gradient calculations.
torch.set_grad_enabled(False)
pretrained = not args.model_file
if pretrained:
print('No model weights file specified, exiting.')
exit()
# Get the data from yogi
db_obj = YogiDB(config.db_url)
imageset = db_obj.get_filtered(ImageSet, name=args.image_set_name)
annotations_source = 'basic-thresholder'
annotations = db_obj.get_annotations(annotation_source=annotations_source)
pts = torch.Tensor(annotations[0]['joint_self'])
num_classes = pts.size(0)
# Initialise the Yogi validation set dataloader.
val_dataset = Generic(image_set=imageset,
inp_res=args.inp_res,
out_res=args.out_res,
annotations=annotations,
mode=args.mode)
val_loader = DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# Create the model, downloading pretrained weights if necessary.
if args.arch == 'hg1':
model = hg1(pretrained=False, num_classes=num_classes)
elif args.arch == 'hg2':
model = hg2(pretrained=False, num_classes=num_classes)
elif args.arch == 'hg8':
model = hg8(pretrained=False, num_classes=num_classes)
else:
raise Exception('unrecognised model architecture: ' + args.model)
model = model.to(device)
# create output dir
os.makedirs(args.output_location, exist_ok=True)
title = args.image_set_name + ' ' + args.arch
filename_pre = title + '_preds_' + datetime.now().strftime("%Y%m%dT%H%M%S")
log_filepath = os.path.join(args.output_location,
filename_pre + '_log.txt')
logger = Logger(log_filepath, title=title)
logger.set_names(['Val Loss', 'Val Acc'])
# Generate predictions for the validation set.
valid_loss, valid_acc, predictions = do_validation_epoch(
val_loader, model, device, args.flip)
# append logger file
logger.append([valid_loss, valid_acc])
# TODO: Report PCKh for the predictions.
# print('\nFinal validation PCKh scores:\n')
# print_mpii_validation_accuracy(predictions)
predictions = to_numpy(predictions)
filepath = os.path.join(args.output_location, filename_pre + '.mat')
scipy.io.savemat(filepath, mdict={'preds': predictions})
from stacked_hourglass import HumanPosePredictor, hg2
from stacked_hourglass.utils.transforms import shufflelr, crop, color_normalize, fliplr, transform
from PIL import Image
import torch
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import collections as mc
import time, cv2, json, os
# ...load image of a person into a PyTorch tensor...
name = "pullup"
model = hg2(pretrained=True)
predictor = HumanPosePredictor(model, device='cpu')
print("==model loaded==")
RGB_MEAN = torch.as_tensor([0.4404, 0.4440, 0.4327])
RGB_STDDEV = torch.as_tensor([0.2458, 0.2410, 0.2468])
images = os.listdir("./video/" + name)
print("frames : ", len(images))
result = {"name": name, "frames": dict()}
for i in images:
orgImg = Image.open("./video/" + name + "/" + i)
im = np.asarray(orgImg)
idx = int(i[:-4])
def test_estimate_heatmaps(device, man_running_image):
model = hg2(pretrained=True)
predictor = HumanPosePredictor(model, device=device)
heatmaps = predictor.estimate_heatmaps(man_running_image)
assert heatmaps.shape == (16, 64, 64)
