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)