feat(keypoints): Keypoints support 1.2 (#60)

Author: greatlog

.gitignore

@@ -31,3 +31,4 @@ wheels/
 
 # vscode editor settings
 .vscode
+

hubconf.py

@@ -47,7 +47,6 @@
 from official.vision.detection.tools.utils import DetEvaluator
 from official.vision.keypoints.inference import KeypointEvaluator
 from official.vision.keypoints.models import (
-    mspn_4stage,
     simplebaseline_res50,
     simplebaseline_res101,
     simplebaseline_res152,

official/vision/keypoints/README.md

@@ -1,14 +1,13 @@
 # Human Pose Esimation
 
-本目录包含了采用MegEngine实现的经典[SimpleBaseline](https://arxiv.org/pdf/1804.06208.pdf)和[MSPN](https://arxiv.org/pdf/1901.00148.pdf)网络结构，同时提供了在COCO数据集上的完整训练和测试代码。
+本目录包含了采用MegEngine实现的经典[SimpleBaseline](https://arxiv.org/pdf/1804.06208.pdf)的网络结构，同时提供了在COCO数据集上的完整训练和测试代码。
 
 本目录使用了在COCO val2017上的Human AP为56.4的人体检测结果，最后在COCO val2017上人体关节点估计结果为
 |Methods|Backbone|Input Size| AP | Ap .5 | AP .75 | AP (M) | AP (L) | AR | AR .5 | AR .75 | AR (M) | AR (L) |
 |---|:---:|---|---|---|---|---|---|---|---|---|---|---|
-| SimpleBaseline |Res50 |256x192| 0.712 | 0.887 | 0.779 | 0.673 | 0.785 | 0.782 | 0.932 | 0.839 | 0.730 | 0.854 |
-| SimpleBaseline |Res101|256x192| 0.722 | 0.891 | 0.795 | 0.687 | 0.795 | 0.794 | 0.936 | 0.855 | 0.745 | 0.863 |
-| SimpleBaseline |Res152|256x192| 0.724 | 0.888 | 0.794 | 0.688 | 0.795 | 0.795 | 0.934 | 0.856 | 0.746 | 0.863 |
-| MSPN_4stage |MSPN|256x192| 0.752 | 0.900 | 0.819 | 0.716 | 0.825 | 0.819 | 0.943 | 0.875 | 0.770 | 0.887 |
+| SimpleBaseline |Res50 |256x192| 0.711 | 0.885 | 0.779 | 0.674 | 0.783 | 0.782 | 0.930 | 0.839 | 0.731 | 0.852 |
+| SimpleBaseline |Res101|256x192| 0.718 | 0.892 | 0.788 | 0.681 | 0.793 | 0.790 | 0.937 | 0.848 | 0.739 | 0.861 |
+| SimpleBaseline |Res152|256x192| 0.723 | 0.888 | 0.794 | 0.688 | 0.795 | 0.795 | 0.934 | 0.856 | 0.746 | 0.863 |
 
 ## 安装和环境配置
 
@@ -65,16 +64,7 @@ ${COCO_DATA_ROOT}
 python3 train.py --arch simplebaseline_res50 \
                  --resume /path/to/model \
                  --ngpus 8 \
-                 --multi_scale_supervision False
 
-```
-训练MSPN:
-```bash
-python3 train.py --arch mspn_4stage \
-                 --resume /path/to/model \
-                 --ngpus 8 \
-                 --multi_scale_supervision True
-
 ```
 
 ## 如何测试
@@ -83,12 +73,10 @@ python3 train.py --arch mspn_4stage \
 ```bash
 python3 test.py --arch name/of/network \
                 --model /path/to/model.pkl \
-                --dt_file /name/human/detection/results
 ```
 `test.py`的命令行参数如下：
 - `--arch`, 网络的名字;
-- `--model`, 待检测的模;
-- `--dt_path`，人体检测结果.
+- `--model`, 待检测的模型;
 
 也可以连续验证多个模型的性能:
 
@@ -119,5 +107,4 @@ python3 inference.py --arch /name/of/tested/network \
 
 ## 参考文献
 
-- [Simple Baselines for Human Pose Estimation and Tracking](https://arxiv.org/pdf/1804.06208.pdf) Bin Xiao, Haiping Wu, and Yichen Wei
-- [Rethinking on Multi-Stage Networks for Human Pose Estimation](https://arxiv.org/pdf/1901.00148.pdf) Wenbo Li1, Zhicheng Wang, Binyi Yin, Qixiang Peng, Yuming Du, Tianzi Xiao, Gang Yu, Hongtao Lu, Yichen Wei and Jian Sun
+- [Simple Baselines for Human Pose Estimation and Tracking](https://arxiv.org/abs/1804.06208) Bin Xiao, Haiping Wu, and Yichen Wei. European Conference on Computer Vision (ECCV), 2018.

official/vision/keypoints/config.py

@@ -6,36 +6,47 @@
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
 # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+
+
 class Config:
-    ##############3## train ##############################################
+    # model
+    model_choices = [
+        "simplebaseline_res50",
+        "simplebaseline_res101",
+        "simplebaseline_res152",
+    ]
+
+    # train
     initial_lr = 3e-4
     lr_ratio = 0.1
 
     batch_size = 32
     epochs = 200
-    warm_epochs = 1
-    weight_decay = 1e-5
+    warm_epochs = 0
+    weight_decay = 0
+
+    report_freq = 10
+    save_freq = 1
 
-    ################## data ###############################################
+    # data
     # path
     data_root = "/data/coco_data/"
 
     # normalize
-    img_mean = [0.485 * 255, 0.456 * 255, 0.406 * 255]
-    img_std = [0.229 * 255, 0.224 * 255, 0.225 * 255]
+    img_mean = [103.530, 116.280, 123.675]
+    img_std = [57.375, 57.120, 58.395]
 
     # shape
     input_shape = (256, 192)
     output_shape = (64, 48)
 
     # heat maps
     keypoint_num = 17
-    heat_kernel = [2.6, 2.0, 1.7, 1.4]
+    heat_kernels = [k * 4 for k in [2.6, 2.0, 1.7, 1.4]]
     heat_thr = 1e-2
     heat_range = 255
 
-    ##################### augumentation #####################################
-
+    # augmentation
     half_body_transform = True
     extend_boxes = True
 
@@ -53,19 +64,21 @@ class Config:
 
     # scale
     scale_prob = 1
-    scale_range = [0.7, 1.3]
+    scale_range = 0.3
 
     # rorate
     rotation_prob = 0.6
-    rotate_range = [-45, 45]
+    rotate_range = 40
 
-    ############## testing settings ##########################################
+    # test settings
     test_aug_border = 10
     test_x_ext = 0.10
     test_y_ext = 0.10
     test_gaussian_kernel = 17
     second_value_aug = True
 
+    # inference settings
+    nms_thr = 0.7
     vis_colors = [
         [255, 0, 0],
         [255, 85, 0],
@@ -95,6 +108,8 @@ class Config:
         [0, 2],
         [1, 3],
         [2, 4],
+        [3, 5],
+        [4, 6],
         [5, 6],
         [5, 7],
         [7, 9],

official/vision/keypoints/dataset.py

@@ -6,21 +6,21 @@
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
 # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-import megengine as mge
-from megengine.data.dataset.vision.meta_vision import VisionDataset
-from megengine.data import Collator
+import json
+import os.path as osp
+from collections import OrderedDict, defaultdict
 
-import numpy as np
 import cv2
-import os.path as osp
-import json
-from collections import defaultdict, OrderedDict
+import numpy as np
+
+from megengine.data import Collator
+from megengine.data.dataset.vision.meta_vision import VisionDataset
 
 
 class COCOJoints(VisionDataset):
     """
     we cannot use the official implementation of COCO dataset here.
-    The output of __getitem__ function here should be a single person instead of a single image. 
+    The output of __getitem__ function here should be a single person instead of a single image.
     """
 
     supported_order = ("image", "keypoints", "boxes", "info")
@@ -47,7 +47,7 @@ class COCOJoints(VisionDataset):
 
     min_bbox_h = 0
     min_bbox_w = 0
-    min_box_area = 1500
+    min_bbox_area = 1500
     min_bbox_score = 1e-10
 
     def __init__(
@@ -87,8 +87,6 @@ def __init__(
 
         selected_anns = []
         for ann in dataset["annotations"]:
-            if "image_id" in ann.keys() and ann["image_id"] not in self.ids:
-                continue
 
             if "iscrowd" in ann.keys() and ann["iscrowd"]:
                 continue
@@ -129,8 +127,8 @@ def __getitem__(self, index):
         img_id = ann["image_id"]
         target = []
         for k in self.order:
-            if k == "image":
 
+            if k == "image":
                 file_name = self.imgs[img_id]["file_name"]
                 img_path = osp.join(self.root, self.image_set, file_name)
                 image = cv2.imread(img_path, cv2.IMREAD_COLOR)
@@ -186,13 +184,9 @@ def __init__(
 
         self.stride = image_shape[1] // heatmap_shape[1]
 
-        x = np.arange(0, heatmap_shape[1], 1)
-        y = np.arange(0, heatmap_shape[0], 1)
-
-        grid_x, grid_y = np.meshgrid(x, y)
-
-        self.grid_x = grid_x[None].repeat(keypoint_num, 0)
-        self.grid_y = grid_y[None].repeat(keypoint_num, 0)
+        ax = (np.arange(0, heatmap_shape[1]) + 0.5) * self.stride - 0.5
+        ay = (np.arange(0, heatmap_shape[0]) + 0.5) * self.stride - 0.5
+        self.grid_x, self.grid_y = np.meshgrid(ax, ay)
 
     def apply(self, inputs):
         """
@@ -204,27 +198,21 @@ def apply(self, inputs):
 
             batch_data["data"].append(image)
 
-            joints = (keypoints[0, :, :2] + 0.5) / self.stride - 0.5
-            heat_valid = np.array(keypoints[0, :, -1]).astype(np.float32)
-            dis = (self.grid_x - joints[:, 0, np.newaxis, np.newaxis]) ** 2 + (
-                self.grid_y - joints[:, 1, np.newaxis, np.newaxis]
+            joint = keypoints[0, :, :2]
+            dis = (self.grid_x[None] - joint[:, 0, None, None]) ** 2 + (
+                self.grid_y[None] - joint[:, 1, None, None]
             ) ** 2
+            heat_valid = np.array(keypoints[0, :, -1]).astype(np.float32)
+
             heatmaps = []
             for k in self.heat_kernel:
+
                 heatmap = np.exp(-dis / 2 / k ** 2)
+                heatmap[heat_valid < 0.1] = 0
                 heatmap[heatmap < self.heat_thr] = 0
-                heatmap[heat_valid == 0] = 0
-                sum_for_norm = heatmap.sum((1, 2))
-                heatmap[sum_for_norm > 0] = (
-                    heatmap[sum_for_norm > 0]
-                    / sum_for_norm[sum_for_norm > 0][:, None, None]
-                )
-                maxi = np.max(heatmap, (1, 2))
-                heatmap[maxi > 1e-5] = (
-                    heatmap[maxi > 1e-5]
-                    / maxi[:, None, None][maxi > 1e-5]
-                    * self.heat_range
-                )
+
+                heatmap *= self.heat_range
+
                 heatmaps.append(heatmap)
 
             batch_data["heatmap"].append(np.array(heatmaps))