add use_nms option (#318)

* add use_nms option

* use default value

* rm unnecessary lines

* rm unnecessary lines

Co-authored-by: jinsheng <jinsheng@sensetime.com>

mmpose/datasets/datasets/top_down/topdown_aic_dataset.py

@@ -54,7 +54,7 @@ def __init__(self,
         self.use_gt_bbox = data_cfg['use_gt_bbox']
         self.bbox_file = data_cfg['bbox_file']
         self.image_thr = data_cfg['image_thr']
-
+        self.use_nms = data_cfg.get('use_nms', True)
         self.soft_nms = data_cfg['soft_nms']
         self.nms_thr = data_cfg['nms_thr']
         self.oks_thr = data_cfg['oks_thr']

mmpose/datasets/datasets/top_down/topdown_coco_dataset.py

@@ -64,7 +64,7 @@ def __init__(self,
         self.use_gt_bbox = data_cfg['use_gt_bbox']
         self.bbox_file = data_cfg['bbox_file']
         self.image_thr = data_cfg['image_thr']
-
+        self.use_nms = data_cfg.get('use_nms', True)
         self.soft_nms = data_cfg['soft_nms']
         self.nms_thr = data_cfg['nms_thr']
         self.oks_thr = data_cfg['oks_thr']
@@ -344,7 +344,7 @@ def evaluate(self, outputs, res_folder, metric='mAP', **kwargs):
         num_joints = self.ann_info['num_joints']
         vis_thr = self.vis_thr
         oks_thr = self.oks_thr
-        oks_nmsed_kpts = []
+        valid_kpts = []
         for image_id in kpts.keys():
             img_kpts = kpts[image_id]
             for n_p in img_kpts:
@@ -361,15 +361,14 @@ def evaluate(self, outputs, res_folder, metric='mAP', **kwargs):
                 # rescoring
                 n_p['score'] = kpt_score * box_score
 
-            nms = soft_oks_nms if self.soft_nms else oks_nms
-            keep = nms(list(img_kpts), oks_thr, sigmas=self.sigmas)
-
-            if len(keep) == 0:
-                oks_nmsed_kpts.append(img_kpts)
+            if self.use_nms:
+                nms = soft_oks_nms if self.soft_nms else oks_nms
+                keep = nms(list(img_kpts), oks_thr, sigmas=self.sigmas)
+                valid_kpts.append([img_kpts[_keep] for _keep in keep])
             else:
-                oks_nmsed_kpts.append([img_kpts[_keep] for _keep in keep])
+                valid_kpts.append(img_kpts)
 
-        self._write_coco_keypoint_results(oks_nmsed_kpts, res_file)
+        self._write_coco_keypoint_results(valid_kpts, res_file)
 
         info_str = self._do_python_keypoint_eval(res_file)
         name_value = OrderedDict(info_str)

mmpose/datasets/datasets/top_down/topdown_coco_wholebody_dataset.py

@@ -49,7 +49,7 @@ def __init__(self,
         self.use_gt_bbox = data_cfg['use_gt_bbox']
         self.bbox_file = data_cfg['bbox_file']
         self.image_thr = data_cfg['image_thr']
-
+        self.use_nms = data_cfg.get('use_nms', True)
         self.soft_nms = data_cfg['soft_nms']
         self.nms_thr = data_cfg['nms_thr']
         self.oks_thr = data_cfg['oks_thr']

mmpose/datasets/datasets/top_down/topdown_crowdpose_dataset.py

@@ -52,7 +52,7 @@ def __init__(self,
         self.use_gt_bbox = data_cfg['use_gt_bbox']
         self.bbox_file = data_cfg['bbox_file']
         self.image_thr = data_cfg['image_thr']
-
+        self.use_nms = data_cfg.get('use_nms', True)
         self.soft_nms = data_cfg['soft_nms']
         self.nms_thr = data_cfg['nms_thr']
         self.oks_thr = data_cfg['oks_thr']

mmpose/datasets/datasets/top_down/topdown_ochuman_dataset.py

@@ -61,7 +61,7 @@ def __init__(self,
         self.use_gt_bbox = data_cfg['use_gt_bbox']
         self.bbox_file = data_cfg['bbox_file']
         self.image_thr = data_cfg['image_thr']
-
+        self.use_nms = data_cfg.get('use_nms', True)
         self.soft_nms = data_cfg['soft_nms']
         self.nms_thr = data_cfg['nms_thr']
         self.oks_thr = data_cfg['oks_thr']

mmpose/datasets/datasets/top_down/topdown_posetrack18_dataset.py

@@ -65,7 +65,7 @@ def __init__(self,
         self.use_gt_bbox = data_cfg['use_gt_bbox']
         self.bbox_file = data_cfg['bbox_file']
         self.image_thr = data_cfg['image_thr']
-
+        self.use_nms = data_cfg.get('use_nms', True)
         self.soft_nms = data_cfg['soft_nms']
         self.nms_thr = data_cfg['nms_thr']
         self.oks_thr = data_cfg['oks_thr']
@@ -165,7 +165,7 @@ def evaluate(self, outputs, res_folder, metric='mAP', **kwargs):
         num_joints = self.ann_info['num_joints']
         vis_thr = self.vis_thr
         oks_thr = self.oks_thr
-        oks_nmsed_kpts = defaultdict(list)
+        valid_kpts = defaultdict(list)
         for image_id in kpts.keys():
             img_kpts = kpts[image_id]
             for n_p in img_kpts:
@@ -181,17 +181,15 @@ def evaluate(self, outputs, res_folder, metric='mAP', **kwargs):
                     kpt_score = kpt_score / valid_num
                 # rescoring
                 n_p['score'] = kpt_score * box_score
-
-            nms = soft_oks_nms if self.soft_nms else oks_nms
-            keep = nms(list(img_kpts), oks_thr, sigmas=self.sigmas)
-
-            if len(keep) == 0:
-                oks_nmsed_kpts[image_id].append(img_kpts)
-            else:
-                oks_nmsed_kpts[image_id].append(
+            if self.use_nms:
+                nms = soft_oks_nms if self.soft_nms else oks_nms
+                keep = nms(list(img_kpts), oks_thr, sigmas=self.sigmas)
+                valid_kpts[image_id].append(
                     [img_kpts[_keep] for _keep in keep])
+            else:
+                valid_kpts[image_id].append(img_kpts)
 
-        self._write_posetrack18_keypoint_results(oks_nmsed_kpts, gt_folder,
+        self._write_posetrack18_keypoint_results(valid_kpts, gt_folder,
                                                  pred_folder)
 
         info_str = self._do_python_keypoint_eval(gt_folder, pred_folder)