Skip to content

Commit

Permalink
Topdown multi batch inference (open-mmlab#390)
Browse files Browse the repository at this point in the history 
* batch_inference

* batch_inference

* batch_inference for topdown

* topdown_multi_batch_inference

* topdown_multi_batch_inference

* topdown_multi_batch_inference

* topdown_multi_batch_inference

* topdown_multi_batch_inference

* some modifications

* topdown_mpii_dataset

* topdown_posetrack18_dataset

* assert bbox_id while bs > 1

Co-authored-by: sensetime <sensetime@sensetime.domain.sensetime.com>
  • Loading branch information
@wusize
wusize and sensetime authored Jan 4, 2021
1 parent 8d3bdfe commit 73130f3
Show file tree
Hide file tree
Showing 14 changed files with 251 additions and 150 deletions.
2 changes: 2 additions & 0 deletions configs/top_down/resnet/coco/res50_coco_256x192.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,8 @@
data = dict(
samples_per_gpu=64,
workers_per_gpu=2,
val_dataloader=dict(samples_per_gpu=32),
test_dataloader=dict(samples_per_gpu=32),
train=dict(
type='TopDownCocoDataset',
ann_file=f'{data_root}/annotations/person_keypoints_train2017.json',
Expand Down
1 change: 0 additions & 1 deletion mmpose/apis/train.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -112,7 +112,6 @@ def train_model(model,
eval_cfg = cfg.get('evaluation', {})
val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
dataloader_setting = dict(
# samples_per_gpu=cfg.data.get('samples_per_gpu', {}),
samples_per_gpu=1,
workers_per_gpu=cfg.data.get('workers_per_gpu', {}),
# cfg.gpus will be ignored if distributed
Expand Down
60 changes: 38 additions & 22 deletions mmpose/datasets/datasets/top_down/topdown_coco_dataset.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -188,6 +188,7 @@ def _load_coco_keypoint_annotation_kernel(self, img_id):
valid_objs.append(obj)
objs = valid_objs

bbox_id = 0
rec = []
for obj in objs:
if 'keypoints' not in obj:
Expand Down Expand Up @@ -215,8 +216,10 @@ def _load_coco_keypoint_annotation_kernel(self, img_id):
'joints_3d': joints_3d,
'joints_3d_visible': joints_3d_visible,
'dataset': self.dataset_name,
'bbox_score': 1
'bbox_score': 1,
'bbox_id': bbox_id
})
bbox_id = bbox_id + 1

return rec

Expand Down Expand Up @@ -264,7 +267,7 @@ def _load_coco_person_detection_results(self):
print(f'=> Total boxes: {len(all_boxes)}')

kpt_db = []
num_boxes = 0
bbox_id = 0
for det_res in all_boxes:
if det_res['category_id'] != 1:
continue
Expand All @@ -277,8 +280,6 @@ def _load_coco_person_detection_results(self):
if score < self.det_bbox_thr:
continue

num_boxes = num_boxes + 1

center, scale = self._xywh2cs(*box[:4])
joints_3d = np.zeros((num_joints, 3), dtype=np.float32)
joints_3d_visible = np.ones((num_joints, 3), dtype=np.float32)
Expand All @@ -291,10 +292,12 @@ def _load_coco_person_detection_results(self):
'bbox_score': score,
'dataset': self.dataset_name,
'joints_3d': joints_3d,
'joints_3d_visible': joints_3d_visible
'joints_3d_visible': joints_3d_visible,
'bbox_id': bbox_id
})
bbox_id = bbox_id + 1
print(f'=> Total boxes after filter '
f'low score@{self.det_bbox_thr}: {num_boxes}')
f'low score@{self.det_bbox_thr}: {bbox_id}')
return kpt_db

