From 9ae434f8a79310c56c652e1c4e6cdbf7ff3b020f Mon Sep 17 00:00:00 2001 From: zhiboniu <31800336+zhiboniu@users.noreply.github.com> Date: Tue, 27 Jul 2021 17:50:41 +0800 Subject: [PATCH] add mpii infer support (#3782) --- deploy/python/visualize.py | 4 +- ppdet/data/source/category.py | 3 +- ppdet/engine/trainer.py | 15 ++++--- ppdet/metrics/keypoint_metrics.py | 2 +- ppdet/utils/visualizer.py | 69 ++++++++++++++++++++++--------- 5 files changed, 66 insertions(+), 27 deletions(-) diff --git a/deploy/python/visualize.py b/deploy/python/visualize.py index f15d1b7a6..402857509 100644 --- a/deploy/python/visualize.py +++ b/deploy/python/visualize.py @@ -240,7 +240,9 @@ def draw_pose(imgfile, raise e skeletons, scores = results['keypoint'] - kpt_nums = skeletons.shape[1] + kpt_nums = 17 + if len(skeletons) > 0: + kpt_nums = skeletons.shape[1] if kpt_nums == 17: #plot coco keypoint EDGES = [(0, 1), (0, 2), (1, 3), (2, 4), (3, 5), (4, 6), (5, 7), (6, 8), (7, 9), (8, 10), (5, 11), (6, 12), (11, 13), (12, 14), diff --git a/ppdet/data/source/category.py b/ppdet/data/source/category.py index 9e612c21b..56ca0f64d 100644 --- a/ppdet/data/source/category.py +++ b/ppdet/data/source/category.py @@ -83,7 +83,8 @@ def get_categories(metric_type, anno_file=None, arch=None): elif metric_type.lower() == 'widerface': return _widerface_category() - elif metric_type.lower() == 'keypointtopdowncocoeval': + elif metric_type.lower() == 'keypointtopdowncocoeval' or metric_type.lower( + ) == 'keypointtopdownmpiieval': return (None, {'id': 'keypoint'}) elif metric_type.lower() in ['mot', 'motdet', 'reid']: diff --git a/ppdet/engine/trainer.py b/ppdet/engine/trainer.py index d8a2803d9..e08d8de07 100644 --- a/ppdet/engine/trainer.py +++ b/ppdet/engine/trainer.py @@ -292,6 +292,7 @@ class Trainer(object): def train(self, validate=False): assert self.mode == 'train', "Model not in 'train' mode" + Init_mark = False # if validation in training is enabled, metrics should be re-init if validate: @@ -394,6 +395,12 @@ class Trainer(object): self._eval_dataset, self.cfg.worker_num, batch_sampler=self._eval_batch_sampler) + # if validation in training is enabled, metrics should be re-init + # Init_mark makes sure this code will only execute once + if validate and Init_mark == False: + Init_mark = True + self._init_metrics(validate=validate) + self._reset_metrics() with paddle.no_grad(): self.status['save_best_model'] = True self._eval_with_loader(self._eval_loader) @@ -558,14 +565,12 @@ class Trainer(object): shape=[None, 3, 192, 64], name='crops') }) - - static_model = paddle.jit.to_static( - self.model, input_spec=input_spec) + static_model = paddle.jit.to_static(self.model, input_spec=input_spec) # NOTE: dy2st do not pruned program, but jit.save will prune program # input spec, prune input spec here and save with pruned input spec pruned_input_spec = self._prune_input_spec( - input_spec, static_model.forward.main_program, - static_model.forward.outputs) + input_spec, static_model.forward.main_program, + static_model.forward.outputs) # dy2st and save model if 'slim' not in self.cfg or self.cfg['slim_type'] != 'QAT': diff --git a/ppdet/metrics/keypoint_metrics.py b/ppdet/metrics/keypoint_metrics.py index 49f0c7a58..150018c61 100644 --- a/ppdet/metrics/keypoint_metrics.py +++ b/ppdet/metrics/keypoint_metrics.py @@ -14,7 +14,7 @@ import os import json -from collections import defaultdict +from collections import defaultdict, OrderedDict import numpy as np from pycocotools.coco import COCO from pycocotools.cocoeval import COCOeval diff --git a/ppdet/utils/visualizer.py b/ppdet/utils/visualizer.py index 505c416b0..fdfd966e2 100644 --- a/ppdet/utils/visualizer.py +++ b/ppdet/utils/visualizer.py @@ -218,23 +218,35 @@ def draw_segm(image, return Image.fromarray(img_array.astype('uint8')) -def map_coco_to_personlab(keypoints): - permute = [0, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3] - return keypoints[:, permute, :] - - -def draw_pose(image, results, visual_thread=0.6, save_name='pose.jpg'): +def draw_pose(image, + results, + visual_thread=0.6, + save_name='pose.jpg', + save_dir='output', + returnimg=False, + ids=None): try: import matplotlib.pyplot as plt import matplotlib plt.switch_backend('agg') except Exception as e: - logger.error('Matplotlib not found, plaese install matplotlib.' + logger.error('Matplotlib not found, please install matplotlib.' 'for example: `pip install matplotlib`.') raise e - EDGES = [(0, 14), (0, 13), (0, 4), (0, 1), (14, 16), (13, 15), (4, 10), - (1, 7), (10, 11), (7, 8), (11, 12), (8, 9), (4, 5), (1, 2), (5, 6), - (2, 3)] + + skeletons = np.array([item['keypoints'] for item in results]) + kpt_nums = 17 + if len(skeletons) > 0: + kpt_nums = int(skeletons.shape[1] / 3) + skeletons = skeletons.reshape(-1, kpt_nums, 3) + if kpt_nums == 17: #plot coco keypoint + EDGES = [(0, 1), (0, 2), (1, 3), (2, 4), (3, 5), (4, 6), (5, 7), (6, 8), + (7, 9), (8, 10), (5, 11), (6, 12), (11, 13), (12, 14), + (13, 15), (14, 16), (11, 12)] + else: #plot mpii keypoint + EDGES = [(0, 1), (1, 2), (3, 4), (4, 5), (2, 6), (3, 6), (6, 7), (7, 8), + (8, 9), (10, 11), (11, 12), (13, 14), (14, 15), (8, 12), + (8, 13)] NUM_EDGES = len(EDGES) colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \ @@ -242,22 +254,36 @@ def draw_pose(image, results, visual_thread=0.6, save_name='pose.jpg'): [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]] cmap = matplotlib.cm.get_cmap('hsv') plt.figure() - skeletons = np.array([item['keypoints'] for item in results]).reshape(-1, - 17, 3) + img = np.array(image).astype('float32') - canvas = img.copy() - for i in range(17): - rgba = np.array(cmap(1 - i / 17. - 1. / 34)) - rgba[0:3] *= 255 + color_set = results['colors'] if 'colors' in results else None + + if 'bbox' in results and ids is None: + bboxs = results['bbox'] + for j, rect in enumerate(bboxs): + xmin, ymin, xmax, ymax = rect + color = colors[0] if color_set is None else colors[color_set[j] % + len(colors)] + cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color, 1) + + canvas = img.copy() + for i in range(kpt_nums): for j in range(len(skeletons)): if skeletons[j][i, 2] < visual_thread: continue + if ids is None: + color = colors[i] if color_set is None else colors[color_set[j] + % + len(colors)] + else: + color = get_color(ids[j]) + cv2.circle( canvas, tuple(skeletons[j][i, 0:2].astype('int32')), 2, - colors[i], + color, thickness=-1) to_plot = cv2.addWeighted(img, 0.3, canvas, 0.7, 0) @@ -265,7 +291,6 @@ def draw_pose(image, results, visual_thread=0.6, save_name='pose.jpg'): stickwidth = 2 - skeletons = map_coco_to_personlab(skeletons) for i in range(NUM_EDGES): for j in range(len(skeletons)): edge = EDGES[i] @@ -283,7 +308,13 @@ def draw_pose(image, results, visual_thread=0.6, save_name='pose.jpg'): polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1) - cv2.fillConvexPoly(cur_canvas, polygon, colors[i]) + if ids is None: + color = colors[i] if color_set is None else colors[color_set[j] + % + len(colors)] + else: + color = get_color(ids[j]) + cv2.fillConvexPoly(cur_canvas, polygon, color) canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0) image = Image.fromarray(canvas.astype('uint8')) plt.close() -- GitLab