[Feature] Support SKPS model

README.md

@@ -342,21 +342,20 @@ This project is released under the [Apache 2.0 license](LICENSE).
 
 - [MMEngine](https://github.com/open-mmlab/mmengine): OpenMMLab foundational library for training deep learning models.
 - [MMCV](https://github.com/open-mmlab/mmcv): OpenMMLab foundational library for computer vision.
-- [MIM](https://github.com/open-mmlab/mim): MIM installs OpenMMLab packages.
-- [MMClassification](https://github.com/open-mmlab/mmclassification): OpenMMLab image classification toolbox and benchmark.
+- [MMPreTrain](https://github.com/open-mmlab/mmpretrain): OpenMMLab pre-training toolbox and benchmark.
+- [MMagic](https://github.com/open-mmlab/mmagic): Open**MM**Lab **A**dvanced, **G**enerative and **I**ntelligent **C**reation toolbox.
 - [MMDetection](https://github.com/open-mmlab/mmdetection): OpenMMLab detection toolbox and benchmark.
 - [MMDetection3D](https://github.com/open-mmlab/mmdetection3d): OpenMMLab's next-generation platform for general 3D object detection.
 - [MMRotate](https://github.com/open-mmlab/mmrotate): OpenMMLab rotated object detection toolbox and benchmark.
+- [MMTracking](https://github.com/open-mmlab/mmtracking): OpenMMLab video perception toolbox and benchmark.
 - [MMSegmentation](https://github.com/open-mmlab/mmsegmentation): OpenMMLab semantic segmentation toolbox and benchmark.
 - [MMOCR](https://github.com/open-mmlab/mmocr): OpenMMLab text detection, recognition, and understanding toolbox.
 - [MMPose](https://github.com/open-mmlab/mmpose): OpenMMLab pose estimation toolbox and benchmark.
 - [MMHuman3D](https://github.com/open-mmlab/mmhuman3d): OpenMMLab 3D human parametric model toolbox and benchmark.
-- [MMSelfSup](https://github.com/open-mmlab/mmselfsup): OpenMMLab self-supervised learning toolbox and benchmark.
-- [MMRazor](https://github.com/open-mmlab/mmrazor): OpenMMLab model compression toolbox and benchmark.
 - [MMFewShot](https://github.com/open-mmlab/mmfewshot): OpenMMLab fewshot learning toolbox and benchmark.
 - [MMAction2](https://github.com/open-mmlab/mmaction2): OpenMMLab's next-generation action understanding toolbox and benchmark.
-- [MMTracking](https://github.com/open-mmlab/mmtracking): OpenMMLab video perception toolbox and benchmark.
 - [MMFlow](https://github.com/open-mmlab/mmflow): OpenMMLab optical flow toolbox and benchmark.
-- [MMEditing](https://github.com/open-mmlab/mmediting): OpenMMLab image and video editing toolbox.
-- [MMGeneration](https://github.com/open-mmlab/mmgeneration): OpenMMLab image and video generative models toolbox.
 - [MMDeploy](https://github.com/open-mmlab/mmdeploy): OpenMMLab Model Deployment Framework.
+- [MMRazor](https://github.com/open-mmlab/mmrazor): OpenMMLab model compression toolbox and benchmark.
+- [MIM](https://github.com/open-mmlab/mim): MIM installs OpenMMLab packages.
+- [Playground](https://github.com/open-mmlab/playground): A central hub for gathering and showcasing amazing projects built upon OpenMMLab.

README_CN.md

@@ -339,24 +339,23 @@ MMPose 是一款由不同学校和公司共同贡献的开源项目。我们感
 
 - [MMEngine](https://github.com/open-mmlab/mmengine): OpenMMLab 深度学习模型训练基础库
 - [MMCV](https://github.com/open-mmlab/mmcv): OpenMMLab 计算机视觉基础库
-- [MIM](https://github.com/open-mmlab/mim): OpenMMlab 项目、算法、模型的统一入口
-- [MMClassification](https://github.com/open-mmlab/mmclassification): OpenMMLab 图像分类工具箱
+- [MMPreTrain](https://github.com/open-mmlab/mmpretrain): OpenMMLab 深度学习预训练工具箱
+- [MMagic](https://github.com/open-mmlab/mmagic): OpenMMLab 新一代人工智能内容生成（AIGC）工具箱
 - [MMDetection](https://github.com/open-mmlab/mmdetection): OpenMMLab 目标检测工具箱
 - [MMDetection3D](https://github.com/open-mmlab/mmdetection3d): OpenMMLab 新一代通用 3D 目标检测平台
 - [MMRotate](https://github.com/open-mmlab/mmrotate): OpenMMLab 旋转框检测工具箱与测试基准
+- [MMTracking](https://github.com/open-mmlab/mmtracking): OpenMMLab 一体化视频目标感知平台
 - [MMSegmentation](https://github.com/open-mmlab/mmsegmentation): OpenMMLab 语义分割工具箱
 - [MMOCR](https://github.com/open-mmlab/mmocr): OpenMMLab 全流程文字检测识别理解工具包
 - [MMPose](https://github.com/open-mmlab/mmpose): OpenMMLab 姿态估计工具箱
 - [MMHuman3D](https://github.com/open-mmlab/mmhuman3d): OpenMMLab 人体参数化模型工具箱与测试基准
-- [MMSelfSup](https://github.com/open-mmlab/mmselfsup): OpenMMLab 自监督学习工具箱与测试基准
-- [MMRazor](https://github.com/open-mmlab/mmrazor): OpenMMLab 模型压缩工具箱与测试基准
 - [MMFewShot](https://github.com/open-mmlab/mmfewshot): OpenMMLab 少样本学习工具箱与测试基准
 - [MMAction2](https://github.com/open-mmlab/mmaction2): OpenMMLab 新一代视频理解工具箱
-- [MMTracking](https://github.com/open-mmlab/mmtracking): OpenMMLab 一体化视频目标感知平台
 - [MMFlow](https://github.com/open-mmlab/mmflow): OpenMMLab 光流估计工具箱与测试基准
-- [MMEditing](https://github.com/open-mmlab/mmediting): OpenMMLab 图像视频编辑工具箱
-- [MMGeneration](https://github.com/open-mmlab/mmgeneration): OpenMMLab 图片视频生成模型工具箱
 - [MMDeploy](https://github.com/open-mmlab/mmdeploy): OpenMMLab 模型部署框架
+- [MMRazor](https://github.com/open-mmlab/mmrazor): OpenMMLab 模型压缩工具箱与测试基准
+- [MIM](https://github.com/open-mmlab/mim): OpenMMlab 项目、算法、模型的统一入口
+- [Playground](https://github.com/open-mmlab/playground): 收集和展示 OpenMMLab 相关的前沿、有趣的社区项目
 
 ## 欢迎加入 OpenMMLab 社区
 

projects/skps/README.md

@@ -0,0 +1,83 @@
+# Simple Keypoints
+
+## Description
+
+Author： @2120140200@mail.nankai.edu.cn
+
+It is a simple keypoints detector model. The model predict a score heatmap and an encoded location map.
+The result in wflw achieves 3.94 NME.
+
+## Usage
+
+### Prerequisites
+
+- Python 3.7
+- PyTorch 1.6 or higher
+- [MIM](https://github.com/open-mmlab/mim) v0.33 or higher
+- [MMPose](https://github.com/open-mmlab/mmpose) v1.0.0rc0 or higher
+
+All the commands below rely on the correct configuration of `PYTHONPATH`, which should point to the project's directory so that Python can locate the module files. In `example_project/` root directory, run the following line to add the current directory to `PYTHONPATH`:
+
+```shell
+export PYTHONPATH=`pwd`:$PYTHONPATH
+```
+
+### Data Preparation
+
+Prepare the COCO dataset according to the [instruction](https://mmpose.readthedocs.io/en/dev-1.x/dataset_zoo/2d_body_keypoint.html#coco).
+
+### Training commands
+
+**To train with single GPU:**
+
+```shell
+mim train mmpose configs/td-hm_hrnetv2-w18_skps-1xb64-80e_wflw-256x256.py
+```
+
+**To train with multiple GPUs:**
+
+```shell
+mim train mmpose configs/td-hm_hrnetv2-w18_skps-1xb64-80e_wflw-256x256.py --launcher pytorch --gpus 8
+```
+
+**To train with multiple GPUs by slurm:**
+
+```shell
+mim train mmpose configs/td-hm_hrnetv2-w18_skps-1xb64-80e_wflw-256x256.py --launcher slurm \
+    --gpus 16 --gpus-per-node 8 --partition $PARTITION
+```
+
+### Testing commands
+
+**To test with single GPU:**
+
+```shell
+mim test mmpose configs/td-hm_hrnetv2-w18_skps-1xb64-80e_wflw-256x256.py -C $CHECKPOINT
+```
+
+**To test with multiple GPUs:**
+
+```shell
+mim test mmpose configs/td-hm_hrnetv2-w18_skps-1xb64-80e_wflw-256x256.py -C $CHECKPOINT --launcher pytorch --gpus 8
+```
+
+**To test with multiple GPUs by slurm:**
+
+```shell
+mim test mmpose configs/td-hm_hrnetv2-w18_skps-1xb64-80e_wflw-256x256.py -C $CHECKPOINT --launcher slurm \
+    --gpus 16 --gpus-per-node 8 --partition $PARTITION
+```
+
+## Results
+
+WFLW
+
+| Arch       | Input Size | NME<sub>*test*</sub> | NME<sub>*pose*</sub> | NME<sub>*illumination*</sub> | NME<sub>*occlusion*</sub> | NME<sub>*blur*</sub> | NME<sub>*makeup*</sub> | NME<sub>*expression*</sub> |    ckpt    |    log    |
+| :--------- | :--------: | :------------------: | :------------------: | :--------------------------: | :-----------------------: | :------------------: | :--------------------: | :------------------------: | :--------: | :-------: |
+| [skps](/configs/td-hm_hrnetv2-w18_skps-1xb64-80e_wflw-256x256.py) |  256x256   |         3.88         |         6.60         |             3.81             |           4.57            |         4.44         |          3.75          |            4.13            | [ckpt](https://drive.google.com/file/d/10U7f_pp0BPMhm575cO7Vg13nZVCy-PfN/view?usp=sharing) | [log](https://drive.google.com/file/d/1fBNcAyMdRr9nTN8wHqvYPnd9AFBinOk4/view?usp=sharing) |
+
+COFW
+
+| Arch                                                           | Input Size | NME  |                              ckpt                              |                              log                               |
+| :------------------------------------------------------------- | :--------: | :--: | :------------------------------------------------------------: | :------------------------------------------------------------: |
+| [skps](/configs/td-hm_hrnetv2-w18_skps-1xb16-160e_cofw-256x256.py) |  256x256   | 3.20 | [ckpt](https://drive.google.com/file/d/1fdYQ0ajg11KAdkeLVCi_jPdSTKf9i3hP/view?usp=sharing) | [log](https://drive.google.com/file/d/1IUjIY_sLbO6YO59G7XZGUHDVC-o_k09y/view?usp=sharing) |

projects/skps/codecs2/__init__.py

@@ -0,0 +1,3 @@
+from .skps_heatmap import SKPSHeatmap
+
+__all__ = ['SKPSHeatmap']

projects/skps/codecs2/skps_heatmap.py

@@ -0,0 +1,166 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from typing import Optional, Tuple
+
+import numpy as np
+
+from mmpose.codecs.base import BaseKeypointCodec
+from mmpose.codecs.utils.gaussian_heatmap import \
+    generate_unbiased_gaussian_heatmaps
+from mmpose.codecs.utils.post_processing import get_heatmap_maximum
+from mmpose.registry import KEYPOINT_CODECS
+
+
+@KEYPOINT_CODECS.register_module()
+class SKPSHeatmap(BaseKeypointCodec):
+    """Generate heatmap the same with MSRAHeatmap, and produce offside within x
+    and y directions.
+
+    Note:
+
+        - instance number: N
+        - keypoint number: K
+        - keypoint dimension: D
+        - image size: [w, h]
+        - heatmap size: [W, H]
+
+    Encoded:
+
+        - heatmaps (np.ndarray): The generated heatmap in shape (K, H, W)
+            where [W, H] is the `heatmap_size`
+        - offside (np.ndarray): The generated offside in x andy direction
+            in shape (2K, H, W) where [W, H] is the `heatmap_size`
+        - keypoint_weights (np.ndarray): The target weights in shape (N, K)
+
+    Args:
+        input_size (tuple): Image size in [w, h]
+        heatmap_size (tuple): Heatmap size in [W, H]
+        sigma (float): The sigma value of the Gaussian heatmap
+    """
+
+    def __init__(self, input_size: Tuple[int, int],
+                 heatmap_size: Tuple[int, int], sigma: float) -> None:
+        super().__init__()
+        self.input_size = input_size
+        self.heatmap_size = heatmap_size
+        self.sigma = sigma
+        self.scale_factor = (np.array(input_size) /
+                             heatmap_size).astype(np.float32)
+
+        self.y_range, self.x_range = np.meshgrid(
+            np.arange(0, self.heatmap_size[1]),
+            np.arange(0, self.heatmap_size[0]),
+            indexing='ij')
+
+    def encode(self,
+               keypoints: np.ndarray,
+               keypoints_visible: Optional[np.ndarray] = None) -> dict:
+        """Encode keypoints into heatmaps. Note that the original keypoint
+        coordinates should be in the input image space.
+
+        Args:
+            keypoints (np.ndarray): Keypoint coordinates in shape (N, K, D)
+            keypoints_visible (np.ndarray): Keypoint visibilities in shape
+                (N, K)
+
+        Returns:
+            dict:
+            - heatmaps (np.ndarray): The generated heatmap in shape
+                (K, H, W) where [W, H] is the `heatmap_size`
+            - offside (np.ndarray): The generated offside in x and y
+                directions in shape (2*K, H, W) where [W, H] is the
+                `heatmap_size`
+            - keypoint_weights (np.ndarray): The target weights in shape
+                (N, K)
+        """
+
+        assert keypoints.shape[0] == 1, (
+            f'{self.__class__.__name__} only support single-instance '
+            'keypoint encoding')
+
+        if keypoints_visible is None:
+            keypoints_visible = np.ones(keypoints.shape[:2], dtype=np.float32)
+
+        heatmaps, keypoint_weights = generate_unbiased_gaussian_heatmaps(
+            heatmap_size=self.heatmap_size,
+            keypoints=keypoints / self.scale_factor,
+            keypoints_visible=keypoints_visible,
+            sigma=self.sigma)
+
+        offside_heatmap = self.generate_offside_heatmap(
+            heatmap_size=self.heatmap_size,
+            keypoints=keypoints / self.scale_factor,
+        )
+
+        encoded = dict(
+            heatmaps=heatmaps,
+            keypoint_weights=keypoint_weights[0],
+            displacements=offside_heatmap)
+
+        return encoded
+
+    def generate_offside_heatmap(
+        self,
+        heatmap_size: Tuple[int, int],
+        keypoints: np.ndarray,
+    ):
+
+        N, K, _ = keypoints.shape
+        W, H = heatmap_size
+
+        # batchsize 1
+        keypoints = keypoints[0]
+
+        # caution: there will be a broadcast which produce
+        # offside_x and offside_y with shape 64x64x98
+
+        offside_x = keypoints[:, 0] - np.expand_dims(self.x_range, axis=-1)
+        offside_y = keypoints[:, 1] - np.expand_dims(self.y_range, axis=-1)
+
+        offside_map = np.concatenate([offside_x, offside_y], axis=-1)
+
+        offside_map = np.transpose(offside_map, axes=[2, 0, 1])
+
+        return offside_map
+
+    def decode(self, encoded: np.ndarray,
+               offside: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        """Decode keypoint coordinates from heatmaps. The decoded keypoint
+        coordinates are in the input image space.
+
+        Args:
+            encoded (np.ndarray): Heatmaps in shape (K, H, W)
+
+        Returns:
+            tuple:
+            - keypoints (np.ndarray): Decoded keypoint coordinates in shape
+                (N, K, D)
+            - scores (np.ndarray): The keypoint scores in shape (N, K). It
+                usually represents the confidence of the keypoint prediction
+        """
+        heatmaps = encoded.copy()
+
+        offside = offside.copy()
+
+        K, H, W = heatmaps.shape
+
+        keypoints, scores = get_heatmap_maximum(heatmaps)
+
+        offside_x = offside[:K, ...]
+        offside_y = offside[K:, ...]
+
+        keypoints_interger = keypoints.astype(np.int32)
+        keypoints_decimal = np.zeros_like(keypoints)
+
+        for i in range(K):
+            [x, y] = keypoints_interger[i]
+            if x < 0 or y < 0:
+                x = y = 0
+
+            # caution: torhc tensor shape is nchw, so indx should be i,y,x
+            keypoints_decimal[i][0] = x + offside_x[i, y, x]
+            keypoints_decimal[i][1] = y + offside_y[i, y, x]
+
+        # Restore the keypoint scale
+        keypoints_decimal = keypoints_decimal * self.scale_factor
+
+        return keypoints_decimal[None], scores[None]