From b54221045f2fd0586baf5b19c353ba6103499c84 Mon Sep 17 00:00:00 2001 From: andretti Date: Mon, 26 Aug 2019 16:05:50 +0300 Subject: [PATCH] Added support for raw dataset inference --- .../mpii/hrnet/w32_256x256_adam_lr1e-3.yaml | 6 +- lib/config/default.py | 6 +- lib/core/function.py | 70 +++++++++++++------ lib/dataset/JointsDataset.py | 4 ++ lib/dataset/mpii.py | 29 ++++++-- lib/utils/vis.py | 4 +- requirements.txt | 1 - tools/test.py | 16 +++-- 8 files changed, 96 insertions(+), 40 deletions(-) diff --git a/experiments/mpii/hrnet/w32_256x256_adam_lr1e-3.yaml b/experiments/mpii/hrnet/w32_256x256_adam_lr1e-3.yaml index 894a8445..a866c1c2 100644 --- a/experiments/mpii/hrnet/w32_256x256_adam_lr1e-3.yaml +++ b/experiments/mpii/hrnet/w32_256x256_adam_lr1e-3.yaml @@ -4,10 +4,10 @@ CUDNN: DETERMINISTIC: false ENABLED: true DATA_DIR: '' -GPUS: (0,1,2,3) +GPUS: (0,) OUTPUT_DIR: 'output' LOG_DIR: 'log' -WORKERS: 24 +WORKERS: 6 PRINT_FREQ: 100 DATASET: @@ -107,7 +107,7 @@ TRAIN: MOMENTUM: 0.9 NESTEROV: false TEST: - BATCH_SIZE_PER_GPU: 32 + BATCH_SIZE_PER_GPU: 128 MODEL_FILE: '' FLIP_TEST: true POST_PROCESS: true diff --git a/lib/config/default.py b/lib/config/default.py index 72d3faf3..a7855ab0 100644 --- a/lib/config/default.py +++ b/lib/config/default.py @@ -20,7 +20,7 @@ _C.LOG_DIR = '' _C.DATA_DIR = '' _C.GPUS = (0,) -_C.WORKERS = 4 +_C.WORKERS = 1 _C.PRINT_FREQ = 20 _C.AUTO_RESUME = False _C.PIN_MEMORY = True @@ -89,14 +89,14 @@ _C.TRAIN.RESUME = False _C.TRAIN.CHECKPOINT = '' -_C.TRAIN.BATCH_SIZE_PER_GPU = 32 +_C.TRAIN.BATCH_SIZE_PER_GPU = 4 _C.TRAIN.SHUFFLE = True # testing _C.TEST = CN() # size of images for each device -_C.TEST.BATCH_SIZE_PER_GPU = 32 +_C.TEST.BATCH_SIZE_PER_GPU = 8 # Test Model Epoch _C.TEST.FLIP_TEST = False _C.TEST.POST_PROCESS = False diff --git a/lib/core/function.py b/lib/core/function.py index dadff0fa..8e3bb482 100755 --- a/lib/core/function.py +++ b/lib/core/function.py @@ -14,6 +14,7 @@ import numpy as np import torch +import cv2 from core.evaluate import accuracy from core.inference import get_final_preds @@ -94,8 +95,23 @@ def train(config, train_loader, model, criterion, optimizer, epoch, prefix) +def compute_joints(batch_image, batch_joints, batch_joints_vis): + for k in range(batch_image.size(0)): + image_tensor = batch_image[k] + image = image_tensor.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy() + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + + joints = batch_joints[k] + joints_vis = batch_joints_vis[k] + for joint in joints: + cv2.circle(image, (int(joint[0]), int(joint[1])), 2, [0, 0, 255], 2) + + cv2.imshow("im", image) + cv2.waitKey() + + def validate(config, val_loader, val_dataset, model, criterion, output_dir, - tb_log_dir, writer_dict=None): + tb_log_dir, writer_dict=None, predict_only=False): batch_time = AverageMeter() losses = AverageMeter() acc = AverageMeter() @@ -116,6 +132,8 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir, with torch.no_grad(): end = time.time() for i, (input, target, target_weight, meta) in enumerate(val_loader): + img = input.data[0].mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy() + # compute output outputs = model(input) if isinstance(outputs, list): @@ -147,12 +165,12 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir, output = (output + output_flipped) * 0.5 + num_images = input.size(0) + target = target.cuda(non_blocking=True) target_weight = target_weight.cuda(non_blocking=True) loss = criterion(output, target, target_weight) - - num_images = input.size(0) # measure accuracy and record loss losses.update(loss.item(), num_images) _, avg_acc, cnt, pred = accuracy(output.cpu().numpy(), @@ -181,6 +199,7 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir, image_path.extend(meta['image']) idx += num_images + compute_joints(input, pred*4, meta['joints_vis']) if i % config.PRINT_FREQ == 0: msg = 'Test: [{0}/{1}]\t' \ @@ -196,18 +215,20 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir, ) save_debug_images(config, input, meta, target, pred*4, output, prefix) + name_values = None + perf_indicator = None + if not predict_only: + name_values, perf_indicator = val_dataset.evaluate( + config, all_preds, output_dir, all_boxes, image_path, + filenames, imgnums + ) - name_values, perf_indicator = val_dataset.evaluate( - config, all_preds, output_dir, all_boxes, image_path, - filenames, imgnums - ) - - model_name = config.MODEL.NAME - if isinstance(name_values, list): - for name_value in name_values: - _print_name_value(name_value, model_name) - else: - _print_name_value(name_values, model_name) + model_name = config.MODEL.NAME + if isinstance(name_values, list): + for name_value in name_values: + _print_name_value(name_value, model_name) + else: + _print_name_value(name_values, model_name) if writer_dict: writer = writer_dict['writer'] @@ -222,19 +243,22 @@ def validate(config, val_loader, val_dataset, model, criterion, output_dir, acc.avg, global_steps ) - if isinstance(name_values, list): - for name_value in name_values: + + if not predict_only: + if isinstance(name_values, list): + for name_value in name_values: + writer.add_scalars( + 'valid', + dict(name_value), + global_steps + ) + else: writer.add_scalars( 'valid', - dict(name_value), + dict(name_values), global_steps ) - else: - writer.add_scalars( - 'valid', - dict(name_values), - global_steps - ) + writer_dict['valid_global_steps'] = global_steps + 1 return perf_indicator diff --git a/lib/dataset/JointsDataset.py b/lib/dataset/JointsDataset.py index 5a8cc3a7..a63aedb5 100755 --- a/lib/dataset/JointsDataset.py +++ b/lib/dataset/JointsDataset.py @@ -174,6 +174,10 @@ def __getitem__(self, idx): if self.transform: input = self.transform(input) + input = cv2.resize(data_numpy, (256, 256), interpolation=cv2.INTER_LINEAR) + input = cv2.normalize(input, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F) + input = torch.from_numpy(input).permute(2, 0, 1).float() + for i in range(self.num_joints): if joints_vis[i, 0] > 0.0: joints[i, 0:2] = affine_transform(joints[i, 0:2], trans) diff --git a/lib/dataset/mpii.py b/lib/dataset/mpii.py index c935001a..6b2fc682 100644 --- a/lib/dataset/mpii.py +++ b/lib/dataset/mpii.py @@ -23,7 +23,7 @@ class MPIIDataset(JointsDataset): - def __init__(self, cfg, root, image_set, is_train, transform=None): + def __init__(self, cfg, root, image_set, is_train, transform=None, unannot_imgs_path=None): super().__init__(cfg, root, image_set, is_train, transform) self.num_joints = 16 @@ -33,10 +33,13 @@ def __init__(self, cfg, root, image_set, is_train, transform=None): self.upper_body_ids = (7, 8, 9, 10, 11, 12, 13, 14, 15) self.lower_body_ids = (0, 1, 2, 3, 4, 5, 6) - self.db = self._get_db() + if unannot_imgs_path is None: + self.db = self._get_db() - if is_train and cfg.DATASET.SELECT_DATA: - self.db = self.select_data(self.db) + if is_train and cfg.DATASET.SELECT_DATA: + self.db = self.select_data(self.db) + else: + self.db = self._get_db_raw(unannot_imgs_path) logger.info('=> load {} samples'.format(len(self.db))) @@ -93,6 +96,24 @@ def _get_db(self): return gt_db + def _get_db_raw(self, unannot_imgs_path): + gt_db = [] + + for image_name in os.listdir(unannot_imgs_path): + gt_db.append( + { + 'image': os.path.join(unannot_imgs_path, image_name), + 'center': np.array([0, 0]), + 'scale': np.array([1., 1.]), + 'joints_3d': np.ones((16, 3)), + 'joints_3d_vis': np.ones((16, 3)), + 'filename': '', + 'imgnum': 0, + } + ) + + return gt_db + def evaluate(self, cfg, preds, output_dir, *args, **kwargs): # convert 0-based index to 1-based index preds = preds[:, :, 0:2] + 1.0 diff --git a/lib/utils/vis.py b/lib/utils/vis.py index adc0947a..4ee65cf5 100755 --- a/lib/utils/vis.py +++ b/lib/utils/vis.py @@ -43,8 +43,8 @@ def save_batch_image_with_joints(batch_image, batch_joints, batch_joints_vis, joints_vis = batch_joints_vis[k] for joint, joint_vis in zip(joints, joints_vis): - joint[0] = x * width + padding + joint[0] - joint[1] = y * height + padding + joint[1] + # joint[0] = x * width + padding + joint[0] + # joint[1] = y * height + padding + joint[1] if joint_vis[0]: cv2.circle(ndarr, (int(joint[0]), int(joint[1])), 2, [255, 0, 0], 2) k = k + 1 diff --git a/requirements.txt b/requirements.txt index 18c9ee11..b3d7aae9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,4 @@ EasyDict==1.7 -opencv-python==3.4.1.15 shapely==1.6.4 Cython scipy diff --git a/tools/test.py b/tools/test.py index cfa333f5..73dcb3ee 100755 --- a/tools/test.py +++ b/tools/test.py @@ -21,7 +21,7 @@ import torch.utils.data.distributed import torchvision.transforms as transforms -import _init_paths +#import _init_paths from config import cfg from config import update_config from core.loss import JointsMSELoss @@ -62,6 +62,12 @@ def parse_args(): type=str, default='') + parser.add_argument('--unnanotImgsPath', + help='annotations availability', + type=str, + default='') + # default='/home/andrettin/Repo/POSE_ESTIMATION/deep-high-resolution-net.pytorch/data/mpii/images') + args = parser.parse_args() return args @@ -95,7 +101,7 @@ def main(): logger.info('=> loading model from {}'.format(model_state_file)) model.load_state_dict(torch.load(model_state_file)) - model = torch.nn.DataParallel(model, device_ids=cfg.GPUS).cuda() + model = torch.nn.DataParallel(model, device_ids=[0]).cuda() # define loss function (criterion) and optimizer criterion = JointsMSELoss( @@ -111,7 +117,8 @@ def main(): transforms.Compose([ transforms.ToTensor(), normalize, - ]) + ]), + args.unnanotImgsPath if args.unnanotImgsPath != '' else None ) valid_loader = torch.utils.data.DataLoader( valid_dataset, @@ -123,7 +130,8 @@ def main(): # evaluate on validation set validate(cfg, valid_loader, valid_dataset, model, criterion, - final_output_dir, tb_log_dir) + final_output_dir, tb_log_dir, + predict_only=True if args.unnanotImgsPath != '' else False) if __name__ == '__main__':