def evaluate(self, outputs, res_folder, metric='mAP', **kwargs):
Expand All @@ -313,10 +316,10 @@ def evaluate(self, outputs, res_folder, metric='mAP', **kwargs):
coordinates, score is the third dimension of the array.
:boxes (np.ndarray[1,6]): [center[0], center[1], scale[0]
, scale[1],area, score]
:image_path (list[str]): For example, [ '/', 'v','a', 'l',
'2', '0', '1', '7', '/', '0', '0', '0', '0', '0',
'0', '3', '9', '7', '1', '3', '3', '.', 'j', 'p', 'g']
:heatmap (np.ndarray[N, K, H, W]): model output heatmap.
:image_path (list[str]): For example, ['data/coco/val2017
/000000393226.jpg']
:heatmap (np.ndarray[N, K, H, W]): model output heatmap
:bbox_id (list(int)).
res_folder (str): Path of directory to save the results.
metric (str | list[str]): Metric to be performed. Defaults: 'mAP'.
Expand All @@ -332,18 +335,20 @@ def evaluate(self, outputs, res_folder, metric='mAP', **kwargs):
res_file = os.path.join(res_folder, 'result_keypoints.json')

kpts = defaultdict(list)
for preds, boxes, image_path, _ in outputs:
str_image_path = ''.join(image_path)
image_id = self.name2id[str_image_path[len(self.img_prefix):]]

kpts[image_id].append({
'keypoints': preds[0],
'center': boxes[0][0:2],
'scale': boxes[0][2:4],
'area': boxes[0][4],
'score': boxes[0][5],
'image_id': image_id,
})
for preds, boxes, image_paths, _, bbox_ids in outputs:
batch_size = len(image_paths)
for i in range(batch_size):
image_id = self.name2id[image_paths[i][len(self.img_prefix):]]
kpts[image_id].append({
'keypoints': preds[i],
'center': boxes[i][0:2],
'scale': boxes[i][2:4],
'area': boxes[i][4],
'score': boxes[i][5],
'image_id': image_id,
'bbox_id': bbox_ids[i]
})
kpts = self._sort_and_unique_bboxes(kpts)

# rescoring and oks nms
num_joints = self.ann_info['num_joints']
Expand Down Expand Up @@ -441,3 +446,14 @@ def _do_python_keypoint_eval(self, res_file):
info_str = list(zip(stats_names, coco_eval.stats))

return info_str

def _sort_and_unique_bboxes(self, kpts, key='bbox_id'):
"""sort kpts and remove the repeated ones."""
for img_id, persons in kpts.items():
num = len(persons)
kpts[img_id] = sorted(kpts[img_id], key=lambda x: x[key])
for i in range(num - 1, 0, -1):
if kpts[img_id][i][key] == kpts[img_id][i - 1][key]:
del kpts[img_id][i]

return kpts
5 changes: 4 additions & 1 deletion mmpose/datasets/datasets/top_down/topdown_coco_wholebody_dataset.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -182,6 +182,7 @@ def _load_coco_keypoint_annotation_kernel(self, img_id):
objs = valid_objs

rec = []
bbox_id = 0
for obj in objs:
if max(obj['keypoints']) == 0:
continue
Expand All @@ -205,8 +206,10 @@ def _load_coco_keypoint_annotation_kernel(self, img_id):
'joints_3d': joints_3d,
'joints_3d_visible': joints_3d_visible,
'dataset': self.dataset_name,
'bbox_score': 1
'bbox_score': 1,
'bbox_id': bbox_id
})
bbox_id = bbox_id + 1

return rec

Expand Down
42 changes: 28 additions & 14 deletions mmpose/datasets/datasets/top_down/topdown_jhmdb_dataset.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,7 @@ def _load_coco_keypoint_annotation_kernel(self, img_id):
objs = valid_objs

rec = []
bbox_id = 0
for obj in objs:
if 'keypoints' not in obj:
continue
Expand Down Expand Up @@ -185,8 +186,10 @@ def _load_coco_keypoint_annotation_kernel(self, img_id):
'joints_3d': joints_3d,
'joints_3d_visible': joints_3d_visible,
'dataset': self.dataset_name,
'bbox_score': 1
'bbox_score': 1,
'bbox_id': f'{img_id}_{bbox_id:03}'
})
bbox_id = bbox_id + 1

return rec

Expand Down Expand Up @@ -327,25 +330,36 @@ def evaluate(self, outputs, res_folder, metric='PCK', **kwargs):

kpts = []

for preds, boxes, image_path, _ in outputs:
for preds, boxes, image_paths, _, bbox_ids in outputs:
# convert 0-based index to 1-based index,
# and get the first two dimensions.
preds[..., :2] += 1.0

str_image_path = ''.join(image_path)
image_id = self.name2id[str_image_path[len(self.img_prefix):]]

kpts.append({
'keypoints': preds[0].tolist(),
'center': boxes[0][0:2].tolist(),
'scale': boxes[0][2:4].tolist(),
'area': float(boxes[0][4]),
'score': float(boxes[0][5]),
'image_id': image_id,
})
batch_size = len(image_paths)
for i in range(batch_size):
image_id = self.name2id[image_paths[i][len(self.img_prefix):]]
kpts.append({
'keypoints': preds[i],
'center': boxes[i][0:2],
'scale': boxes[i][2:4],
'area': boxes[i][4],
'score': boxes[i][5],
'image_id': image_id,
'bbox_id': bbox_ids[i]
})
kpts = self._sort_and_unique_bboxes(kpts)

self._write_keypoint_results(kpts, res_file)
info_str = self._report_metric(res_file, metrics)
name_value = OrderedDict(info_str)

return name_value

def _sort_and_unique_bboxes(self, kpts, key='bbox_id'):
"""sort kpts and remove the repeated ones."""
kpts = sorted(kpts, key=lambda x: x[key])
num = len(kpts)
for i in range(num - 1, 0, -1):
if kpts[i][key] == kpts[i - 1][key]:
del kpts[i]

return kpts
53 changes: 33 additions & 20 deletions mmpose/datasets/datasets/top_down/topdown_mpii_dataset.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@ def _get_db(self):
anno = json.load(anno_file)

gt_db = []
bbox_id = 0
for a in anno:
image_name = a['image']

Expand Down Expand Up @@ -111,25 +112,20 @@ def _get_db(self):

joints_3d[:, 0:2] = joints[:, 0:2] - 1
joints_3d_visible[:, :2] = joints_vis[:, None]

image_file = os.path.join(self.img_prefix, image_name)
gt_db.append({
'image_file':
os.path.join(self.img_prefix, image_name),
'center':
center,
'scale':
scale,
'rotation':
0,
'joints_3d':
joints_3d,
'joints_3d_visible':
joints_3d_visible,
'dataset':
'mpii',
'bbox_score':
1
'image_file': image_file,
'bbox_id': bbox_id,
'center': center,
'scale': scale,
'rotation': 0,
'joints_3d': joints_3d,
'joints_3d_visible': joints_3d_visible,
'dataset': 'mpii',
'bbox_score': 1
})
bbox_id = bbox_id + 1
gt_db = sorted(gt_db, key=lambda x: x['bbox_id'])

return gt_db

Expand All @@ -151,8 +147,8 @@ def evaluate(self, outputs, res_folder, metric='PCKh', **kwargs):
coordinates, score is the third dimension of the array.
* boxes(np.ndarray[1,6]): [center[0], center[1], scale[0]
, scale[1],area, score]
* image_path(list[str]): For example, ['0', '0',
'0', '0', '0', '1', '1', '6', '3', '.', 'j', 'p', 'g']
* image_path(list[str]): For example, ['/val2017/000000
397133.jpg']
* heatmap (np.ndarray[N, K, H, W]): model output heatmap.
res_folder(str): Path of directory to save the results.
Expand All @@ -169,7 +165,14 @@ def evaluate(self, outputs, res_folder, metric='PCKh', **kwargs):
if metric not in allowed_metrics:
raise KeyError(f'metric {metric} is not supported')

preds = np.stack([kpts[0] for kpts, _, _, _ in outputs])
kpts = []
for preds, _, _, _, bbox_ids in outputs:
batch_size = len(bbox_ids)
for i in range(batch_size):
kpts.append({'keypoints': preds[i], 'bbox_id': bbox_ids[i]})
kpts = self._sort_and_unique_bboxes(kpts)

preds = np.stack([kpt['keypoints'] for kpt in kpts])

# convert 0-based index to 1-based index,
# and get the first two dimensions.
Expand Down Expand Up @@ -248,3 +251,13 @@ def evaluate(self, outputs, res_folder, metric='PCKh', **kwargs):
name_value = OrderedDict(name_value)

return name_value

def _sort_and_unique_bboxes(self, kpts, key='bbox_id'):
"""sort kpts and remove the repeated ones."""
kpts = sorted(kpts, key=lambda x: x[key])
num = len(kpts)
for i in range(num - 1, 0, -1):
if kpts[i][key] == kpts[i - 1][key]:
del kpts[i]

return kpts
Loading

0 comments on commit 73130f3

Please sign in to comment.