ADD file via upload

README.md

@@ -1,16 +1,69 @@
 # 基于 OpenPifPaf 的多摄像头、多人实时跌倒检测模型
-利用 OpenPifPaf 对输入视频进行人体姿势估计,然后通过长短时记忆神经网络（LSTM）从前面得到的姿势信息中提取五个时间和空间特征以预测"跌倒"动作,支持多摄像头和多人实时检测。
-## 检测实例见 examples 文件夹
-## 安装
+<p align="center">
+<img src="https://git.trustie.net/pkwhiuqat/HumanFallDetectionLSTM/raw/branch/master/examples/outfallingdown.gif?raw=true" alt="outfallingdown"/>
 
+利用 OpenPifPaf 对输入视频进行人体姿势估计,然后通过长短时记忆神经网络（LSTM）从前面得到的姿势信息中提取五个时间和空间特征（作为当前的 *X*<sub>n</sub> 输入）以预测"跌倒"动作,支持多摄像头和多人实时检测。模型在 UP-Fall Detection 数据集上训练，基于 PyTorch 实现。
+
+<p align="center">
+<img src="https://git.trustie.net/pkwhiuqat/HumanFallDetectionLSTM/raw/branch/master/flowchart.png?raw=true" alt="LSTM" style="zoom:68%;" />
+<p align="center">
+<img src="https://git.trustie.net/pkwhiuqat/HumanFallDetectionLSTM/raw/branch/master/LSTM.png?raw=true" alt="LSTM" style="zoom:45%;" />
+
+## 安装
 ```shell script
 pip install -r requirements.txt
 ```
 
 ## 使用
 ```shell script
-python3 fall_detector.py --num_cams=1
+python fall_detector.py --num_cams=1
 ```
 
+# Alphapose部分说明(AlphaPose-new分支)
+
+## Alphapose介绍
+
+AlphaPose采用自顶向下的方法，提出了RMPE（区域多人姿态检测）框架。该框架主要包括symmetric spatial transformer network (SSTN)、Parametric Pose Non- Maximum-Suppression (NMS)和Pose-Guided Proposals Generator (PGPG)。并且使用symmetric spatial transformer network (SSTN)、deep proposals generator (DPG) 、parametric pose nonmaximum suppression (p-NMS) 三个技术来解决野外场景下多人姿态估计问题。
+
+在SPPE结构上添加SSTN，能够在不精准的区域框中提取到高质量的人体区域。并行的SPPE分支（SSTN）来优化自身网络。使用parametric pose NMS来解决冗余检测问题，在该结构中，使用了自创的姿态距离度量方案比较姿态之间的相似度。用数据驱动的方法优化姿态距离参数。最后我们使用PGPG来强化训练数据，通过学习输出结果中不同姿态的描述信息，来模仿人体区域框的生成过程，进一步产生一个更大的训练集。
+
+## Alphapose在该项目的作用
+
+可以向Alphapose输入人体行为视频，输出.avi格式的附带人体骨架的视频，这种视频或许可以用来作为OpenPifPaf的输入素材使用。
+
+## Alphapose项目运行环境
+
+pytorch1.4
+
+cuda10.1
+
+cudnn7
+
+ubuntu18.04
+
+## Alphapose配置步骤
+
+1、在命令提示符里输入cuda的路径。
+
+export PATH=/usr/local/cuda/bin/:$PATH
+export LD_LIBRARY_PATH=/usr/local/cuda/lib64/:$LD_LIBRARY_PATH
+
+2、之后输入pip install .安装配置文件
+
+## Alphapose使用流程
+
+1、将视频存入特定文件夹，并记下视频路径。（建议存储在AlphaPose/data里面）
+
+2、在命令提示符里输入如下内容运行Alphapose：
+
+bash ./scripts/inference.sh configs/coco/resnet/256x192_res152_lr1e-3_1x-duc.yaml pretrained_models/fast_res50_256x192.pth 视频路径 outputs
+
+3、在AlphaPose/outputs文件夹里找到输出的内容
+
+注：若outputs文件夹里只生成了.json文件，则用记事本打开scripts/inference.sh，将里面最后一行的“#--save_video”中的“#”删除，之后再执行第二步即可。
+
 ## 参考
-https://github.com/openpifpaf/openpifpaf
+- [OpenPifPaf](https://github.com/openpifpaf/openpifpaf)
+- [UP-fall detection Dataset](https://dx.doi.org/10.3390/s19091988)
+- [Multi-camera, multi-person, and real-time fall detection using long short term memory](https://doi.org/10.1117/12.2580700)
+- [AlphaPose](https://github.com/MVIG-SJTU/AlphaPose)

alphapose/__init__.py

@@ -0,0 +1,3 @@
+from .version import __version__, short_version
+
+__all__ = ['__version__', 'short_version']

alphapose/datasets

@@ -0,0 +1,10 @@
+from .coco_det import Mscoco_det
+from .concat_dataset import ConcatDataset
+from .custom import CustomDataset
+from .mscoco import Mscoco
+from .mpii import Mpii
+from .halpe_26 import Halpe_26
+from .halpe_136 import Halpe_136
+from .halpe_136_det import  Halpe_136_det
+from .halpe_26_det import  Halpe_26_det
+__all__ = ['CustomDataset', 'Halpe_136', 'Halpe_26_det', 'Halpe_136_det', 'Halpe_26', 'Mscoco', 'Mscoco_det', 'Mpii', 'ConcatDataset', 'coco_wholebody', 'coco_wholebody_det']

alphapose/datasets1/__init__.py

@@ -0,0 +1,10 @@
+from .coco_det import Mscoco_det
+from .concat_dataset import ConcatDataset
+from .custom import CustomDataset
+from .mscoco import Mscoco
+from .mpii import Mpii
+from .halpe_26 import Halpe_26
+from .halpe_136 import Halpe_136
+from .halpe_136_det import  Halpe_136_det
+from .halpe_26_det import  Halpe_26_det
+__all__ = ['CustomDataset', 'Halpe_136', 'Halpe_26_det', 'Halpe_136_det', 'Halpe_26', 'Mscoco', 'Mscoco_det', 'Mpii', 'ConcatDataset', 'coco_wholebody', 'coco_wholebody_det']

alphapose/datasets1/coco_det.py

@@ -0,0 +1,107 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+
+"""MS COCO Human Detection Box dataset."""
+import json
+import os
+
+import cv2
+import torch
+import torch.utils.data as data
+from tqdm import tqdm
+
+from alphapose.utils.presets import SimpleTransform
+from detector.apis import get_detector
+from alphapose.models.builder import DATASET
+
+
+@DATASET.register_module
+class Mscoco_det(data.Dataset):
+    """ COCO human detection box dataset.
+
+    """
+    EVAL_JOINTS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
+
+    def __init__(self,
+                 det_file=None,
+                 opt=None,
+                 **cfg):
+
+        self._cfg = cfg
+        self._opt = opt
+        self._preset_cfg = cfg['PRESET']
+        self._root = cfg['ROOT']
+        self._img_prefix = cfg['IMG_PREFIX']
+        if not det_file:
+            det_file = cfg['DET_FILE']
+        self._ann_file = os.path.join(self._root, cfg['ANN'])
+
+        if os.path.exists(det_file):
+            print("Detection results exist, will use it")
+        else:
+            print("Will create detection results to {}".format(det_file))
+            self.write_coco_json(det_file)
+
+        assert os.path.exists(det_file), "Error: no detection results found"
+        with open(det_file, 'r') as fid:
+            self._det_json = json.load(fid)
+
+        self._input_size = self._preset_cfg['IMAGE_SIZE']
+        self._output_size = self._preset_cfg['HEATMAP_SIZE']
+
+        self._sigma = self._preset_cfg['SIGMA']
+
+        if self._preset_cfg['TYPE'] == 'simple':
+            self.transformation = SimpleTransform(
+                self, scale_factor=0,
+                input_size=self._input_size,
+                output_size=self._output_size,
+                rot=0, sigma=self._sigma,
+                train=False, add_dpg=False)
+
+    def __getitem__(self, index):
+        det_res = self._det_json[index]
+        if not isinstance(det_res['image_id'], int):
+            img_id, _ = os.path.splitext(os.path.basename(det_res['image_id']))
+            img_id = int(img_id)
+        else:
+            img_id = det_res['image_id']
+        img_path = './data/coco/val2017/%012d.jpg' % img_id
+
+        # Load image
+        image = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB) #scipy.misc.imread(img_path, mode='RGB')
+
+        imght, imgwidth = image.shape[0], image.shape[1]
+        x1, y1, w, h = det_res['bbox']
+        bbox = [x1, y1, x1 + w, y1 + h]
+        inp, bbox = self.transformation.test_transform(image, bbox)
+        return inp, torch.Tensor(bbox), torch.Tensor([det_res['bbox']]), torch.Tensor([det_res['image_id']]), torch.Tensor([det_res['score']]), torch.Tensor([imght]), torch.Tensor([imgwidth])
+
+    def __len__(self):
+        return len(self._det_json)
+
+    def write_coco_json(self, det_file):
+        from pycocotools.coco import COCO
+        import pathlib
+
+        _coco = COCO(self._ann_file)
+        image_ids = sorted(_coco.getImgIds())
+        det_model = get_detector(self._opt)
+        dets = []
+        for entry in tqdm(_coco.loadImgs(image_ids)):
+            abs_path = os.path.join(
+                self._root, self._img_prefix, entry['file_name'])
+            det = det_model.detect_one_img(abs_path)
+            if det:
+                dets += det
+        pathlib.Path(os.path.split(det_file)[0]).mkdir(parents=True, exist_ok=True)
+        json.dump(dets, open(det_file, 'w'))
+
+    @property
+    def joint_pairs(self):
+        """Joint pairs which defines the pairs of joint to be swapped
+        when the image is flipped horizontally."""
+        return [[1, 2], [3, 4], [5, 6], [7, 8],
+                [9, 10], [11, 12], [13, 14], [15, 16]]

alphapose/datasets1/coco_wholebody.py

@@ -0,0 +1,160 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Haoyi Zhu and Hao-Shu Fang
+# -----------------------------------------------------
+
+"""Halpe Full-Body(136 points) Human keypoint dataset."""
+import os
+
+import numpy as np
+from tkinter import _flatten
+
+from alphapose.models.builder import DATASET
+from alphapose.utils.bbox import bbox_clip_xyxy, bbox_xywh_to_xyxy
+
+from .custom import CustomDataset
+
+
+@DATASET.register_module
+class coco_wholebody(CustomDataset):
+    """ Halpe Full-Body(136 points) Person dataset.
+
+    Parameters
+    ----------
+    train: bool, default is True
+        If true, will set as training mode.
+    skip_empty: bool, default is False
+        Whether skip entire image if no valid label is found. Use `False` if this dataset is
+        for validation to avoid COCO metric error.
+    dpg: bool, default is False
+        If true, will activate `dpg` for data augmentation.
+    """
+    CLASSES = ['person']
+    EVAL_JOINTS = list(range(133))
+    num_joints = 133
+    CustomDataset.lower_body_ids = (11, 12, 13, 14, 15, 16, 17, 21-3, 22-3, 23-3, 24-3, 25-3)
+    """Joint pairs which defines the pairs of joint to be swapped
+        when the image is flipped horizontally."""
+    joint_pairs =  [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12], [13, 14], [15, 16], #17 body keypoints
+        [20-3, 23-3], [21-3, 24-3], [22-3, 25-3], [26-3, 42-3], [27-3, 41-3], [28-3, 40-3], [29-3, 39-3], [30-3, 38-3], 
+        [31-3, 37-3], [32-3, 36-3], [33-3, 35-3], [43-3, 52-3], [44-3, 51-3], [45-3, 50-3], [46-3, 49-3], [47-3, 48-3], 
+        [62-3, 71-3], [63-3, 70-3], [64-3, 69-3], [65-3, 68-3], [66-3, 73-3], [67-3, 72-3], [57-3, 61-3], [58-3, 60-3],
+        [74-3, 80-3], [75-3, 79-3], [76-3, 78-3], [87-3, 89-3], [93-3, 91-3], [86-3, 90-3], [85-3, 81-3], [84-3, 82-3],
+        [94-3, 115-3], [95-3, 116-3], [96-3, 117-3], [97-3, 118-3], [98-3, 119-3], [99-3, 120-3], [100-3, 121-3],
+        [101-3, 122-3], [102-3, 123-3], [103-3, 124-3], [104-3, 125-3], [105-3, 126-3], [106-3, 127-3], [107-3, 128-3],
+        [108-3, 129-3], [109-3, 130-3], [110-3, 131-3], [111-3, 132-3], [112-3, 133-3], [113-3, 134-3], [114-3, 135-3]]
+                
+
+    def _load_jsons(self):
+        """Load all image paths and labels from JSON annotation files into buffer."""
+        items = []
+        labels = []
+
+        _coco = self._lazy_load_ann_file()
+
+        classes = [c['name'] for c in _coco.loadCats(_coco.getCatIds())]
+        assert classes == self.CLASSES, "Incompatible category names with COCO. "
+
+        self.json_id_to_contiguous = {
+            v: k for k, v in enumerate(_coco.getCatIds())}
+
+        # iterate through the annotations
+        image_ids = sorted(_coco.getImgIds())
+        for entry in _coco.loadImgs(image_ids):
+            dirname, filename = entry['coco_url'].split('/')[-2:]
+            abs_path = os.path.join('/DATA1/Benchmark/coco', dirname, filename)
+            if not os.path.exists(abs_path):
+                raise IOError('Image: {} not exists.'.format(abs_path))
+            label = self._check_load_keypoints(_coco, entry)
+            if not label:
+                continue
+            for obj in label:
+                items.append(abs_path)
+                labels.append(obj)
+
+        return items, labels
+
+    def _check_load_keypoints(self, coco, entry):
+        """Check and load ground-truth keypoints"""
+        ann_ids = coco.getAnnIds(imgIds=entry['id'], iscrowd=False)
+        objs = coco.loadAnns(ann_ids)
+        # check valid bboxes
+        valid_objs = []
+        width = entry['width']
+        height = entry['height']
+
+        for obj in objs:
+            #obj['keypoints'].extend([0,0,0, 0,0,0, 0,0,0])
+            obj['keypoints'].extend(obj['foot_kpts'])
+            obj['keypoints'].extend(obj['face_kpts'])
+            obj['keypoints'].extend(obj['lefthand_kpts'])
+            obj['keypoints'].extend(obj['righthand_kpts'])
+            contiguous_cid = self.json_id_to_contiguous[obj['category_id']]
+            if contiguous_cid >= self.num_class:
+                # not class of interest
+                continue
+            if max(obj['keypoints']) == 0:
+                continue
+            # convert from (x, y, w, h) to (xmin, ymin, xmax, ymax) and clip bound
+            xmin, ymin, xmax, ymax = bbox_clip_xyxy(bbox_xywh_to_xyxy(obj['bbox']), width, height)
+            # require non-zero box area
+            #if obj['area'] <= 0 or xmax <= xmin or ymax <= ymin:
+            if (xmax-xmin)*(ymax-ymin) <= 0 or xmax <= xmin or ymax <= ymin:
+                continue
+            if 'num_keypoints' in obj and obj['num_keypoints'] == 0:
+                continue
+            # joints 3d: (num_joints, 3, 2); 3 is for x, y, z; 2 is for position, visibility
+            joints_3d = np.zeros((self.num_joints, 3, 2), dtype=np.float32)
+            for i in range(self.num_joints):
+                joints_3d[i, 0, 0] = obj['keypoints'][i * 3 + 0]
+                joints_3d[i, 1, 0] = obj['keypoints'][i * 3 + 1]
+                # joints_3d[i, 2, 0] = 0
+                if obj['keypoints'][i * 3 + 2] >= 0.35:
+                    visible = 1
+                else:
+                    visible = 0
+                #visible = min(1, visible)
+                joints_3d[i, :2, 1] = visible
+                # joints_3d[i, 2, 1] = 0
+
+            if np.sum(joints_3d[:, 0, 1]) < 1:
+                # no visible keypoint
+                continue
+
+            if self._check_centers and self._train:
+                bbox_center, bbox_area = self._get_box_center_area((xmin, ymin, xmax, ymax))
+                kp_center, num_vis = self._get_keypoints_center_count(joints_3d)
+                ks = np.exp(-2 * np.sum(np.square(bbox_center - kp_center)) / bbox_area)
+                if (num_vis / 80.0 + 47 / 80.0) > ks:
+                    continue
+
+            valid_objs.append({
+                'bbox': (xmin, ymin, xmax, ymax),
+                'width': width,
+                'height': height,
+                'joints_3d': joints_3d
+            })
+
+        if not valid_objs:
+            if not self._skip_empty:
+                # dummy invalid labels if no valid objects are found
+                valid_objs.append({
+                    'bbox': np.array([-1, -1, 0, 0]),
+                    'width': width,
+                    'height': height,
+                    'joints_3d': np.zeros((self.num_joints, 2, 2), dtype=np.float32)
+                })
+        return valid_objs
+
+    def _get_box_center_area(self, bbox):
+        """Get bbox center"""
+        c = np.array([(bbox[0] + bbox[2]) / 2.0, (bbox[1] + bbox[3]) / 2.0])
+        area = (bbox[3] - bbox[1]) * (bbox[2] - bbox[0])
+        return c, area
+
+    def _get_keypoints_center_count(self, keypoints):
+        """Get geometric center of all keypoints"""
+        keypoint_x = np.sum(keypoints[:, 0, 0] * (keypoints[:, 0, 1] > 0))
+        keypoint_y = np.sum(keypoints[:, 1, 0] * (keypoints[:, 1, 1] > 0))
+        num = float(np.sum(keypoints[:, 0, 1]))
+        return np.array([keypoint_x / num, keypoint_y / num]), num

alphapose/datasets1/coco_wholebody_det.py

@@ -0,0 +1,113 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Haoyi Zhu
+# -----------------------------------------------------
+
+"""Haple_136 Human Detection Box dataset."""
+import json
+import os
+
+import cv2
+import torch
+import torch.utils.data as data
+from tqdm import tqdm
+
+from alphapose.utils.presets import SimpleTransform
+from detector.apis import get_detector
+from alphapose.models.builder import DATASET
+
+
+@DATASET.register_module
+class coco_wholebody_det(data.Dataset):
+    """ Halpe_136 human detection box dataset.
+
+    """
+    EVAL_JOINTS = list(range(133))
+
+    def __init__(self,
+                 det_file=None,
+                 opt=None,
+                 **cfg):
+
+        self._cfg = cfg
+        self._opt = opt
+        self._preset_cfg = cfg['PRESET']
+        self._root = cfg['ROOT']
+        self._img_prefix = cfg['IMG_PREFIX']
+        if not det_file:
+            det_file = cfg['DET_FILE']
+        self._ann_file = os.path.join(self._root, cfg['ANN'])
+
+        if os.path.exists(det_file):
+            print("Detection results exist, will use it")
+        else:
+            print("Will create detection results to {}".format(det_file))
+            self.write_coco_json(det_file)
+
+        assert os.path.exists(det_file), "Error: no detection results found"
+        with open(det_file, 'r') as fid:
+            self._det_json = json.load(fid)
+
+        self._input_size = self._preset_cfg['IMAGE_SIZE']
+        self._output_size = self._preset_cfg['HEATMAP_SIZE']
+
+        self._sigma = self._preset_cfg['SIGMA']
+
+        if self._preset_cfg['TYPE'] == 'simple':
+            self.transformation = SimpleTransform(
+                self, scale_factor=0,
+                input_size=self._input_size,
+                output_size=self._output_size,
+                rot=0, sigma=self._sigma,
+                train=False, add_dpg=False)
+
+    def __getitem__(self, index):
+        det_res = self._det_json[index]
+        if not isinstance(det_res['image_id'], int):
+            img_id, _ = os.path.splitext(os.path.basename(det_res['image_id']))
+            img_id = int(img_id)
+        else:
+            img_id = det_res['image_id']
+        img_path = '/DATA1/Benchmark/coco/val2017/%012d.jpg' % img_id
+
+        # Load image
+        image = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB) #scipy.misc.imread(img_path, mode='RGB')
+
+        imght, imgwidth = image.shape[1], image.shape[2]
+        x1, y1, w, h = det_res['bbox']
+        bbox = [x1, y1, x1 + w, y1 + h]
+        inp, bbox = self.transformation.test_transform(image, bbox)
+        return inp, torch.Tensor(bbox), torch.Tensor([det_res['bbox']]), torch.Tensor([det_res['image_id']]), torch.Tensor([det_res['score']]), torch.Tensor([imght]), torch.Tensor([imgwidth])
+
+    def __len__(self):
+        return len(self._det_json)
+
+    def write_coco_json(self, det_file):
+        from pycocotools.coco import COCO
+        import pathlib
+
+        _coco = COCO(self._ann_file)
+        image_ids = sorted(_coco.getImgIds())
+        det_model = get_detector(self._opt)
+        dets = []
+        for entry in tqdm(_coco.loadImgs(image_ids)):
+            abs_path = os.path.join(
+                '/DATA1/Benchmark/coco', self._img_prefix, entry['file_name'])
+            det = det_model.detect_one_img(abs_path)
+            if det:
+                dets += det
+        pathlib.Path(os.path.split(det_file)[0]).mkdir(parents=True, exist_ok=True)
+        json.dump(dets, open(det_file, 'w'))
+
+    @property
+    def joint_pairs(self):
+        """Joint pairs which defines the pairs of joint to be swapped
+        when the image is flipped horizontally."""
+        return [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12], [13, 14], [15, 16], #17 body keypoints
+        [20-3, 23-3], [21-3, 24-3], [22-3, 25-3], [26-3, 42-3], [27-3, 41-3], [28-3, 40-3], [29-3, 39-3], [30-3, 38-3], 
+        [31-3, 37-3], [32-3, 36-3], [33-3, 35-3], [43-3, 52-3], [44-3, 51-3], [45-3, 50-3], [46-3, 49-3], [47-3, 48-3], 
+        [62-3, 71-3], [63-3, 70-3], [64-3, 69-3], [65-3, 68-3], [66-3, 73-3], [67-3, 72-3], [57-3, 61-3], [58-3, 60-3],
+        [74-3, 80-3], [75-3, 79-3], [76-3, 78-3], [87-3, 89-3], [93-3, 91-3], [86-3, 90-3], [85-3, 81-3], [84-3, 82-3],
+        [94-3, 115-3], [95-3, 116-3], [96-3, 117-3], [97-3, 118-3], [98-3, 119-3], [99-3, 120-3], [100-3, 121-3],
+        [101-3, 122-3], [102-3, 123-3], [103-3, 124-3], [104-3, 125-3], [105-3, 126-3], [106-3, 127-3], [107-3, 128-3],
+        [108-3, 129-3], [109-3, 130-3], [110-3, 131-3], [111-3, 132-3], [112-3, 133-3], [113-3, 134-3], [114-3, 135-3]]

alphapose/datasets1/concat_dataset.py

@@ -0,0 +1,77 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+
+import bisect
+
+import torch
+import torch.utils.data as data
+
+from alphapose.models.builder import DATASET, build_dataset
+
+
+@DATASET.register_module
+class ConcatDataset(data.Dataset):
+    """Custom Concat dataset.
+    Annotation file must be in `coco` format.
+
+    Parameters
+    ----------
+    train: bool, default is True
+        If true, will set as training mode.
+    dpg: bool, default is False
+        If true, will activate `dpg` for data augmentation.
+    skip_empty: bool, default is False
+        Whether skip entire image if no valid label is found.
+    cfg: dict, dataset configuration.
+    """
+
+    def __init__(self,
+                 train=True,
+                 dpg=False,
+                 skip_empty=True,
+                 **cfg):
+
+        self._cfg = cfg
+        self._subset_cfg_list = cfg['SET_LIST']
+        self._preset_cfg = cfg['PRESET']
+        self._mask_id = [item['MASK_ID'] for item in self._subset_cfg_list]
+
+        self.num_joints = self._preset_cfg['NUM_JOINTS']
+
+        self._subsets = []
+        self._subset_size = [0]
+        for _subset_cfg in self._subset_cfg_list:
+            subset = build_dataset(_subset_cfg, preset_cfg=self._preset_cfg, train=train)
+            self._subsets.append(subset)
+            self._subset_size.append(len(subset))
+        self.cumulative_sizes = self.cumsum(self._subset_size)
+
+    def __getitem__(self, idx):
+        assert idx >= 0
+        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
+        dataset_idx -= 1
+        sample_idx = idx - self.cumulative_sizes[dataset_idx]
+
+        sample = self._subsets[dataset_idx][sample_idx]
+        img, label, label_mask, img_id, bbox = sample
+
+        K = label.shape[0]  # num_joints from `_subsets[dataset_idx]`
+        expend_label = torch.zeros((self.num_joints, *label.shape[1:]), dtype=label.dtype)
+        expend_label_mask = torch.zeros((self.num_joints, *label_mask.shape[1:]), dtype=label_mask.dtype)
+        expend_label[self._mask_id[dataset_idx]:self._mask_id[dataset_idx] + K] = label
+        expend_label_mask[self._mask_id[dataset_idx]:self._mask_id[dataset_idx] + K] = label_mask
+
+        return img, expend_label, expend_label_mask, img_id, bbox
+
+    def __len__(self):
+        return self.cumulative_sizes[-1]
+
+    @staticmethod
+    def cumsum(sequence):
+        r, s = [], 0
+        for e in sequence:
+            r.append(e + s)
+            s += e
+        return r

alphapose/datasets1/custom.py

@@ -0,0 +1,292 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com), Haoyi Zhu
+# -----------------------------------------------------
+
+"""Custum training dataset."""
+import copy
+import os
+import pickle as pk
+from abc import abstractmethod, abstractproperty
+
+import torch.utils.data as data
+from pycocotools.coco import COCO
+
+from alphapose.utils.presets import SimpleTransform
+
+import cv2
+import json
+import numpy as np
+import random
+
+class CustomDataset(data.Dataset):
+    """Custom dataset.
+    Annotation file must be in `coco` format.
+
+    Parameters
+    ----------
+    train: bool, default is True
+        If true, will set as training mode.
+    dpg: bool, default is False
+        If true, will activate `dpg` for data augmentation.
+    skip_empty: bool, default is False
+        Whether skip entire image if no valid label is found.
+    cfg: dict, dataset configuration.
+    """
+
+    CLASSES = None
+
+    def __init__(self,
+                 train=True,
+                 dpg=False,
+                 skip_empty=True,
+                 lazy_import=False,
+                 **cfg):
+        if os.path.exists('/home/group3/background.json'):
+            self.bgim = json.load(open('/home/group3/background.json','r'))
+        else:
+            self.bgim = None
+
+        self._cfg = cfg
+        self._preset_cfg = cfg['PRESET']
+        self._root = cfg['ROOT']
+        self._img_prefix = cfg['IMG_PREFIX']
+        self._ann_file = os.path.join(self._root, cfg['ANN'])
+
+        self._lazy_import = lazy_import
+        self._skip_empty = skip_empty
+        self._train = train
+        self._dpg = dpg
+
+        if 'AUG' in cfg.keys():
+            self._scale_factor = cfg['AUG']['SCALE_FACTOR']
+            self._rot = cfg['AUG']['ROT_FACTOR']
+            self.num_joints_half_body = cfg['AUG']['NUM_JOINTS_HALF_BODY']
+            self.prob_half_body = cfg['AUG']['PROB_HALF_BODY']
+        else:
+            self._scale_factor = 0
+            self._rot = 0
+            self.num_joints_half_body = -1
+            self.prob_half_body = -1
+
+        self._input_size = self._preset_cfg['IMAGE_SIZE']
+        self._output_size = self._preset_cfg['HEATMAP_SIZE']
+
+        self._sigma = self._preset_cfg['SIGMA']
+
+        self._check_centers = False
+
+        self.num_class = len(self.CLASSES)
+
+        self._loss_type = self._preset_cfg.get('LOSS_TYPE', 'MSELoss')
+
+        self.upper_body_ids = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10)
+        self.lower_body_ids = (11, 12, 13, 14, 15, 16)
+
+        if self._preset_cfg['TYPE'] == 'simple':
+            self.transformation = SimpleTransform(
+                self, scale_factor=self._scale_factor,
+                input_size=self._input_size,
+                output_size=self._output_size,
+                rot=self._rot, sigma=self._sigma,
+                train=self._train, add_dpg=self._dpg,
+                loss_type=self._loss_type)
+        else:
+            raise NotImplementedError
+
+        self._items, self._labels = self._lazy_load_json()
+
+    def __getitem__(self, idx):
+        source = None
+        # get image id
+        if type(self._items[idx]) == dict:
+            img_path = self._items[idx]['path']
+            img_id = self._items[idx]['id']
+            source = self._items[idx]['source']
+        else:
+            img_path = self._items[idx]
+            img_id = int(os.path.splitext(os.path.basename(img_path))[0])
+
+        # load ground truth, including bbox, keypoints, image size
+        label = copy.deepcopy(self._labels[idx])
+        img = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)
+        
+        if self.bgim and (source == 'frei' or source == 'partX' or source == 'OneHand' or source == 'interhand'): # hand
+            img, label = self.hand_augmentation(img, label)
+            
+        if source == 'hand_labels_synth' or source == 'hand143_panopticdb': # hand
+            if not self.skip_augmentation(0.8):
+                img = self.motion_blur(img)
+        
+        # transform ground truth into training label and apply data augmentation
+        img, label, label_mask, bbox = self.transformation(img, label, source)
+        return img, label, label_mask, img_id, bbox
+
+    def __len__(self):
+        return len(self._items)
+
+    def _lazy_load_ann_file(self):
+        if os.path.exists(self._ann_file + '.pkl') and self._lazy_import:
+            print('Lazy load json...')
+            with open(self._ann_file + '.pkl', 'rb') as fid:
+                return pk.load(fid)
+        else:
+            _database = COCO(self._ann_file)
+            if os.access(self._ann_file + '.pkl', os.W_OK):
+                with open(self._ann_file + '.pkl', 'wb') as fid:
+                    pk.dump(_database, fid, pk.HIGHEST_PROTOCOL)
+            return _database
+
+    def _lazy_load_json(self):
+        if os.path.exists(self._ann_file + '_annot_keypoint.pkl') and self._lazy_import:
+            print('Lazy load annot...')
+            with open(self._ann_file + '_annot_keypoint.pkl', 'rb') as fid:
+                items, labels = pk.load(fid)
+        else:
+            items, labels = self._load_jsons()
+            if os.access(self._ann_file + '_annot_keypoint.pkl', os.W_OK):
+                with open(self._ann_file + '_annot_keypoint.pkl', 'wb') as fid:
+                    pk.dump((items, labels), fid, pk.HIGHEST_PROTOCOL)
+
+        return items, labels
+
+    def motion_blur(self, image, degree=12, angle=45):
+        image = np.array(image)
+     
+        M = cv2.getRotationMatrix2D((degree / 2, degree / 2), angle, 1)
+        motion_blur_kernel = np.diag(np.ones(degree))
+        motion_blur_kernel = cv2.warpAffine(motion_blur_kernel, M, (degree, degree))
+     
+        motion_blur_kernel = motion_blur_kernel / degree
+        blurred = cv2.filter2D(image, -1, motion_blur_kernel)
+     
+        # convert to uint8
+        cv2.normalize(blurred, blurred, 0, 255, cv2.NORM_MINMAX)
+        blurred = np.array(blurred, dtype=np.uint8)
+        return blurred
+
+    def skip_augmentation(self, p):
+        x = np.random.rand()
+        if x < p:
+            return True
+        else:
+            return False
+
+    def get_bgimg(self, box_h, box_w):
+        bgimgpath = random.choice(self.bgim)
+        file_name = bgimgpath['file_name']
+        img_name = file_name.split('/')[-1]
+        img_path = '/home/group3/coco/train2017/' + img_name
+        img = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)
+        img_h, img_w = img.shape[0], img.shape[1]
+
+        if img_h <= box_h or img_w <= box_w:
+            img = cv2.resize(img, (int(max(img_w, (np.random.rand()*1.8+1.2) * box_w)), int(max(img_h, (np.random.rand()*1.8+1.2) * box_h))))
+            img_h, img_w = img.shape[0], img.shape[1]
+
+        # crop the img
+        if (img_h <= img_w) and (img_h - 4 > box_w) and (np.random.rand() > 0.25):
+            crop_w = int((img_h - 2 - max((box_w + 2), (img_h / 3))) * np.random.rand() + max((box_w + 2), (img_h / 3)))
+            start_p = (img_w - crop_w) * np.random.rand()
+            img = img[:, int(start_p):int(start_p + crop_w + 1), :]
+            assert img.shape[1] > box_w and img.shape[0] > img.shape[1], (img.shape, (box_w, box_h))
+
+        assert img.shape[0] > box_h and img.shape[1] > box_w, (img.shape, (box_h, box_w))
+
+        return img
+
+    def hand_augmentation(self, img, label):
+        # some images are too big (mainly in OneHand)
+        if img.shape[0] > 640 or img.shape[1] > 640:
+            h, w, c = img.shape
+            resize_scale = 640 / h
+            img = cv2.resize(img, (int(w * resize_scale), int(h * resize_scale)))
+            handkp = label['joints_3d'][:,0:2,0][115:136,:]
+            assert handkp.shape == (21, 2)
+            handkp = handkp * resize_scale
+            label['joints_3d'][:,0:2,0][115:136,:] = handkp
+            label['bbox'] = list(np.array(label['bbox']) * resize_scale)
+            label['height'], label['width'] = img.shape[0:2]
+        
+        if not self.skip_augmentation(0.8):
+            img = self.motion_blur(img)
+
+        label['bbox'] = list(label['bbox'])
+        # print(img_path, 'hand augmentation')
+        if not self.skip_augmentation(0.3):
+            handkp = label['joints_3d'][:,0:2,0][115:136,:]
+            assert handkp.shape == (21, 2)
+            
+            # resize the hand img (random scale between 40% and 100%)
+            resize_scale = 0.6 * np.random.rand() + 0.4
+            handkp = handkp * resize_scale
+            img = cv2.resize(img, dsize=None, fx=resize_scale, fy=resize_scale)
+            label['height'], label['width'] = img.shape[0:2]
+            label['bbox'] = list(np.array(label['bbox']) * resize_scale)
+
+            h, w = img.shape[0:2]
+            hand_xmin, hand_xmax, hand_ymin, hand_ymax = int(round(min(handkp[:,0]))), int(round(max(handkp[:,0]))),int(round(min(handkp[:,1]))), int(round(max(handkp[:,1])))
+            boxw_time, boxh_time = float(np.random.rand()*2+4), float(np.random.rand()*4+5)
+            box_w, box_h = max(int((hand_xmax - hand_xmin)*boxw_time), w+1), max(int((hand_ymax - hand_ymin)*boxh_time),h+1)
+
+            background= self.get_bgimg(box_h, box_w)
+
+            bh, bw, bc = background.shape
+            # print(bw - box_w, bh - box_h)
+            x, y = int(np.random.randint(0,int(bw - box_w),size=1)), int(np.random.randint(0,int(bh - box_h), size=1))
+            hd = copy.deepcopy(img)
+            new_image = copy.deepcopy(background)
+            ralative_x, ralative_y = int(np.random.randint(0,int(box_w-w),size=1)), int(np.random.randint(0,int(box_h-h), size=1))
+            new_loc_x, new_loc_y = x + ralative_x, y + ralative_y
+            assert  (new_loc_x+w < x+box_w) and (new_loc_y+h < y+box_h)
+
+            handkp[(handkp[:, 0] + handkp[:, 1]) > 0] += [new_loc_x, new_loc_y]
+
+            if new_loc_x < 0:
+                hd = hd[:,-new_loc_x:,:]
+                new_loc_x = 0
+            if new_loc_y < 0:
+                hd = hd[-new_loc_y:,:,:]
+                new_loc_y = 0
+            if new_loc_x+hd.shape[1]>new_image.shape[1]:
+                hd = hd[:, :new_image.shape[1]-new_loc_x, :]
+            if new_loc_y+hd.shape[0]>new_image.shape[0]:
+                hd = hd[:new_image.shape[0]-new_loc_y, :, :]
+
+            new_image[new_loc_y:new_loc_y+h,new_loc_x:new_loc_x+w,:] = hd
+            label['bbox'][0] = label['bbox'][0] + new_loc_x
+            label['bbox'][1] = label['bbox'][1] + new_loc_y
+            label['bbox'][2] = label['bbox'][2] + new_loc_x
+            label['bbox'][3] = label['bbox'][3] + new_loc_y
+
+            max_length = max(new_image.shape[0], new_image.shape[1])
+            if max_length > 1200:
+                scale = 640 / max_length
+                new_image = cv2.resize(new_image, (int(round(new_image.shape[1] * scale)), int(round(new_image.shape[0] * scale))))
+                handkp = handkp * scale
+            label['joints_3d'][:,0:2,0][115:136,:] = handkp
+            img = new_image
+            label['height'], label['width'] = img.shape[0:2]
+
+        label['bbox'] = tuple(label['bbox'])
+        assert label['height'] == img.shape[0] and label['width'] == img.shape[1], (img.shape, (label['height'], label['width']), flag)
+
+        return img, label
+
+    @abstractmethod
+    def _load_jsons(self):
+        pass
+
+    @abstractproperty
+    def CLASSES(self):
+        return None
+
+    @abstractproperty
+    def num_joints(self):
+        return None
+
+    @abstractproperty
+    def joint_pairs(self):
+        """Joint pairs which defines the pairs of joint to be swapped
+        when the image is flipped horizontally."""
+        return None

alphapose/datasets1/halpe_136.py

@@ -0,0 +1,185 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Haoyi Zhu and Hao-Shu Fang
+# -----------------------------------------------------
+
+"""Halpe Full-Body(136 points) Human keypoint dataset."""
+import os
+
+import numpy as np
+from tkinter import _flatten
+
+from alphapose.models.builder import DATASET
+from alphapose.utils.bbox import bbox_clip_xyxy, bbox_xywh_to_xyxy
+
+from .custom import CustomDataset
+
+
+@DATASET.register_module
+class Halpe_136(CustomDataset):
+    """ Halpe Full-Body(136 points) Person dataset.
+
+    Parameters
+    ----------
+    train: bool, default is True
+        If true, will set as training mode.
+    skip_empty: bool, default is False
+        Whether skip entire image if no valid label is found. Use `False` if this dataset is
+        for validation to avoid COCO metric error.
+    dpg: bool, default is False
+        If true, will activate `dpg` for data augmentation.
+    """
+    CLASSES = ['person']
+    EVAL_JOINTS = list(range(136))
+    num_joints = 136
+    CustomDataset.lower_body_ids = (11, 12, 13, 14, 15, 16, 20, 21, 22, 23, 24, 25)
+    """Joint pairs which defines the pairs of joint to be swapped
+        when the image is flipped horizontally."""
+    joint_pairs =  [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12], [13, 14], [15, 16], #17 body keypoints
+        [20, 21], [22, 23], [24, 25], [26, 42], [27, 41], [28, 40], [29, 39], [30, 38], 
+        [31, 37], [32, 36], [33, 35], [43, 52], [44, 51], [45, 50],[46, 49], [47, 48], 
+        [62, 71], [63, 70], [64, 69], [65, 68], [66, 73], [67, 72], [57, 61], [58, 60],
+        [74, 80], [75, 79], [76, 78], [87, 89], [93, 91], [86, 90], [85, 81], [84, 82],
+        [94, 115], [95, 116], [96, 117], [97, 118], [98, 119], [99, 120], [100, 121],
+        [101, 122], [102, 123], [103, 124], [104, 125], [105, 126], [106, 127], [107, 128],
+        [108, 129], [109, 130], [110, 131], [111, 132], [112, 133], [113, 134], [114, 135]]
+                
+
+    def _load_jsons(self):
+        """Load all image paths and labels from JSON annotation files into buffer."""
+        items = []
+        labels = []
+
+        _coco = self._lazy_load_ann_file()
+
+        classes = [c['name'] for c in _coco.loadCats(_coco.getCatIds())]
+        assert classes == self.CLASSES, "Incompatible category names with COCO. "
+
+        self.json_id_to_contiguous = {
+            v: k for k, v in enumerate(_coco.getCatIds())}
+
+        # iterate through the annotations
+        image_ids = sorted(_coco.getImgIds())
+        for entry in _coco.loadImgs(image_ids):
+            
+            if 'source' not in entry: # coco
+                dirname, filename = entry['coco_url'].split('/')[-2:]
+                abs_path = os.path.join('/DATA1/Benchmark/coco', dirname, filename)
+                if not os.path.exists(abs_path):
+                    raise IOError('Image: {} not exists.'.format(abs_path))
+                label = self._check_load_keypoints(_coco, entry)
+                if not label:
+                    continue
+                for obj in label:
+                    items.append(abs_path)
+                    labels.append(obj)
+            else:
+                source = entry['source']
+                if source == 'hico':
+                    abs_path = os.path.join('/DATA1/Benchmark/hico_20160224_det/images/train2015', entry['file_name'])
+                elif source == '300wLP':
+                    abs_path = os.path.join('/DATA1/Benchmark/300W_LP', entry['file_name'])
+                elif source == 'frei':
+                    abs_path = os.path.join('/DATA1/Benchmark/FreiHand/training/rgb', entry['file_name'])
+                
+                if not os.path.exists(abs_path):
+                    raise IOError('Image: {} not exists.'.format(abs_path))
+                label = self._check_load_keypoints(_coco, entry)
+                if not label:
+                    continue
+
+                # num of items are relative to person, not image
+                if source == 'hico':
+                    for i in range(6):
+                        for obj in label:
+                            items.append({'path': abs_path, 'id': entry['id'], 'source':source})
+                            labels.append(obj)
+                elif source == 'frei':
+                    for obj in label:
+                        items.append({'path': abs_path, 'id': entry['id'], 'source':source})
+                        labels.append(obj)
+                for obj in label:
+                    items.append({'path': abs_path, 'id': entry['id'], 'source':source})
+                    labels.append(obj)
+
+        return items, labels
+
+    def _check_load_keypoints(self, coco, entry):
+        """Check and load ground-truth keypoints"""
+        ann_ids = coco.getAnnIds(imgIds=entry['id'], iscrowd=False)
+        objs = coco.loadAnns(ann_ids)
+        # check valid bboxes
+        valid_objs = []
+        width = entry['width']
+        height = entry['height']
+
+        for obj in objs:
+            contiguous_cid = self.json_id_to_contiguous[obj['category_id']]
+            if contiguous_cid >= self.num_class:
+                # not class of interest
+                continue
+            if max(obj['keypoints']) == 0:
+                continue
+            # convert from (x, y, w, h) to (xmin, ymin, xmax, ymax) and clip bound
+            xmin, ymin, xmax, ymax = bbox_clip_xyxy(bbox_xywh_to_xyxy(obj['bbox']), width, height)
+            # require non-zero box area
+            #if obj['area'] <= 0 or xmax <= xmin or ymax <= ymin:
+            if (xmax-xmin)*(ymax-ymin) <= 0 or xmax <= xmin or ymax <= ymin:
+                continue
+            if 'num_keypoints' in obj and obj['num_keypoints'] == 0:
+                continue
+            # joints 3d: (num_joints, 3, 2); 3 is for x, y, z; 2 is for position, visibility
+            joints_3d = np.zeros((self.num_joints, 3, 2), dtype=np.float32)
+            for i in range(self.num_joints):
+                joints_3d[i, 0, 0] = obj['keypoints'][i * 3 + 0]
+                joints_3d[i, 1, 0] = obj['keypoints'][i * 3 + 1]
+                # joints_3d[i, 2, 0] = 0
+                if obj['keypoints'][i * 3 + 2] >= 0.35:
+                    visible = 1
+                else:
+                    visible = 0
+                #visible = min(1, visible)
+                joints_3d[i, :2, 1] = visible
+                # joints_3d[i, 2, 1] = 0
+
+            if np.sum(joints_3d[:, 0, 1]) < 1:
+                # no visible keypoint
+                continue
+
+            if self._check_centers and self._train:
+                bbox_center, bbox_area = self._get_box_center_area((xmin, ymin, xmax, ymax))
+                kp_center, num_vis = self._get_keypoints_center_count(joints_3d)
+                ks = np.exp(-2 * np.sum(np.square(bbox_center - kp_center)) / bbox_area)
+                if (num_vis / 80.0 + 47 / 80.0) > ks:
+                    continue
+
+            valid_objs.append({
+                'bbox': (xmin, ymin, xmax, ymax),
+                'width': width,
+                'height': height,
+                'joints_3d': joints_3d
+            })
+
+        if not valid_objs:
+            if not self._skip_empty:
+                # dummy invalid labels if no valid objects are found
+                valid_objs.append({
+                    'bbox': np.array([-1, -1, 0, 0]),
+                    'width': width,
+                    'height': height,
+                    'joints_3d': np.zeros((self.num_joints, 2, 2), dtype=np.float32)
+                })
+        return valid_objs
+
+    def _get_box_center_area(self, bbox):
+        """Get bbox center"""
+        c = np.array([(bbox[0] + bbox[2]) / 2.0, (bbox[1] + bbox[3]) / 2.0])
+        area = (bbox[3] - bbox[1]) * (bbox[2] - bbox[0])
+        return c, area
+
+    def _get_keypoints_center_count(self, keypoints):
+        """Get geometric center of all keypoints"""
+        keypoint_x = np.sum(keypoints[:, 0, 0] * (keypoints[:, 0, 1] > 0))
+        keypoint_y = np.sum(keypoints[:, 1, 0] * (keypoints[:, 1, 1] > 0))
+        num = float(np.sum(keypoints[:, 0, 1]))
+        return np.array([keypoint_x / num, keypoint_y / num]), num

alphapose/datasets1/halpe_136_det.py

@@ -0,0 +1,113 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Haoyi Zhu
+# -----------------------------------------------------
+
+"""Haple_136 Human Detection Box dataset."""
+import json
+import os
+
+import cv2
+import torch
+import torch.utils.data as data
+from tqdm import tqdm
+
+from alphapose.utils.presets import SimpleTransform
+from detector.apis import get_detector
+from alphapose.models.builder import DATASET
+
+
+@DATASET.register_module
+class Halpe_136_det(data.Dataset):
+    """ Halpe_136 human detection box dataset.
+
+    """
+    EVAL_JOINTS = list(range(136))
+
+    def __init__(self,
+                 det_file=None,
+                 opt=None,
+                 **cfg):
+
+        self._cfg = cfg
+        self._opt = opt
+        self._preset_cfg = cfg['PRESET']
+        self._root = cfg['ROOT']
+        self._img_prefix = cfg['IMG_PREFIX']
+        if not det_file:
+            det_file = cfg['DET_FILE']
+        self._ann_file = os.path.join(self._root, cfg['ANN'])
+
+        if os.path.exists(det_file):
+            print("Detection results exist, will use it")
+        else:
+            print("Will create detection results to {}".format(det_file))
+            self.write_coco_json(det_file)
+
+        assert os.path.exists(det_file), "Error: no detection results found"
+        with open(det_file, 'r') as fid:
+            self._det_json = json.load(fid)
+
+        self._input_size = self._preset_cfg['IMAGE_SIZE']
+        self._output_size = self._preset_cfg['HEATMAP_SIZE']
+
+        self._sigma = self._preset_cfg['SIGMA']
+
+        if self._preset_cfg['TYPE'] == 'simple':
+            self.transformation = SimpleTransform(
+                self, scale_factor=0,
+                input_size=self._input_size,
+                output_size=self._output_size,
+                rot=0, sigma=self._sigma,
+                train=False, add_dpg=False)
+
+    def __getitem__(self, index):
+        det_res = self._det_json[index]
+        if not isinstance(det_res['image_id'], int):
+            img_id, _ = os.path.splitext(os.path.basename(det_res['image_id']))
+            img_id = int(img_id)
+        else:
+            img_id = det_res['image_id']
+        img_path = '/DATA1/Benchmark/coco/val2017/%012d.jpg' % img_id
+
+        # Load image
+        image = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB) #scipy.misc.imread(img_path, mode='RGB')
+
+        imght, imgwidth = image.shape[1], image.shape[2]
+        x1, y1, w, h = det_res['bbox']
+        bbox = [x1, y1, x1 + w, y1 + h]
+        inp, bbox = self.transformation.test_transform(image, bbox)
+        return inp, torch.Tensor(bbox), torch.Tensor([det_res['bbox']]), torch.Tensor([det_res['image_id']]), torch.Tensor([det_res['score']]), torch.Tensor([imght]), torch.Tensor([imgwidth])
+
+    def __len__(self):
+        return len(self._det_json)
+
+    def write_coco_json(self, det_file):
+        from pycocotools.coco import COCO
+        import pathlib
+
+        _coco = COCO(self._ann_file)
+        image_ids = sorted(_coco.getImgIds())
+        det_model = get_detector(self._opt)
+        dets = []
+        for entry in tqdm(_coco.loadImgs(image_ids)):
+            abs_path = os.path.join(
+                '/DATA1/Benchmark/coco', self._img_prefix, entry['file_name'])
+            det = det_model.detect_one_img(abs_path)
+            if det:
+                dets += det
+        pathlib.Path(os.path.split(det_file)[0]).mkdir(parents=True, exist_ok=True)
+        json.dump(dets, open(det_file, 'w'))
+
+    @property
+    def joint_pairs(self):
+        """Joint pairs which defines the pairs of joint to be swapped
+        when the image is flipped horizontally."""
+        return[[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12], [13, 14], [15, 16], 
+        [20, 21], [22, 23], [24, 25], [26, 42], [27, 41], [28, 40], [29, 39], [30, 38], 
+        [31, 37], [32, 36], [33, 35], [43, 52], [44, 51], [45, 50],[46, 49], [47, 48], 
+        [62, 71], [63, 70], [64, 69], [65, 68], [66, 73], [67, 72], [57, 61], [58, 60],
+        [74, 80], [75, 79], [76, 78], [87, 89], [93, 91], [86, 90], [85, 81], [84, 82],
+        [94, 115], [95, 116], [96, 117], [97, 118], [98, 119], [99, 120], [100, 121],
+        [101, 122], [102, 123], [103, 124], [104, 125], [105, 126], [106, 127], [107, 128],
+        [108, 129], [109, 130], [110, 131], [111, 132], [112, 133], [113, 134], [114, 135]]

alphapose/datasets1/halpe_26.py

@@ -0,0 +1,173 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Haoyi Zhu and Hao-Shu Fang
+# -----------------------------------------------------
+
+"""Halpe Human keypoint(26 points version) dataset."""
+import os
+
+import numpy as np
+from tkinter import _flatten
+
+from alphapose.models.builder import DATASET
+from alphapose.utils.bbox import bbox_clip_xyxy, bbox_xywh_to_xyxy
+
+from .custom import CustomDataset
+
+
+@DATASET.register_module
+class Halpe_26(CustomDataset):
+    """ Halpe_simple 26 keypoints Person Pose dataset.
+
+    Parameters
+    ----------
+    train: bool, default is True
+        If true, will set as training mode.
+    skip_empty: bool, default is False
+        Whether skip entire image if no valid label is found. Use `False` if this dataset is
+        for validation to avoid COCO metric error.
+    dpg: bool, default is False
+        If true, will activate `dpg` for data augmentation.
+    """
+    CLASSES = ['person']
+    EVAL_JOINTS = list(range(26))
+    num_joints = 26
+    CustomDataset.lower_body_ids = (11, 12, 13, 14, 15, 16, 20, 21, 22, 23, 24, 25)
+    joint_pairs = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12], [13, 14], [15, 16],
+        [20, 21], [22, 23], [24, 25]]
+
+    def _load_jsons(self):
+        """Load all image paths and labels from JSON annotation files into buffer."""
+        items = []
+        labels = []
+
+        _coco = self._lazy_load_ann_file()
+
+        classes = [c['name'] for c in _coco.loadCats(_coco.getCatIds())]
+        assert classes == self.CLASSES, "Incompatible category names with COCO. "
+
+        self.json_id_to_contiguous = {
+            v: k for k, v in enumerate(_coco.getCatIds())}
+
+        # iterate through the annotations
+        image_ids = sorted(_coco.getImgIds())
+        for entry in _coco.loadImgs(image_ids):
+            
+            if 'source' not in entry: # coco
+                dirname, filename = entry['coco_url'].split('/')[-2:]
+                abs_path = os.path.join('/DATA1/Benchmark/coco', dirname, filename)
+                if not os.path.exists(abs_path):
+                    raise IOError('Image: {} not exists.'.format(abs_path))
+                label = self._check_load_keypoints(_coco, entry)
+                if not label:
+                    continue
+                for obj in label:
+                    items.append(abs_path)
+                    labels.append(obj)
+            else:
+                source = entry['source']
+                if source == 'hico':
+                    abs_path = os.path.join('/DATA1/Benchmark/hico_20160224_det/images/train2015', entry['file_name'])
+                elif source == '300wLP':
+                    abs_path = os.path.join('/DATA1/Benchmark/300W_LP', entry['file_name'])
+                elif source == 'frei':
+                    abs_path = os.path.join('/DATA1/Benchmark/FreiHand/training/rgb', entry['file_name'])
+                
+                if not os.path.exists(abs_path):
+                    raise IOError('Image: {} not exists.'.format(abs_path))
+                label = self._check_load_keypoints(_coco, entry)
+                if not label:
+                    continue
+
+                # num of items are relative to person, not image
+                if source == 'hico':
+                    for i in range(6):
+                        for obj in label:
+                            items.append({'path': abs_path, 'id': entry['id'], 'source':source})
+                            labels.append(obj)
+                    
+                for obj in label:
+                    items.append({'path': abs_path, 'id': entry['id'], 'source':source})
+                    labels.append(obj)
+
+        return items, labels
+
+    def _check_load_keypoints(self, coco, entry):
+        """Check and load ground-truth keypoints"""
+        ann_ids = coco.getAnnIds(imgIds=entry['id'], iscrowd=False)
+        objs = coco.loadAnns(ann_ids)
+        # check valid bboxes
+        valid_objs = []
+        width = entry['width']
+        height = entry['height']
+
+        for obj in objs:
+            contiguous_cid = self.json_id_to_contiguous[obj['category_id']]
+            if contiguous_cid >= self.num_class:
+                # not class of interest
+                continue
+            if max(obj['keypoints']) == 0:
+                continue
+            # convert from (x, y, w, h) to (xmin, ymin, xmax, ymax) and clip bound
+            xmin, ymin, xmax, ymax = bbox_clip_xyxy(bbox_xywh_to_xyxy(obj['bbox']), width, height)
+            # require non-zero box area
+            #if obj['area'] <= 0 or xmax <= xmin or ymax <= ymin:
+            if (xmax-xmin)*(ymax-ymin) <= 0 or xmax <= xmin or ymax <= ymin:
+                continue
+            if 'num_keypoints' in obj and obj['num_keypoints'] == 0:
+                continue
+            # joints 3d: (num_joints, 3, 2); 3 is for x, y, z; 2 is for position, visibility
+            joints_3d = np.zeros((self.num_joints, 3, 2), dtype=np.float32)
+            for i in range(self.num_joints):
+                joints_3d[i, 0, 0] = obj['keypoints'][i * 3 + 0]
+                joints_3d[i, 1, 0] = obj['keypoints'][i * 3 + 1]
+                # joints_3d[i, 2, 0] = 0
+                if obj['keypoints'][i * 3 + 2] >= 0.35:
+                    visible = 1
+                else:
+                    visible = 0
+                #visible = min(1, visible)
+                joints_3d[i, :2, 1] = visible
+                # joints_3d[i, 2, 1] = 0
+
+            if np.sum(joints_3d[:, 0, 1]) < 1:
+                # no visible keypoint
+                continue
+
+            if self._check_centers and self._train:
+                bbox_center, bbox_area = self._get_box_center_area((xmin, ymin, xmax, ymax))
+                kp_center, num_vis = self._get_keypoints_center_count(joints_3d)
+                ks = np.exp(-2 * np.sum(np.square(bbox_center - kp_center)) / bbox_area)
+                if (num_vis / 80.0 + 47 / 80.0) > ks:
+                    continue
+
+            valid_objs.append({
+                'bbox': (xmin, ymin, xmax, ymax),
+                'width': width,
+                'height': height,
+                'joints_3d': joints_3d
+            })
+
+        if not valid_objs:
+            if not self._skip_empty:
+                # dummy invalid labels if no valid objects are found
+                valid_objs.append({
+                    'bbox': np.array([-1, -1, 0, 0]),
+                    'width': width,
+                    'height': height,
+                    'joints_3d': np.zeros((self.num_joints, 2, 2), dtype=np.float32)
+                })
+        return valid_objs
+
+    def _get_box_center_area(self, bbox):
+        """Get bbox center"""
+        c = np.array([(bbox[0] + bbox[2]) / 2.0, (bbox[1] + bbox[3]) / 2.0])
+        area = (bbox[3] - bbox[1]) * (bbox[2] - bbox[0])
+        return c, area
+
+    def _get_keypoints_center_count(self, keypoints):
+        """Get geometric center of all keypoints"""
+        keypoint_x = np.sum(keypoints[:, 0, 0] * (keypoints[:, 0, 1] > 0))
+        keypoint_y = np.sum(keypoints[:, 1, 0] * (keypoints[:, 1, 1] > 0))
+        num = float(np.sum(keypoints[:, 0, 1]))
+        return np.array([keypoint_x / num, keypoint_y / num]), num

alphapose/datasets1/halpe_26_det.py

@@ -0,0 +1,107 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by HaoyiZhu
+# -----------------------------------------------------
+
+"""Haple_26 Human Detection Box dataset."""
+import json
+import os
+
+import cv2
+import torch
+import torch.utils.data as data
+from tqdm import tqdm
+
+from alphapose.utils.presets import SimpleTransform
+from detector.apis import get_detector
+from alphapose.models.builder import DATASET
+
+
+@DATASET.register_module
+class Halpe_26_det(data.Dataset):
+    """ Halpe_26 human detection box dataset.
+
+    """
+    EVAL_JOINTS = list(range(26))
+
+    def __init__(self,
+                 det_file=None,
+                 opt=None,
+                 **cfg):
+
+        self._cfg = cfg
+        self._opt = opt
+        self._preset_cfg = cfg['PRESET']
+        self._root = cfg['ROOT']
+        self._img_prefix = cfg['IMG_PREFIX']
+        if not det_file:
+            det_file = cfg['DET_FILE']
+        self._ann_file = os.path.join(self._root, cfg['ANN'])
+
+        if os.path.exists(det_file):
+            print("Detection results exist, will use it")
+        else:
+            print("Will create detection results to {}".format(det_file))
+            self.write_coco_json(det_file)
+
+        assert os.path.exists(det_file), "Error: no detection results found"
+        with open(det_file, 'r') as fid:
+            self._det_json = json.load(fid)
+
+        self._input_size = self._preset_cfg['IMAGE_SIZE']
+        self._output_size = self._preset_cfg['HEATMAP_SIZE']
+
+        self._sigma = self._preset_cfg['SIGMA']
+
+        if self._preset_cfg['TYPE'] == 'simple':
+            self.transformation = SimpleTransform(
+                self, scale_factor=0,
+                input_size=self._input_size,
+                output_size=self._output_size,
+                rot=0, sigma=self._sigma,
+                train=False, add_dpg=False)
+
+    def __getitem__(self, index):
+        det_res = self._det_json[index]
+        if not isinstance(det_res['image_id'], int):
+            img_id, _ = os.path.splitext(os.path.basename(det_res['image_id']))
+            img_id = int(img_id)
+        else:
+            img_id = det_res['image_id']
+        img_path = '/DATA1/Benchmark/coco/val2017/%012d.jpg' % img_id
+
+        # Load image
+        image = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB) #scipy.misc.imread(img_path, mode='RGB')
+
+        imght, imgwidth = image.shape[1], image.shape[2]
+        x1, y1, w, h = det_res['bbox']
+        bbox = [x1, y1, x1 + w, y1 + h]
+        inp, bbox = self.transformation.test_transform(image, bbox)
+        return inp, torch.Tensor(bbox), torch.Tensor([det_res['bbox']]), torch.Tensor([det_res['image_id']]), torch.Tensor([det_res['score']]), torch.Tensor([imght]), torch.Tensor([imgwidth])
+
+    def __len__(self):
+        return len(self._det_json)
+
+    def write_coco_json(self, det_file):
+        from pycocotools.coco import COCO
+        import pathlib
+
+        _coco = COCO(self._ann_file)
+        image_ids = sorted(_coco.getImgIds())
+        det_model = get_detector(self._opt)
+        dets = []
+        for entry in tqdm(_coco.loadImgs(image_ids)):
+            abs_path = os.path.join(
+                '/DATA1/Benchmark/coco', self._img_prefix, entry['file_name'])
+            det = det_model.detect_one_img(abs_path)
+            if det:
+                dets += det
+        pathlib.Path(os.path.split(det_file)[0]).mkdir(parents=True, exist_ok=True)
+        json.dump(dets, open(det_file, 'w'))
+
+    @property
+    def joint_pairs(self):
+        """Joint pairs which defines the pairs of joint to be swapped
+        when the image is flipped horizontally."""
+        return[[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12], [13, 14], [15, 16], 
+        [20, 21], [22, 23], [24, 25]]

alphapose/datasets1/layers/DUC.py

@@ -0,0 +1,29 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+
+import torch.nn as nn
+
+
+class DUC(nn.Module):
+    '''
+    Initialize: inplanes, planes, upscale_factor
+    OUTPUT: (planes // upscale_factor^2) * ht * wd
+    '''
+
+    def __init__(self, inplanes, planes,
+                 upscale_factor=2, norm_layer=nn.BatchNorm2d):
+        super(DUC, self).__init__()
+        self.conv = nn.Conv2d(
+            inplanes, planes, kernel_size=3, padding=1, bias=False)
+        self.bn = norm_layer(planes, momentum=0.1)
+        self.relu = nn.ReLU(inplace=True)
+        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        x = self.pixel_shuffle(x)
+        return x

alphapose/datasets1/layers/PixelUnshuffle.py

@@ -0,0 +1,28 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+
+import torch.nn as nn
+
+
+class PixelUnshuffle(nn.Module):
+    '''
+    Initialize: inplanes, planes, upscale_factor
+    OUTPUT: (planes // upscale_factor^2) * ht * wd
+    '''
+
+    def __init__(self, downscale_factor=2):
+        super(PixelUnshuffle, self).__init__()
+        self._r = downscale_factor
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        out_c = c * (self._r * self._r)
+        out_h = h // self._r
+        out_w = w // self._r
+
+        x_view = x.contiguous().view(b, c, out_h, self._r, out_w, self._r)
+        x_prime = x_view.permute(0, 1, 3, 5, 2, 4).contiguous().view(b, out_c, out_h, out_w)
+
+        return x_prime

alphapose/datasets1/layers/Resnet.py

@@ -0,0 +1,200 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1,
+                 downsample=None, norm_layer=nn.BatchNorm2d, dcn=None):
+        super(Bottleneck, self).__init__()
+        self.dcn = dcn
+        self.with_dcn = dcn is not None
+
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = norm_layer(planes, momentum=0.1)
+        if self.with_dcn:
+            fallback_on_stride = dcn.get('FALLBACK_ON_STRIDE', False)
+            self.with_modulated_dcn = dcn.get('MODULATED', False)
+        if not self.with_dcn or fallback_on_stride:
+            self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                                   padding=1, bias=False)
+        else:
+            from .dcn import DeformConv, ModulatedDeformConv
+            self.deformable_groups = dcn.get('DEFORM_GROUP', 1)
+            if not self.with_modulated_dcn:
+                conv_op = DeformConv
+                offset_channels = 18
+            else:
+                conv_op = ModulatedDeformConv
+                offset_channels = 27
+
+            self.conv2_offset = nn.Conv2d(
+                planes,
+                self.deformable_groups * offset_channels,
+                kernel_size=3,
+                stride=stride,
+                padding=1)
+            self.conv2 = conv_op(
+                planes,
+                planes,
+                kernel_size=3,
+                stride=stride,
+                padding=1,
+                deformable_groups=self.deformable_groups,
+                bias=False)
+
+        self.bn2 = norm_layer(planes, momentum=0.1)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = norm_layer(planes * 4, momentum=0.1)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = F.relu(self.bn1(self.conv1(x)), inplace=True)
+        if not self.with_dcn:
+            out = F.relu(self.bn2(self.conv2(out)), inplace=True)
+        elif self.with_modulated_dcn:
+            offset_mask = self.conv2_offset(out)
+            offset = offset_mask[:, :18 * self.deformable_groups, :, :]
+            mask = offset_mask[:, -9 * self.deformable_groups:, :, :]
+            mask = mask.sigmoid()
+            out = F.relu(self.bn2(self.conv2(out, offset, mask)))
+        else:
+            offset = self.conv2_offset(out)
+            out = F.relu(self.bn2(self.conv2(out, offset)), inplace=True)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = F.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+    """ ResNet """
+
+    def __init__(self, architecture, norm_layer=nn.BatchNorm2d, dcn=None, stage_with_dcn=(False, False, False, False)):
+        super(ResNet, self).__init__()
+        self._norm_layer = norm_layer
+        assert architecture in ["resnet18", "resnet50", "resnet101", 'resnet152']
+        layers = {
+            'resnet18': [2, 2, 2, 2],
+            'resnet34': [3, 4, 6, 3],
+            'resnet50': [3, 4, 6, 3],
+            'resnet101': [3, 4, 23, 3],
+            'resnet152': [3, 8, 36, 3],
+        }
+        self.inplanes = 64
+        if architecture == "resnet18" or architecture == 'resnet34':
+            self.block = BasicBlock
+        else:
+            self.block = Bottleneck
+        self.layers = layers[architecture]
+
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7,
+                               stride=2, padding=3, bias=False)
+        self.bn1 = norm_layer(64, eps=1e-5, momentum=0.1, affine=True)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        stage_dcn = [dcn if with_dcn else None for with_dcn in stage_with_dcn]
+
+        self.layer1 = self.make_layer(
+            self.block, 64, self.layers[0], dcn=stage_dcn[0])
+        self.layer2 = self.make_layer(
+            self.block, 128, self.layers[1], stride=2, dcn=stage_dcn[1])
+        self.layer3 = self.make_layer(
+            self.block, 256, self.layers[2], stride=2, dcn=stage_dcn[2])
+
+        self.layer4 = self.make_layer(
+            self.block, 512, self.layers[3], stride=2, dcn=stage_dcn[3])
+
+    def forward(self, x):
+        x = self.maxpool(self.relu(self.bn1(self.conv1(x))))  # 64 * h/4 * w/4
+        x = self.layer1(x)  # 256 * h/4 * w/4
+        x = self.layer2(x)  # 512 * h/8 * w/8
+        x = self.layer3(x)  # 1024 * h/16 * w/16
+        x = self.layer4(x)  # 2048 * h/32 * w/32
+        return x
+
+    def stages(self):
+        return [self.layer1, self.layer2, self.layer3, self.layer4]
+
+    def make_layer(self, block, planes, blocks, stride=1, dcn=None):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                self._norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample,
+                            norm_layer=self._norm_layer, dcn=dcn))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes,
+                                norm_layer=self._norm_layer, dcn=dcn))
+
+        return nn.Sequential(*layers)

alphapose/datasets1/layers/SE_Resnet.py

@@ -0,0 +1,215 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .SE_module import SELayer
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None,
+                 reduction=False, norm_layer=nn.BatchNorm2d):
+        super(BasicBlock, self).__init__()
+
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+        if reduction:
+            self.se = SELayer(planes)
+        self.reduc = reduction
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.reduc:
+            out = self.se(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1,
+                 downsample=None, reduction=False,
+                 norm_layer=nn.BatchNorm2d,
+                 dcn=None):
+        super(Bottleneck, self).__init__()
+        self.dcn = dcn
+        self.with_dcn = dcn is not None
+
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = norm_layer(planes, momentum=0.1)
+        if self.with_dcn:
+            fallback_on_stride = dcn.get('FALLBACK_ON_STRIDE', False)
+            self.with_modulated_dcn = dcn.get('MODULATED', False)
+        if not self.with_dcn or fallback_on_stride:
+            self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                                   padding=1, bias=False)
+        else:
+            from .dcn import DeformConv, ModulatedDeformConv
+            self.deformable_groups = dcn.get('DEFORM_GROUP', 1)
+            if not self.with_modulated_dcn:
+                conv_op = DeformConv
+                offset_channels = 18
+            else:
+                conv_op = ModulatedDeformConv
+                offset_channels = 27
+
+            self.conv2_offset = nn.Conv2d(
+                planes,
+                self.deformable_groups * offset_channels,
+                kernel_size=3,
+                stride=stride,
+                padding=1)
+            self.conv2 = conv_op(
+                planes,
+                planes,
+                kernel_size=3,
+                stride=stride,
+                padding=1,
+                deformable_groups=self.deformable_groups,
+                bias=False)
+
+        self.bn2 = norm_layer(planes, momentum=0.1)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = norm_layer(planes * 4, momentum=0.1)
+        if reduction:
+            self.se = SELayer(planes * 4)
+
+        self.reduc = reduction
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = F.relu(self.bn1(self.conv1(x)), inplace=True)
+        if not self.with_dcn:
+            out = F.relu(self.bn2(self.conv2(out)), inplace=True)
+        elif self.with_modulated_dcn:
+            offset_mask = self.conv2_offset(out)
+            offset = offset_mask[:, :18 * self.deformable_groups, :, :]
+            mask = offset_mask[:, -9 * self.deformable_groups:, :, :]
+            mask = mask.sigmoid()
+            out = F.relu(self.bn2(self.conv2(out, offset, mask)))
+        else:
+            offset = self.conv2_offset(out)
+            out = F.relu(self.bn2(self.conv2(out, offset)), inplace=True)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+        if self.reduc:
+            out = self.se(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = F.relu(out)
+
+        return out
+
+
+class SEResnet(nn.Module):
+    """ SEResnet """
+
+    def __init__(self, architecture, norm_layer=nn.BatchNorm2d,
+                 dcn=None, stage_with_dcn=(False, False, False, False)):
+        super(SEResnet, self).__init__()
+        self._norm_layer = norm_layer
+        assert architecture in ["resnet18", "resnet50", "resnet101", 'resnet152']
+        layers = {
+            'resnet18': [2, 2, 2, 2],
+            'resnet34': [3, 4, 6, 3],
+            'resnet50': [3, 4, 6, 3],
+            'resnet101': [3, 4, 23, 3],
+            'resnet152': [3, 8, 36, 3],
+        }
+        self.inplanes = 64
+        if architecture == "resnet18" or architecture == 'resnet34':
+            self.block = BasicBlock
+        else:
+            self.block = Bottleneck
+        self.layers = layers[architecture]
+
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7,
+                               stride=2, padding=3, bias=False)
+        self.bn1 = norm_layer(64, eps=1e-5, momentum=0.1, affine=True)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        stage_dcn = [dcn if with_dcn else None for with_dcn in stage_with_dcn]
+
+        self.layer1 = self.make_layer(self.block, 64, self.layers[0], dcn=stage_dcn[0])
+        self.layer2 = self.make_layer(
+            self.block, 128, self.layers[1], stride=2, dcn=stage_dcn[1])
+        self.layer3 = self.make_layer(
+            self.block, 256, self.layers[2], stride=2, dcn=stage_dcn[2])
+
+        self.layer4 = self.make_layer(
+            self.block, 512, self.layers[3], stride=2, dcn=stage_dcn[3])
+
+    def forward(self, x):
+        x = self.maxpool(self.relu(self.bn1(self.conv1(x))))  # 64 * h/4 * w/4
+        x = self.layer1(x)  # 256 * h/4 * w/4
+        x = self.layer2(x)  # 512 * h/8 * w/8
+        x = self.layer3(x)  # 1024 * h/16 * w/16
+        x = self.layer4(x)  # 2048 * h/32 * w/32
+        return x
+
+    def stages(self):
+        return [self.layer1, self.layer2, self.layer3, self.layer4]
+
+    def make_layer(self, block, planes, blocks, stride=1, dcn=None):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                self._norm_layer(planes * block.expansion, momentum=0.1),
+            )
+
+        layers = []
+        if downsample is not None:
+            layers.append(block(self.inplanes, planes,
+                                stride, downsample, reduction=True,
+                                norm_layer=self._norm_layer, dcn=dcn))
+        else:
+            layers.append(block(self.inplanes, planes, stride, downsample,
+                                norm_layer=self._norm_layer, dcn=dcn))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes,
+                                norm_layer=self._norm_layer, dcn=dcn))
+
+        return nn.Sequential(*layers)

alphapose/datasets1/layers/SE_module.py

@@ -0,0 +1,24 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+
+from torch import nn
+
+
+class SELayer(nn.Module):
+    def __init__(self, channel, reduction=1):
+        super(SELayer, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Linear(channel, channel // reduction),
+            nn.ReLU(inplace=True),
+            nn.Linear(channel // reduction, channel),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        b, c, _, _ = x.size()
+        y = self.avg_pool(x).view(b, c)
+        y = self.fc(y).view(b, c, 1, 1)
+        return x * y

alphapose/datasets1/layers/ShuffleResnet.py

@@ -0,0 +1,197 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .dcn import DCN
+from .PixelUnshuffle import PixelUnshuffle
+from .SE_module import SELayer
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None,
+                 reduction=False, norm_layer=nn.BatchNorm2d):
+        super(BasicBlock, self).__init__()
+
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+        if reduction:
+            self.se = SELayer(planes)
+        self.reduc = reduction
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.reduc:
+            out = self.se(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1,
+                 downsample=None, reduction=False,
+                 norm_layer=nn.BatchNorm2d, dcn=None):
+        super(Bottleneck, self).__init__()
+        self.dcn = dcn
+        self.with_dcn = dcn is not None
+
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = norm_layer(planes, momentum=0.1)
+        if self.with_dcn:
+            fallback_on_stride = dcn.get('FALLBACK_ON_STRIDE', False)
+            self.with_modulated_dcn = dcn.get('MODULATED', False)
+
+        if stride > 1:
+            conv_layers = []
+            conv_layers.append(PixelUnshuffle(stride))
+            if not self.with_dcn or fallback_on_stride:
+                conv_layers.append(nn.Conv2d(planes * 4, planes, kernel_size=3, stride=1,
+                                             padding=1, bias=False))
+            else:
+                conv_layers.append(DCN(planes * 4, planes, dcn, kernel_size=3, stride=1,
+                                       padding=1, bias=False))
+            self.conv2 = nn.Sequential(*conv_layers)
+        else:
+            if not self.with_dcn or fallback_on_stride:
+                self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                                       padding=1, bias=False)
+            else:
+                self.conv2 = DCN(planes, planes, dcn, kernel_size=3, stride=stride,
+                                 padding=1, bias=False)
+
+        self.bn2 = norm_layer(planes, momentum=0.1)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = norm_layer(planes * 4, momentum=0.1)
+        if reduction:
+            self.se = SELayer(planes * 4)
+
+        self.reduc = reduction
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = F.relu(self.bn1(self.conv1(x)), inplace=True)
+        out = F.relu(self.bn2(self.conv2(out)), inplace=True)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+        if self.reduc:
+            out = self.se(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = F.relu(out, inplace=True)
+
+        return out
+
+
+class ShuffleResnet(nn.Module):
+    """ ShuffleResnet """
+
+    def __init__(self, architecture, norm_layer=nn.BatchNorm2d, dcn=None, stage_with_dcn=(False, False, False, False)):
+        super(ShuffleResnet, self).__init__()
+        self._norm_layer = norm_layer
+        assert architecture in ["resnet18", "resnet50", "resnet101", 'resnet152']
+        layers = {
+            'resnet18': [2, 2, 2, 2],
+            'resnet34': [3, 4, 6, 3],
+            'resnet50': [3, 4, 6, 3],
+            'resnet101': [3, 4, 23, 3],
+            'resnet152': [3, 8, 36, 3],
+        }
+        self.inplanes = 64
+        if architecture == "resnet18" or architecture == 'resnet34':
+            self.block = BasicBlock
+        else:
+            self.block = Bottleneck
+        self.layers = layers[architecture]
+
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7,
+                               stride=2, padding=3, bias=False)
+        self.bn1 = norm_layer(64, eps=1e-5, momentum=0.1, affine=True)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        stage_dcn = [dcn if with_dcn else None for with_dcn in stage_with_dcn]
+
+        self.layer1 = self.make_layer(
+            self.block, 64, self.layers[0], dcn=stage_dcn[0])
+        self.layer2 = self.make_layer(
+            self.block, 128, self.layers[1], stride=2, dcn=stage_dcn[1])
+        self.layer3 = self.make_layer(
+            self.block, 256, self.layers[2], stride=2, dcn=stage_dcn[2])
+
+        self.layer4 = self.make_layer(
+            self.block, 512, self.layers[3], stride=2, dcn=stage_dcn[3])
+
+    def forward(self, x):
+        x = self.maxpool(self.relu(self.bn1(self.conv1(x))))  # 64 * h/4 * w/4
+        x = self.layer1(x)  # 256 * h/4 * w/4
+        x = self.layer2(x)  # 512 * h/8 * w/8
+        x = self.layer3(x)  # 1024 * h/16 * w/16
+        x = self.layer4(x)  # 2048 * h/32 * w/32
+        return x
+
+    def stages(self):
+        return [self.layer1, self.layer2, self.layer3, self.layer4]
+
+    def make_layer(self, block, planes, blocks, stride=1, dcn=None):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                self._norm_layer(planes * block.expansion, momentum=0.1),
+            )
+
+        layers = []
+        if downsample is not None:
+            layers.append(block(self.inplanes, planes,
+                                stride, downsample, reduction=True,
+                                norm_layer=self._norm_layer, dcn=dcn))
+        else:
+            layers.append(block(self.inplanes, planes, stride, downsample,
+                                norm_layer=self._norm_layer, dcn=dcn))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes,
+                                norm_layer=self._norm_layer, dcn=dcn))
+
+        return nn.Sequential(*layers)

alphapose/datasets1/mpii.py

@@ -0,0 +1,124 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+
+"""MPII Human Pose Dataset."""
+import os
+
+import numpy as np
+
+from alphapose.models.builder import DATASET
+from alphapose.utils.bbox import bbox_clip_xyxy, bbox_xywh_to_xyxy
+
+from .custom import CustomDataset
+
+
+@DATASET.register_module
+class Mpii(CustomDataset):
+    """ MPII Human Pose Dataset.
+
+    Parameters
+    ----------
+    root: str, default './data/mpii'
+        Path to the mpii dataset.
+    train: bool, default is True
+        If true, will set as training mode.
+    dpg: bool, default is False
+        If true, will activate `dpg` for data augmentation.
+    """
+
+    CLASSES = ['person']
+    num_joints = 16
+
+    @property
+    def joint_pairs(self):
+        """Joint pairs which defines the pairs of joint to be swapped
+        when the image is flipped horizontally."""
+        return [[0, 5], [1, 4], [2, 3],
+                [10, 15], [11, 14], [12, 13]]
+
+    def _load_jsons(self):
+        """Load all image paths and labels from annotation files into buffer."""
+        items = []
+        labels = []
+
+        _mpii = self._lazy_load_ann_file()
+        classes = [c['name'] for c in _mpii.loadCats(_mpii.getCatIds())]
+        assert classes == self.CLASSES, "Incompatible category names with MPII. "
+
+        # iterate through the annotations
+        image_ids = sorted(_mpii.getImgIds())
+        for entry in _mpii.loadImgs(image_ids):
+            filename = entry['file_name']
+            abs_path = os.path.join(self._root, self._img_prefix, filename)
+            if not os.path.exists(abs_path):
+                raise IOError('Image: {} not exists.'.format(abs_path))
+            label = self._check_load_keypoints(_mpii, entry)
+            if not label:
+                continue
+
+            # num of items are relative to person, not image
+            for obj in label:
+                items.append(abs_path)
+                labels.append(obj)
+
+        return items, labels
+
+    def _check_load_keypoints(self, _mpii, entry):
+        """Check and load ground-truth keypoints"""
+        ann_ids = _mpii.getAnnIds(imgIds=entry['id'], iscrowd=False)
+        objs = _mpii.loadAnns(ann_ids)
+        # check valid bboxes
+        valid_objs = []
+        width = entry['width']
+        height = entry['height']
+
+        for obj in objs:
+            if max(obj['keypoints']) == 0:
+                continue
+            # convert from (x, y, w, h) to (xmin, ymin, xmax, ymax) and clip bound
+            xmin, ymin, xmax, ymax = bbox_clip_xyxy(bbox_xywh_to_xyxy(obj['bbox']), width, height)
+            # require non-zero box area
+            if xmax <= xmin or ymax <= ymin:
+                continue
+            if obj['num_keypoints'] == 0:
+                continue
+            # joints 3d: (num_joints, 3, 2); 3 is for x, y, z; 2 is for position, visibility
+            joints_3d = np.zeros((self.num_joints, 3, 2), dtype=np.float32)
+            for i in range(self.num_joints):
+                joints_3d[i, 0, 0] = obj['keypoints'][i * 3 + 0]
+                joints_3d[i, 1, 0] = obj['keypoints'][i * 3 + 1]
+                # joints_3d[i, 2, 0] = 0
+                visible = min(1, obj['keypoints'][i * 3 + 2])
+                joints_3d[i, :2, 1] = visible
+                # joints_3d[i, 2, 1] = 0
+
+            if np.sum(joints_3d[:, 0, 1]) < 1:
+                # no visible keypoint
+                continue
+
+            if self._check_centers and self._train:
+                bbox_center, bbox_area = self._get_box_center_area((xmin, ymin, xmax, ymax))
+                kp_center, num_vis = self._get_keypoints_center_count(joints_3d)
+                ks = np.exp(-2 * np.sum(np.square(bbox_center - kp_center)) / bbox_area)
+                if (num_vis / 80.0 + 47 / 80.0) > ks:
+                    continue
+
+            valid_objs.append({
+                'bbox': (xmin, ymin, xmax, ymax),
+                'width': width,
+                'height': height,
+                'joints_3d': joints_3d
+            })
+
+        if not valid_objs:
+            if not self._skip_empty:
+                # dummy invalid labels if no valid objects are found
+                valid_objs.append({
+                    'bbox': np.array([-1, -1, 0, 0]),
+                    'width': width,
+                    'height': height,
+                    'joints_3d': np.zeros((self.num_joints, 2, 2), dtype=np.float32)
+                })
+        return valid_objs

alphapose/datasets1/mscoco.py

@@ -0,0 +1,140 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+
+"""MS COCO Human keypoint dataset."""
+import os
+
+import numpy as np
+
+from alphapose.models.builder import DATASET
+from alphapose.utils.bbox import bbox_clip_xyxy, bbox_xywh_to_xyxy
+
+from .custom import CustomDataset
+
+
+@DATASET.register_module
+class Mscoco(CustomDataset):
+    """ COCO Person dataset.
+
+    Parameters
+    ----------
+    train: bool, default is True
+        If true, will set as training mode.
+    skip_empty: bool, default is False
+        Whether skip entire image if no valid label is found. Use `False` if this dataset is
+        for validation to avoid COCO metric error.
+    dpg: bool, default is False
+        If true, will activate `dpg` for data augmentation.
+    """
+    CLASSES = ['person']
+    EVAL_JOINTS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
+    num_joints = 17
+    joint_pairs = [[1, 2], [3, 4], [5, 6], [7, 8],
+                   [9, 10], [11, 12], [13, 14], [15, 16]]
+
+    def _load_jsons(self):
+        """Load all image paths and labels from JSON annotation files into buffer."""
+        items = []
+        labels = []
+
+        _coco = self._lazy_load_ann_file()
+        classes = [c['name'] for c in _coco.loadCats(_coco.getCatIds())]
+        assert classes == self.CLASSES, "Incompatible category names with COCO. "
+
+        self.json_id_to_contiguous = {
+            v: k for k, v in enumerate(_coco.getCatIds())}
+
+        # iterate through the annotations
+        image_ids = sorted(_coco.getImgIds())
+        for entry in _coco.loadImgs(image_ids):
+            dirname, filename = entry['coco_url'].split('/')[-2:]
+            abs_path = os.path.join(self._root, dirname, filename)
+            if not os.path.exists(abs_path):
+                raise IOError('Image: {} not exists.'.format(abs_path))
+            label = self._check_load_keypoints(_coco, entry)
+            if not label:
+                continue
+
+            # num of items are relative to person, not image
+            for obj in label:
+                items.append(abs_path)
+                labels.append(obj)
+
+        return items, labels
+
+    def _check_load_keypoints(self, coco, entry):
+        """Check and load ground-truth keypoints"""
+        ann_ids = coco.getAnnIds(imgIds=entry['id'], iscrowd=False)
+        objs = coco.loadAnns(ann_ids)
+        # check valid bboxes
+        valid_objs = []
+        width = entry['width']
+        height = entry['height']
+
+        for obj in objs:
+            contiguous_cid = self.json_id_to_contiguous[obj['category_id']]
+            if contiguous_cid >= self.num_class:
+                # not class of interest
+                continue
+            if max(obj['keypoints']) == 0:
+                continue
+            # convert from (x, y, w, h) to (xmin, ymin, xmax, ymax) and clip bound
+            xmin, ymin, xmax, ymax = bbox_clip_xyxy(bbox_xywh_to_xyxy(obj['bbox']), width, height)
+            # require non-zero box area
+            if obj['area'] <= 0 or xmax <= xmin or ymax <= ymin:
+                continue
+            if obj['num_keypoints'] == 0:
+                continue
+            # joints 3d: (num_joints, 3, 2); 3 is for x, y, z; 2 is for position, visibility
+            joints_3d = np.zeros((self.num_joints, 3, 2), dtype=np.float32)
+            for i in range(self.num_joints):
+                joints_3d[i, 0, 0] = obj['keypoints'][i * 3 + 0]
+                joints_3d[i, 1, 0] = obj['keypoints'][i * 3 + 1]
+                # joints_3d[i, 2, 0] = 0
+                visible = min(1, obj['keypoints'][i * 3 + 2])
+                joints_3d[i, :2, 1] = visible
+                # joints_3d[i, 2, 1] = 0
+
+            if np.sum(joints_3d[:, 0, 1]) < 1:
+                # no visible keypoint
+                continue
+
+            if self._check_centers and self._train:
+                bbox_center, bbox_area = self._get_box_center_area((xmin, ymin, xmax, ymax))
+                kp_center, num_vis = self._get_keypoints_center_count(joints_3d)
+                ks = np.exp(-2 * np.sum(np.square(bbox_center - kp_center)) / bbox_area)
+                if (num_vis / 80.0 + 47 / 80.0) > ks:
+                    continue
+
+            valid_objs.append({
+                'bbox': (xmin, ymin, xmax, ymax),
+                'width': width,
+                'height': height,
+                'joints_3d': joints_3d
+            })
+
+        if not valid_objs:
+            if not self._skip_empty:
+                # dummy invalid labels if no valid objects are found
+                valid_objs.append({
+                    'bbox': np.array([-1, -1, 0, 0]),
+                    'width': width,
+                    'height': height,
+                    'joints_3d': np.zeros((self.num_joints, 2, 2), dtype=np.float32)
+                })
+        return valid_objs
+
+    def _get_box_center_area(self, bbox):
+        """Get bbox center"""
+        c = np.array([(bbox[0] + bbox[2]) / 2.0, (bbox[1] + bbox[3]) / 2.0])
+        area = (bbox[3] - bbox[1]) * (bbox[2] - bbox[0])
+        return c, area
+
+    def _get_keypoints_center_count(self, keypoints):
+        """Get geometric center of all keypoints"""
+        keypoint_x = np.sum(keypoints[:, 0, 0] * (keypoints[:, 0, 1] > 0))
+        keypoint_y = np.sum(keypoints[:, 1, 0] * (keypoints[:, 1, 1] > 0))
+        num = float(np.sum(keypoints[:, 0, 1]))
+        return np.array([keypoint_x / num, keypoint_y / num]), num

alphapose/models/__init__.py

@@ -0,0 +1,10 @@
+from .fastpose import FastPose
+from .fastpose_duc import FastPose_DUC
+from .hrnet import PoseHighResolutionNet
+from .simplepose import SimplePose
+from .fastpose_duc_dense import FastPose_DUC_Dense
+from .hardnet import HarDNetPose
+from .criterion import L1JointRegression
+
+__all__ = ['FastPose', 'SimplePose', 'PoseHighResolutionNet',
+           'FastPose_DUC', 'FastPose_DUC_Dense', 'HarDNetPose', 'L1JointRegression']

alphapose/models/builder.py

@@ -0,0 +1,46 @@
+from torch import nn
+
+from alphapose.utils import Registry, build_from_cfg, retrieve_from_cfg
+
+
+SPPE = Registry('sppe')
+LOSS = Registry('loss')
+DATASET = Registry('dataset')
+
+
+def build(cfg, registry, default_args=None):
+    if isinstance(cfg, list):
+        modules = [
+            build_from_cfg(cfg_, registry, default_args) for cfg_ in cfg
+        ]
+        return nn.Sequential(*modules)
+    else:
+        return build_from_cfg(cfg, registry, default_args)
+
+
+def build_sppe(cfg, preset_cfg, **kwargs):
+    default_args = {
+        'PRESET': preset_cfg,
+    }
+    for key, value in kwargs.items():
+        default_args[key] = value
+    return build(cfg, SPPE, default_args=default_args)
+
+
+def build_loss(cfg):
+    return build(cfg, LOSS)
+
+
+def build_dataset(cfg, preset_cfg, **kwargs):
+    exec(f'from ..datasets import {cfg.TYPE}')
+    default_args = {
+        'PRESET': preset_cfg,
+    }
+    for key, value in kwargs.items():
+        default_args[key] = value
+    return build(cfg, DATASET, default_args=default_args)
+
+
+def retrieve_dataset(cfg):
+    exec(f'from ..datasets import {cfg.TYPE}')
+    return retrieve_from_cfg(cfg, DATASET)

alphapose/models/criterion.py

@@ -0,0 +1,97 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+import torch
+import torch.nn as nn
+
+from .builder import LOSS
+
+from alphapose.utils.transforms import _integral_tensor
+
+
+class IngetralCoordinate(torch.autograd.Function):
+    ''' Symmetry integral regression function.
+    '''
+    AMPLITUDE = 2
+
+    @staticmethod
+    def forward(ctx, input):
+        assert isinstance(
+            input, torch.Tensor), 'IngetralCoordinate only takes input as torch.Tensor'
+        input_size = input.size()
+        weight = torch.arange(
+            input_size[-1], dtype=input.dtype, layout=input.layout, device=input.device)
+        ctx.input_size = input_size
+        output = input.mul(weight)
+        ctx.save_for_backward(weight, output)
+
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        weight, output = ctx.saved_tensors
+        output_coord = output.sum(dim=2, keepdim=True)
+        weight = weight[None, None, :].repeat(
+            output_coord.shape[0], output_coord.shape[1], 1)
+        weight_mask = torch.ones(weight.shape, dtype=grad_output.dtype,
+                                 layout=grad_output.layout, device=grad_output.device)
+        weight_mask[weight < output_coord] = -1
+        weight_mask[output_coord.repeat(
+            1, 1, weight.shape[-1]) > ctx.input_size[-1]] = 1
+        weight_mask *= IngetralCoordinate.AMPLITUDE
+        return grad_output.mul(weight_mask)
+
+
+@LOSS.register_module
+class L1JointRegression(nn.Module):
+    ''' L1 Joint Regression Loss
+    '''
+    def __init__(self, OUTPUT_3D=False, size_average=True, reduce=True, NORM_TYPE='softmax'):
+        super(L1JointRegression, self).__init__()
+        self.size_average = size_average
+        self.reduce = reduce
+        self.output_3d = OUTPUT_3D
+        self.norm_type = NORM_TYPE
+
+        self.integral_operation = IngetralCoordinate.apply
+
+    def forward(self, preds, *args):
+        gt_joints = args[0]
+        gt_joints_vis = args[1]
+
+        if self.output_3d:
+            num_joints = int(gt_joints_vis.shape[1] / 3)
+        else:
+            num_joints = int(gt_joints_vis.shape[1] / 2)
+        hm_width = preds.shape[-1]
+        hm_height = preds.shape[-2]
+        hm_depth = preds.shape[-3] // num_joints if self.output_3d else 1
+
+        pred_jts, pred_scores = _integral_tensor(
+            preds, num_joints, self.output_3d, hm_width, hm_height, hm_depth, integral_operation=self.integral_operation, norm_type=self.norm_type)
+
+        _assert_no_grad(gt_joints)
+        _assert_no_grad(gt_joints_vis)
+        return weighted_l1_loss(pred_jts, pred_scores, gt_joints, gt_joints_vis, self.size_average)
+
+
+def _assert_no_grad(tensor):
+    assert not tensor.requires_grad, \
+        "nn criterions don't compute the gradient w.r.t. targets - please " \
+        "mark these tensors as not requiring gradients"
+
+
+def weighted_l1_loss(input, scores, target, weights, size_average):
+    out = torch.abs(input - target)
+    out = out * weights
+    #out_of_scores = torch.abs(scores - torch.ones_like(scores))
+    #out_of_scores = out_of_scores.reshape((out_of_scores.shape[0], -1))
+    #out_of_scores = out_of_scores * weights[:, 0::2]
+    if size_average:
+        return out.sum() / len(input)
+    else:
+        return out.sum()
+
+
+LOSS.register_module(torch.nn.MSELoss)

alphapose/models/fastpose.py

@@ -0,0 +1,67 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+
+import torch.nn as nn
+
+from .builder import SPPE
+from .layers.DUC import DUC
+from .layers.SE_Resnet import SEResnet
+
+
+@SPPE.register_module
+class FastPose(nn.Module):
+
+    def __init__(self, norm_layer=nn.BatchNorm2d, **cfg):
+        super(FastPose, self).__init__()
+        self._preset_cfg = cfg['PRESET']
+        if 'CONV_DIM' in cfg.keys():
+            self.conv_dim = cfg['CONV_DIM']
+        else:
+            self.conv_dim = 128
+        if 'DCN' in cfg.keys():
+            stage_with_dcn = cfg['STAGE_WITH_DCN']
+            dcn = cfg['DCN']
+            self.preact = SEResnet(
+                f"resnet{cfg['NUM_LAYERS']}", dcn=dcn, stage_with_dcn=stage_with_dcn)
+        else:
+            self.preact = SEResnet(f"resnet{cfg['NUM_LAYERS']}")
+
+        # Imagenet pretrain model
+        import torchvision.models as tm   # noqa: F401,F403
+        assert cfg['NUM_LAYERS'] in [18, 34, 50, 101, 152]
+        x = eval(f"tm.resnet{cfg['NUM_LAYERS']}(pretrained=True)")
+
+        model_state = self.preact.state_dict()
+        state = {k: v for k, v in x.state_dict().items()
+                 if k in self.preact.state_dict() and v.size() == self.preact.state_dict()[k].size()}
+        model_state.update(state)
+        self.preact.load_state_dict(model_state)
+
+        self.suffle1 = nn.PixelShuffle(2)
+        self.duc1 = DUC(512, 1024, upscale_factor=2, norm_layer=norm_layer)
+        if self.conv_dim == 256:
+            self.duc2 = DUC(256, 1024, upscale_factor=2, norm_layer=norm_layer)
+        else:
+            self.duc2 = DUC(256, 512, upscale_factor=2, norm_layer=norm_layer)
+        self.conv_out = nn.Conv2d(
+            self.conv_dim, self._preset_cfg['NUM_JOINTS'], kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x):
+        out = self.preact(x)
+        out = self.suffle1(out)
+        out = self.duc1(out)
+        out = self.duc2(out)
+
+        out = self.conv_out(out)
+        return out
+
+    def _initialize(self):
+        for m in self.conv_out.modules():
+            if isinstance(m, nn.Conv2d):
+                # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                # logger.info('=> init {}.weight as normal(0, 0.001)'.format(name))
+                # logger.info('=> init {}.bias as 0'.format(name))
+                nn.init.normal_(m.weight, std=0.001)
+                nn.init.constant_(m.bias, 0)

alphapose/models/fastpose_duc.py

@@ -0,0 +1,97 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+
+import torch.nn as nn
+
+from .builder import SPPE
+from .layers.Resnet import ResNet
+from .layers.SE_Resnet import SEResnet
+from .layers.ShuffleResnet import ShuffleResnet
+
+
+@SPPE.register_module
+class FastPose_DUC(nn.Module):
+    conv_dim = 256
+
+    def __init__(self, norm_layer=nn.BatchNorm2d, **cfg):
+        super(FastPose_DUC, self).__init__()
+        self._preset_cfg = cfg['PRESET']
+        if cfg['BACKBONE'] == 'shuffle':
+            print('Load shuffle backbone...')
+            backbone = ShuffleResnet
+        elif cfg['BACKBONE'] == 'se-resnet':
+            print('Load SE Resnet...')
+            backbone = SEResnet
+        else:
+            print('Load Resnet...')
+            backbone = ResNet
+
+        if 'DCN' in cfg.keys():
+            stage_with_dcn = cfg['STAGE_WITH_DCN']
+            dcn = cfg['DCN']
+            self.preact = backbone(
+                f"resnet{cfg['NUM_LAYERS']}", dcn=dcn, stage_with_dcn=stage_with_dcn)
+        else:
+            self.preact = backbone(f"resnet{cfg['NUM_LAYERS']}")
+
+        # Imagenet pretrain model
+        import torchvision.models as tm   # noqa: F401,F403
+        assert cfg['NUM_LAYERS'] in [18, 34, 50, 101, 152]
+        x = eval(f"tm.resnet{cfg['NUM_LAYERS']}(pretrained=True)")
+
+        model_state = self.preact.state_dict()
+        state = {k: v for k, v in x.state_dict().items()
+                 if k in self.preact.state_dict() and v.size() == self.preact.state_dict()[k].size()}
+        model_state.update(state)
+        self.preact.load_state_dict(model_state)
+        self.norm_layer = norm_layer
+
+        stage1_cfg = cfg['STAGE1']
+        stage2_cfg = cfg['STAGE2']
+        stage3_cfg = cfg['STAGE3']
+
+        self.duc1 = self._make_duc_stage(stage1_cfg, 2048, 1024)
+        self.duc2 = self._make_duc_stage(stage2_cfg, 1024, 512)
+        self.duc3 = self._make_duc_stage(stage3_cfg, 512, self.conv_dim)
+
+        self.conv_out = nn.Conv2d(
+            self.conv_dim, self._preset_cfg['NUM_JOINTS'], kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x):
+        out = self.preact(x)
+        out = self.duc1(out)
+        out = self.duc2(out)
+        out = self.duc3(out)
+
+        out = self.conv_out(out)
+        return out
+
+    def _make_duc_stage(self, layer_config, inplanes, outplanes):
+        layers = []
+
+        shuffle = nn.PixelShuffle(2)
+        inplanes //= 4
+        layers.append(shuffle)
+        for i in range(layer_config.NUM_CONV - 1):
+            conv = nn.Conv2d(inplanes, inplanes, kernel_size=3,
+                             padding=1, bias=False)
+            norm_layer = self.norm_layer(inplanes, momentum=0.1)
+            relu = nn.ReLU(inplace=True)
+            layers += [conv, norm_layer, relu]
+        conv = nn.Conv2d(inplanes, outplanes, kernel_size=3,
+                         padding=1, bias=False)
+        norm_layer = self.norm_layer(outplanes, momentum=0.1)
+        relu = nn.ReLU(inplace=True)
+        layers += [conv, norm_layer, relu]
+        return nn.Sequential(*layers)
+
+    def _initialize(self):
+        for m in self.conv_out.modules():
+            if isinstance(m, nn.Conv2d):
+                # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                # logger.info('=> init {}.weight as normal(0, 0.001)'.format(name))
+                # logger.info('=> init {}.bias as 0'.format(name))
+                nn.init.normal_(m.weight, std=0.001)
+                nn.init.constant_(m.bias, 0)

alphapose/models/fastpose_duc_dense.py

@@ -0,0 +1,155 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+
+import torch
+import torch.nn as nn
+from .builder import SPPE
+from .layers.Resnet import ResNet
+from .layers.SE_Resnet import SEResnet
+from .layers.ShuffleResnet import ShuffleResnet
+
+@SPPE.register_module
+class FastPose_DUC_Dense(nn.Module):
+    conv_dim = 256
+
+    def __init__(self,norm_layer=nn.BatchNorm2d,**cfg):
+        super(FastPose_DUC_Dense, self).__init__()
+        self._preset_cfg = cfg['PRESET']
+        if cfg['BACKBONE'] == 'shuffle':
+            print('Load shuffle backbone...')
+            backbone = ShuffleResnet
+        elif cfg['BACKBONE'] == 'se-resnet':
+            print('Load SE Resnet...')
+            backbone = SEResnet
+        else:
+            print('Load Resnet...')
+            backbone = ResNet
+
+        if 'DCN' in cfg.keys():
+            stage_with_dcn = cfg['STAGE_WITH_DCN']
+            dcn = cfg['DCN']
+            self.preact = backbone(
+                f"resnet{cfg['NUM_LAYERS']}", dcn=dcn, stage_with_dcn=stage_with_dcn)
+        else:
+            self.preact = backbone(f"resnet{cfg['NUM_LAYERS']}")
+
+        # Init Backbone
+        for m in self.preact.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.normal_(m.weight, std=0.001)
+                for name, _ in m.named_parameters():
+                    if name in ['bias']:
+                        nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                # nn.init.constant_(m.weight, 1)
+                nn.init.uniform_(m.weight, 0, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Imagenet pretrain model
+        import torchvision.models as tm
+        if cfg['NUM_LAYERS'] == 152:
+            ''' Load pretrained model '''
+            x = tm.resnet152(pretrained=True)
+        elif cfg['NUM_LAYERS'] == 101:
+            ''' Load pretrained model '''
+            x = tm.resnet101(pretrained=True)
+        elif cfg['NUM_LAYERS'] == 50:
+            x = tm.resnet50(pretrained=True)
+        elif cfg['NUM_LAYERS'] == 18:
+            x = tm.resnet18(pretrained=True)
+        else:
+            raise NotImplementedError
+        model_state = self.preact.state_dict()
+        state = {k: v for k, v in x.state_dict().items()
+                 if k in self.preact.state_dict() and v.size() == self.preact.state_dict()[k].size()}
+        model_state.update(state)
+        self.preact.load_state_dict(model_state)
+        self.norm_layer = norm_layer
+
+        stage1_cfg = cfg['STAGE1']
+        stage2_cfg = cfg['STAGE2']
+        stage3_cfg = cfg['STAGE3']
+
+        duc1 = self._make_duc_stage(stage1_cfg, 2048, 1024)
+        duc2 = self._make_duc_stage(stage2_cfg, 1024, 512)
+        duc3 = self._make_duc_stage(stage3_cfg, 512, self.conv_dim)
+
+        self.duc = nn.Sequential(duc1, duc2, duc3)
+
+        duc1_dense = self._make_duc_stage(stage1_cfg,2048,1024)
+        duc2_dense = self._make_duc_stage(stage2_cfg,1024,512)
+        duc3_dense = self._make_duc_stage(stage3_cfg,512,self.conv_dim)
+
+        self.duc_dense = nn.Sequential(duc1_dense,duc2_dense,duc3_dense)
+
+        self.conv_out = nn.Conv2d(
+            self.conv_dim, self._preset_cfg['NUM_JOINTS'], kernel_size=3, stride=1, padding=1)
+
+        self.conv_out_dense = nn.Conv2d(
+            self.conv_dim,(self._preset_cfg['NUM_JOINTS_DENSE']-self._preset_cfg['NUM_JOINTS']),kernel_size=3,stride=1,padding=1)
+        for params in self.preact.parameters():
+            params.requires_grad = False
+        for params in self.duc.parameters():
+            params.requires_grad = False
+
+    def forward(self, x):
+        bk_out = self.preact(x)
+        out = self.duc(bk_out)
+        out_dense = self.duc_dense(bk_out)
+        out = self.conv_out(out)
+        out_dense = self.conv_out_dense(out_dense)
+        out = torch.cat((out,out_dense),1)
+        return out
+
+    def _make_duc_stage(self, layer_config, inplanes, outplanes):
+        layers = []
+
+        shuffle = nn.PixelShuffle(2)
+        inplanes //= 4
+        layers.append(shuffle)
+        for i in range(layer_config.NUM_CONV - 1):
+            conv = nn.Conv2d(inplanes, inplanes, kernel_size=3,
+                             padding=1, bias=False)
+            norm_layer = self.norm_layer(inplanes, momentum=0.1)
+            relu = nn.ReLU(inplace=True)
+            layers += [conv, norm_layer, relu]
+        conv = nn.Conv2d(inplanes, outplanes, kernel_size=3,
+                         padding=1, bias=False)
+        norm_layer = self.norm_layer(outplanes, momentum=0.1)
+        relu = nn.ReLU(inplace=True)
+        layers += [conv, norm_layer, relu]
+        return nn.Sequential(*layers)
+
+    def _initialize(self):
+        for m in self.duc.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.normal_(m.weight, std=0.001)
+                for name, _ in m.named_parameters():
+                    if name in ['bias']:
+                        nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                # nn.init.constant_(m.weight, 1)
+                nn.init.uniform_(m.weight, 0, 1)
+                nn.init.constant_(m.bias, 0)
+        for m in self.conv_out.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.normal_(m.weight, std=0.001)
+                nn.init.constant_(m.bias, 0)
+
+        #init dense-branch
+        for m in self.duc_dense.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.normal_(m.weight, std=0.001)
+                for name, _ in m.named_parameters():
+                    if name in ['bias']:
+                        nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                # nn.init.constant_(m.weight, 1)
+                nn.init.uniform_(m.weight, 0, 1)
+                nn.init.constant_(m.bias, 0)
+        for m in self.conv_out_dense.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.normal_(m.weight, std=0.001)
+                nn.init.constant_(m.bias, 0)

alphapose/models/hardnet.py

@@ -0,0 +1,569 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import collections
+import numpy as np
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import torch.nn.init as init
+
+
+from .builder import SPPE
+from .layers.Resnet import ResNet
+from .layers.SE_Resnet import SEResnet
+from .layers.ShuffleResnet import ShuffleResnet
+
+BN_MOMENTUM = 0.1
+DEBUG = False
+
+def conv3x3(in_planes, out_planes, stride=1):
+    "3x3 convolution with padding"
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class Identity(nn.Module):
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, x):
+        return x
+
+
+class Flatten(nn.Module):
+    def __init__(self):
+        super().__init__()
+    def forward(self, x):
+        return x.view(x.data.size(0),-1)
+
+
+class CombConvLayer(nn.Sequential):
+    def __init__(self, in_channels, out_channels, norm_layer, kernel=1, stride=1, dropout=0.1, bias=False):
+        super().__init__()
+        self.add_module('layer1',ConvLayer(in_channels, out_channels, kernel))
+        self.add_module('layer2',DWConvLayer(out_channels, out_channels, norm_layer, stride=stride))
+
+    def forward(self, x):
+        return super().forward(x)
+
+
+class DWConvLayer(nn.Sequential):
+    def __init__(self, in_channels, out_channels,  norm_layer, stride=1,  bias=False):
+        super().__init__()
+        out_ch = out_channels
+
+        groups = in_channels
+        kernel = 3
+        if DEBUG:
+          print(kernel, 'x', kernel, 'x', out_channels, 'x', out_channels, 'DepthWise')
+
+        self.add_module('dwconv', nn.Conv2d(groups, groups, kernel_size=3,
+                                          stride=stride, padding=1, groups=groups, bias=bias))
+
+        self.add_module('norm', norm_layer(groups, momentum=BN_MOMENTUM))
+    def forward(self, x):
+        return super().forward(x)
+
+
+class ConvLayer(nn.Sequential):
+    def __init__(self, in_channels, out_channels, norm_layer, kernel=3, stride=1, padding=0, bias=False):
+        super().__init__()
+        self.out_channels = out_channels
+        out_ch = out_channels
+        groups = 1
+        if DEBUG:
+          print(kernel, 'x', kernel, 'x', in_channels, 'x', out_channels)
+        pad = kernel//2 if padding == 0 else padding
+        self.add_module('conv', nn.Conv2d(in_channels, out_ch, kernel_size=kernel,
+                                          stride=stride, padding=pad, groups=groups, bias=bias))
+        self.add_module('norm', norm_layer(out_ch, momentum=BN_MOMENTUM))
+        self.add_module('relu', nn.ReLU(True))
+    def forward(self, x):
+        return super().forward(x)
+
+
+class BRLayer(nn.Sequential):
+    def __init__(self, in_channels, norm_layer):
+        super().__init__()
+
+        self.add_module('norm', norm_layer(in_channels))
+        self.add_module('relu', nn.ReLU(True))
+    def forward(self, x):
+        return super().forward(x)
+
+
+class HarDBlock(nn.Module):
+    def get_link(self, layer, base_ch, growth_rate, grmul):
+        if layer == 0:
+          return base_ch, 0, []
+        out_channels = growth_rate
+        link = []
+        for i in range(10):
+          dv = 2 ** i
+          if layer % dv == 0:
+            k = layer - dv
+            link.append(k)
+            if i > 0:
+                out_channels *= grmul
+        out_channels = int(int(out_channels + 1) / 2) * 2
+        in_channels = 0
+        for i in link:
+          ch,_,_ = self.get_link(i, base_ch, growth_rate, grmul)
+          in_channels += ch
+        return out_channels, in_channels, link
+
+    def get_out_ch(self):
+        return self.out_channels
+
+    def __init__(self, in_channels, growth_rate, grmul, n_layers, norm_layer, keepBase=False, residual_out=False, dwconv=False):
+        super().__init__()
+        self.in_channels = in_channels
+        self.growth_rate = growth_rate
+        self.grmul = grmul
+        self.n_layers = n_layers
+        self.norm_layer = norm_layer
+        self.keepBase = keepBase
+        self.links = []
+        layers_ = []
+        self.out_channels = 0
+
+        for i in range(n_layers):
+          outch, inch, link = self.get_link(i+1, in_channels, growth_rate, grmul)
+          self.links.append(link)
+          use_relu = residual_out
+          if dwconv:
+            layers_.append(CombConvLayer(inch, outch, norm_layer))
+          else:
+            layers_.append(ConvLayer(inch, outch, norm_layer))
+
+          if (i % 2 == 0) or (i == n_layers - 1):
+            self.out_channels += outch
+        if DEBUG:
+          print("Blk out =",self.out_channels)
+        self.layers = nn.ModuleList(layers_)
+
+    def forward(self, x):
+        layers_ = [x]
+        for layer in range(len(self.layers)):
+            link = self.links[layer]
+            tin = []
+            for i in link:
+                tin.append(layers_[i])
+            if len(tin) > 1:
+                x = torch.cat(tin, 1)
+            else:
+                x = tin[0]
+            out = self.layers[layer](x)
+            layers_.append(out)
+        t = len(layers_)
+        out_ = []
+        for i in range(t):
+          if (i == 0 and self.keepBase) or \
+             (i == t-1) or (i%2 == 1):
+              out_.append(layers_[i])
+        out = torch.cat(out_, 1)
+        return out
+
+
+class HarDBlock_v2(nn.Module):
+    def get_link(self, layer, base_ch, growth_rate, grmul):
+        if layer == 0:
+          return base_ch, 0, []
+        out_channels = growth_rate
+        link = []
+        for i in range(10):
+          dv = 2 ** i
+          if layer % dv == 0:
+            k = layer - dv
+            link.insert(0, k)
+            if i > 0:
+                out_channels *= grmul
+        out_channels = int(int(out_channels + 1) / 2) * 2
+        in_channels = 0
+        for i in link:
+          ch,_,_ = self.get_link(i, base_ch, growth_rate, grmul)
+          in_channels += ch
+        return out_channels, in_channels, link
+
+    def get_out_ch(self):
+        return self.out_channels
+
+    def __init__(self, in_channels, growth_rate, grmul, n_layers, norm_layer, dwconv=False):
+        super().__init__()
+        self.links = []
+        conv_layers_ = []
+        bnrelu_layers_ = []
+        self.layer_bias = []
+        self.out_channels = 0
+        self.norm_layer = norm_layer
+        self.out_partition = collections.defaultdict(list)
+
+        for i in range(n_layers):
+          outch, inch, link = self.get_link(i+1, in_channels, growth_rate, grmul)
+          self.links.append(link)
+          for j in link:
+            self.out_partition[j].append(outch)
+
+        cur_ch = in_channels
+        for i in range(n_layers):
+          accum_out_ch = sum( self.out_partition[i] )
+          real_out_ch = self.out_partition[i][0]
+          conv_layers_.append( nn.Conv2d(cur_ch, accum_out_ch, kernel_size=3, stride=1, padding=1, bias=True) )
+          bnrelu_layers_.append( BRLayer(real_out_ch, norm_layer) )
+          cur_ch = real_out_ch
+          if (i % 2 == 0) or (i == n_layers - 1):
+            self.out_channels += real_out_ch
+        self.conv_layers = nn.ModuleList(conv_layers_)
+        self.bnrelu_layers = nn.ModuleList(bnrelu_layers_)
+
+    def transform(self, blk, trt=False):
+        # Transform weight matrix from a pretrained HarDBlock v1
+        in_ch = blk.layers[0][0].weight.shape[1]
+        for i in range(len(self.conv_layers)):
+            link = self.links[i].copy()
+            link_ch = [blk.layers[k-1][0].weight.shape[0] if k > 0 else
+                       blk.layers[0  ][0].weight.shape[1] for k in link]
+            part = self.out_partition[i]
+            w_src = blk.layers[i][0].weight
+            b_src = blk.layers[i][0].bias
+
+
+            self.conv_layers[i].weight[0:part[0], :, :,:] = w_src[:, 0:in_ch, :,:]
+            self.layer_bias.append(b_src)
+            #if b_src is not None:
+            #  self.layer_bias[i] = b_src.view(1,-1,1,1)
+            if b_src is not None:
+                if trt:
+                    self.conv_layers[i].bias[1:part[0]] = b_src[1:]
+                    self.conv_layers[i].bias[0] = b_src[0]
+                    self.conv_layers[i].bias[part[0]:] = 0
+                    self.layer_bias[i] = None
+                else:
+                    #for pytorch, add bias with standalone tensor is more efficient than within conv.bias
+                    #this is because the amount of non-zero bias is small, 
+                    #but if we use conv.bias, the number of bias will be much larger
+                    self.conv_layers[i].bias = None
+            else:
+                self.conv_layers[i].bias = None 
+
+
+            in_ch = part[0]
+            link_ch.reverse()
+            link.reverse()
+            if len(link) > 1:
+                for j in range(1, len(link) ):
+                    ly  = link[j]
+                    part_id  = self.out_partition[ly].index(part[0])
+                    chos = sum( self.out_partition[ly][0:part_id] )
+                    choe = chos + part[0]
+                    chis = sum( link_ch[0:j] )
+                    chie = chis + link_ch[j]
+                    self.conv_layers[ly].weight[chos:choe, :,:,:] = w_src[:, chis:chie,:,:]
+
+            #update BatchNorm or remove it if there is no BatchNorm in the v1 block
+            self.bnrelu_layers[i] = None
+            if isinstance(blk.layers[i][1], self.norm_layer):
+                self.bnrelu_layers[i] = nn.Sequential(
+                         blk.layers[i][1],
+                         blk.layers[i][2])
+            else:
+                self.bnrelu_layers[i] = blk.layers[i][1]
+
+    def forward(self, x):
+        layers_ = []
+        outs_ = []
+        xin = x
+        for i in range(len(self.conv_layers)):
+            link = self.links[i]
+            part = self.out_partition[i]
+
+            xout = self.conv_layers[i](xin)
+            layers_.append(xout)
+
+            xin = xout[:,0:part[0],:,:] if len(part) > 1 else xout
+            if self.layer_bias[i] is not None:
+                xin += self.layer_bias[i].view(1,-1,1,1)
+
+            if len(link) > 1:
+                for j in range( len(link) - 1 ):
+                    ly  = link[j]
+                    part_id  = self.out_partition[ly].index(part[0])
+                    chs = sum( self.out_partition[ly][0:part_id] )
+                    che = chs + part[0]
+
+                    xin += layers_[ly][:,chs:che,:,:]
+
+            xin = self.bnrelu_layers[i](xin)
+
+            if i%2 == 0 or i == len(self.conv_layers)-1:
+              outs_.append(xin)
+
+        out = torch.cat(outs_, 1)
+        return out
+
+
+class HarDNetBase(nn.Module):
+    def __init__(self, arch, norm_layer, depth_wise=False):
+        super().__init__()
+        if arch == 85:
+          first_ch  = [48, 96]
+          second_kernel = 3
+
+          ch_list = [  192, 256, 320, 480, 720]
+          grmul = 1.7
+          gr       = [  24, 24, 28, 36, 48]
+          n_layers = [   8, 16, 16, 16, 16]
+        elif arch == 68:
+          first_ch  = [32, 64]
+          second_kernel = 3
+
+          ch_list = [  128, 256, 320, 640]
+          grmul = 1.7
+          gr       = [  14, 16, 20, 40]
+          n_layers = [   8, 16, 16, 16]
+        else:
+          print("Error: HarDNet",arch," has no implementation.")
+          exit()
+
+        blks = len(n_layers)
+        self.base = nn.ModuleList([])
+
+        # First Layer: Standard Conv3x3, Stride=2
+        self.base.append (
+             ConvLayer(in_channels=3, out_channels=first_ch[0], norm_layer=norm_layer, kernel=3,
+                       stride=2,  bias=False) )
+
+        # Second Layer
+        self.base.append ( ConvLayer(first_ch[0], first_ch[1], norm_layer, kernel=second_kernel) )
+
+        # Maxpooling or DWConv3x3 downsampling
+        self.base.append(nn.AvgPool2d(kernel_size=3, stride=2, padding=1))
+
+        # Build all HarDNet blocks
+        ch = first_ch[1]
+        for i in range(blks):
+            blk = HarDBlock(ch, gr[i], grmul, n_layers[i], norm_layer, dwconv=depth_wise)
+            ch = blk.get_out_ch()
+            self.base.append ( blk )
+
+            if i != blks-1:            
+              self.base.append ( ConvLayer(ch, ch_list[i], norm_layer, kernel=1) )
+            ch = ch_list[i]
+            if i== 0:
+                self.base.append(nn.AvgPool2d(kernel_size=2, stride=2, ceil_mode=True))
+            elif i != blks-1 and i != 1 and i != 3:
+                self.base.append(nn.AvgPool2d(kernel_size=2, stride=2))
+        
+
+def fill_fc_weights(layers):
+    for m in layers.modules():
+        if isinstance(m, nn.Conv2d):
+            if m.weight is not None:
+                init.kaiming_uniform_(m.weight, nonlinearity='relu')
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+
+def weights_init(m):
+        for key in m.state_dict():
+            if key.split('.')[-1] == 'weight':
+                if 'conv' in key:
+                    init.kaiming_uniform_(m.state_dict()[key], nonlinearity='relu')
+                if 'bn' in key:
+                    m.state_dict()[key][...] = 1
+            elif key.split('.')[-1] == 'bias':
+                m.state_dict()[key][...] = 0
+
+
+class TransitionUp(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+
+    def forward(self, x, skip, concat=True):
+        out = F.interpolate(
+                x,
+                size=(skip.size(2), skip.size(3)),
+                mode="bilinear",
+                align_corners=True)
+        if concat:
+          out = torch.cat([out, skip], 1)
+        return out
+
+@SPPE.register_module
+class HarDNetPose(nn.Module):
+    def __init__(self, norm_layer=nn.BatchNorm2d, **cfg):
+        super(HarDNetPose, self).__init__()
+        assert cfg['DOWN_RATIO'] in [2, 4, 8, 16]
+        self.norm_layer = norm_layer
+        self._preset_cfg = cfg['PRESET']
+        self.first_level = int(np.log2(cfg['DOWN_RATIO']))-1
+        self.trt = cfg['TRT']
+        
+        self.base = HarDNetBase(cfg['NUM_LAYERS'], norm_layer).base
+        self.last_pool = nn.AvgPool2d(kernel_size=2, stride=2)
+        
+        if cfg['NUM_LAYERS'] == 85:
+          self.last_proj = ConvLayer(784, 256, norm_layer, kernel=1)
+          self.last_blk = HarDBlock(768, 80, 1.7, 8, norm_layer)
+          self.skip_nodes = [1,3,8,13]
+          self.SC = [32, 32, 0]
+          gr = [64, 48, 28]
+          layers = [8, 8, 4]
+          ch_list2 = [224 + self.SC[0], 160 + self.SC[1], 96 + self.SC[2]]
+          channels = [96, 214, 458, 784]
+          self.skip_lv = 3
+          scales = [2 ** i for i in range(len(channels[self.first_level:]))]
+
+        elif cfg['NUM_LAYERS'] == 68:
+          self.last_proj = ConvLayer(654, 192, norm_layer, kernel=1)
+          self.last_blk = HarDBlock(576, 72, 1.7, 8, norm_layer)
+          self.skip_nodes = [1,3,8,11]
+          self.SC = [32, 32, 0 ]  
+          gr = [48, 32, 20]
+          layers = [8, 8, 4]
+          ch_list2 = [224+self.SC[0], 96+self.SC[1], 64+self.SC[2]]
+          channels = [64, 124, 328, 654]
+          self.skip_lv = 2
+          scales = [2 ** i for i in range(len(channels[self.first_level:]))]
+
+        
+
+        self.transUpBlocks = nn.ModuleList([])
+        self.denseBlocksUp = nn.ModuleList([])
+        self.conv1x1_up    = nn.ModuleList([])
+        self.avg9x9   = nn.AvgPool2d(kernel_size=(9,9), stride=1, padding=(4,4))
+        prev_ch = self.last_blk.get_out_ch()
+
+        for i in range(3):
+            skip_ch = channels[3-i]
+            self.transUpBlocks.append(TransitionUp(prev_ch, prev_ch))
+            if i < self.skip_lv:
+              cur_ch = prev_ch + skip_ch
+            else:
+              cur_ch = prev_ch
+            self.conv1x1_up.append(ConvLayer(cur_ch, ch_list2[i], norm_layer, kernel=1))
+            cur_ch = ch_list2[i]
+            cur_ch -= self.SC[i]
+            cur_ch *= 3
+
+            blk = HarDBlock(cur_ch, gr[i], 1.7, layers[i], norm_layer)
+
+            self.denseBlocksUp.append(blk)
+            prev_ch = blk.get_out_ch()
+        
+        prev_ch += self.SC[0] + self.SC[1] + self.SC[2]
+
+        weights_init(self.denseBlocksUp) 
+        weights_init(self.conv1x1_up) 
+        weights_init(self.last_blk)
+        weights_init(self.last_proj)
+        
+        out_channel = self._preset_cfg['NUM_JOINTS']
+
+        ch = max(128, out_channel*4)
+        self.conv_out = nn.Sequential(
+            nn.Conv2d(prev_ch, ch,
+              kernel_size=3, padding=1, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(ch, out_channel, 
+              kernel_size=cfg['FINAL_CONV_KERNEL'], stride=1, 
+              padding=cfg['FINAL_CONV_KERNEL'] // 2, bias=True))
+        fill_fc_weights(self.conv_out)
+        self.conv_out[-1].bias.data.fill_(-2.19)
+
+
+    def v2_transform(self):
+        print('Transform HarDBlock v2..')
+        for i in range( len(self.base)):
+            if isinstance(self.base[i], HarDBlock):
+                blk = self.base[i]
+                self.base[i] = HarDBlock_v2(blk.in_channels, blk.growth_rate, blk.grmul, blk.n_layers, blk.norm_layer)
+                self.base[i].transform(blk, self.trt)
+        blk = self.last_blk
+        self.last_blk = HarDBlock_v2(blk.in_channels, blk.growth_rate, blk.grmul, blk.n_layers, blk.norm_layer)
+        self.last_blk.transform(blk, self.trt)
+        for i in range(3):
+            blk = self.denseBlocksUp[i]
+            self.denseBlocksUp[i] = HarDBlock_v2(blk.in_channels, blk.growth_rate, blk.grmul, blk.n_layers, blk.norm_layer)
+            self.denseBlocksUp[i].transform(blk, self.trt)
+
+    def forward(self, x):
+        xs = []
+        x_sc = []
+        
+        for i in range(len(self.base)):
+            x = self.base[i](x)
+            if i in self.skip_nodes:
+                xs.append(x)
+        
+        x = self.last_proj(x)
+        x = self.last_pool(x)
+        x2 = self.avg9x9(x)
+        x3 = x/(x.sum((2,3),keepdim=True) + 0.1)
+        x = torch.cat([x,x2,x3],1)
+        x = self.last_blk(x)
+        
+        for i in range(3):
+            skip_x = xs[3-i]
+            x = self.transUpBlocks[i](x, skip_x, (i<self.skip_lv))
+            x = self.conv1x1_up[i](x)
+            if self.SC[i] > 0:
+              end = x.shape[1]
+              x_sc.append( x[:,end-self.SC[i]:,:,:].contiguous() )
+              x = x[:,:end-self.SC[i],:,:].contiguous()
+            x2 = self.avg9x9(x)
+            x3 = x/(x.sum((2,3),keepdim=True) + 0.1)
+            x = torch.cat([x,x2,x3],1)
+            x = self.denseBlocksUp[i](x)
+
+        scs = [x]
+        for i in range(3):
+          if self.SC[i] > 0:
+            scs.insert(0,  F.interpolate(
+                            x_sc[i], size=(x.size(2), x.size(3)), 
+                            mode="bilinear", align_corners=True) )
+        x = torch.cat(scs,1)
+        x = self.conv_out(x)
+        return x
+    
+    def _initialize(self, pretrained=''):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                nn.init.normal_(m.weight, std=0.001)
+                for name, _ in m.named_parameters():
+                    if name in ['bias']:
+                        nn.init.constant_(m.bias, 0)
+            elif isinstance(m, self.norm_layer):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.ConvTranspose2d):
+                nn.init.normal_(m.weight, std=0.001)
+                for name, _ in m.named_parameters():
+                    if name in ['bias']:
+                        nn.init.constant_(m.bias, 0)
+
+        if os.path.isfile(pretrained):
+            pretrained_state_dict = torch.load(pretrained)
+
+            need_init_state_dict = {}
+            for name, m in pretrained_state_dict.items():
+                if name.split('.')[0] in self.pretrained_layers \
+                   or self.pretrained_layers[0] == '*':
+                    need_init_state_dict[name] = m
+            self.load_state_dict(need_init_state_dict, strict=False)
+        elif pretrained:
+            raise ValueError('{} is not exist!'.format(pretrained))
+
+
+def get_pose_net(cfg, is_train, **kwargs):
+  model = HarDNetPose(cfg, **kwargs)
+  if is_train and cfg.MODEL.INIT_WEIGHTS:
+        model._initialize(cfg.MODEL.INIT_WEIGHTS)
+  total_params = sum(p.numel() for p in model.parameters())
+  print( "Parameters=", total_params )
+  return model

alphapose/models/hrnet.py

@@ -0,0 +1,494 @@
+# ------------------------------------------------------------------------------
+# Copyright (c) Microsoft
+# Licensed under the MIT License.
+# Written by Bin Xiao (Bin.Xiao@microsoft.com)
+# ------------------------------------------------------------------------------
+
+from __future__ import absolute_import, division, print_function
+
+import os
+
+import torch
+import torch.nn as nn
+
+from .builder import SPPE
+
+BN_MOMENTUM = 0.1
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                               padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
+                               bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion,
+                                  momentum=BN_MOMENTUM)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class HighResolutionModule(nn.Module):
+    def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
+                 num_channels, fuse_method, multi_scale_output=True):
+        super(HighResolutionModule, self).__init__()
+        self._check_branches(
+            num_branches, blocks, num_blocks, num_inchannels, num_channels)
+
+        self.num_inchannels = num_inchannels
+        self.fuse_method = fuse_method
+        self.num_branches = num_branches
+
+        self.multi_scale_output = multi_scale_output
+
+        self.branches = self._make_branches(
+            num_branches, blocks, num_blocks, num_channels)
+        self.fuse_layers = self._make_fuse_layers()
+        self.relu = nn.ReLU(True)
+
+    def _check_branches(self, num_branches, blocks, num_blocks,
+                        num_inchannels, num_channels):
+        if num_branches != len(num_blocks):
+            error_msg = 'NUM_BRANCHES({}) <> NUM_BLOCKS({})'.format(
+                num_branches, len(num_blocks))
+            raise ValueError(error_msg)
+
+        if num_branches != len(num_channels):
+            error_msg = 'NUM_BRANCHES({}) <> NUM_CHANNELS({})'.format(
+                num_branches, len(num_channels))
+            raise ValueError(error_msg)
+
+        if num_branches != len(num_inchannels):
+            error_msg = 'NUM_BRANCHES({}) <> NUM_INCHANNELS({})'.format(
+                num_branches, len(num_inchannels))
+            raise ValueError(error_msg)
+
+    def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
+                         stride=1):
+        downsample = None
+        if stride != 1 or \
+           self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(
+                    self.num_inchannels[branch_index],
+                    num_channels[branch_index] * block.expansion,
+                    kernel_size=1, stride=stride, bias=False
+                ),
+                nn.BatchNorm2d(
+                    num_channels[branch_index] * block.expansion,
+                    momentum=BN_MOMENTUM
+                ),
+            )
+
+        layers = []
+        layers.append(
+            block(
+                self.num_inchannels[branch_index],
+                num_channels[branch_index],
+                stride,
+                downsample
+            )
+        )
+        self.num_inchannels[branch_index] = \
+            num_channels[branch_index] * block.expansion
+        for i in range(1, num_blocks[branch_index]):
+            layers.append(
+                block(
+                    self.num_inchannels[branch_index],
+                    num_channels[branch_index]
+                )
+            )
+
+        return nn.Sequential(*layers)
+
+    def _make_branches(self, num_branches, block, num_blocks, num_channels):
+        branches = []
+
+        for i in range(num_branches):
+            branches.append(
+                self._make_one_branch(i, block, num_blocks, num_channels)
+            )
+
+        return nn.ModuleList(branches)
+
+    def _make_fuse_layers(self):
+        if self.num_branches == 1:
+            return None
+
+        num_branches = self.num_branches
+        num_inchannels = self.num_inchannels
+        fuse_layers = []
+        for i in range(num_branches if self.multi_scale_output else 1):
+            fuse_layer = []
+            for j in range(num_branches):
+                if j > i:
+                    fuse_layer.append(
+                        nn.Sequential(
+                            nn.Conv2d(
+                                num_inchannels[j],
+                                num_inchannels[i],
+                                1, 1, 0, bias=False
+                            ),
+                            nn.BatchNorm2d(num_inchannels[i]),
+                            nn.Upsample(scale_factor=2 **
+                                        (j - i), mode='nearest')
+                        )
+                    )
+                elif j == i:
+                    fuse_layer.append(None)
+                else:
+                    conv3x3s = []
+                    for k in range(i - j):
+                        if k == i - j - 1:
+                            num_outchannels_conv3x3 = num_inchannels[i]
+                            conv3x3s.append(
+                                nn.Sequential(
+                                    nn.Conv2d(
+                                        num_inchannels[j],
+                                        num_outchannels_conv3x3,
+                                        3, 2, 1, bias=False
+                                    ),
+                                    nn.BatchNorm2d(num_outchannels_conv3x3)
+                                )
+                            )
+                        else:
+                            num_outchannels_conv3x3 = num_inchannels[j]
+                            conv3x3s.append(
+                                nn.Sequential(
+                                    nn.Conv2d(
+                                        num_inchannels[j],
+                                        num_outchannels_conv3x3,
+                                        3, 2, 1, bias=False
+                                    ),
+                                    nn.BatchNorm2d(num_outchannels_conv3x3),
+                                    nn.ReLU(True)
+                                )
+                            )
+                    fuse_layer.append(nn.Sequential(*conv3x3s))
+            fuse_layers.append(nn.ModuleList(fuse_layer))
+
+        return nn.ModuleList(fuse_layers)
+
+    def get_num_inchannels(self):
+        return self.num_inchannels
+
+    def forward(self, x):
+        if self.num_branches == 1:
+            return [self.branches[0](x[0])]
+
+        for i in range(self.num_branches):
+            x[i] = self.branches[i](x[i])
+
+        x_fuse = []
+
+        for i in range(len(self.fuse_layers)):
+            y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
+            for j in range(1, self.num_branches):
+                if i == j:
+                    y = y + x[j]
+                else:
+                    y = y + self.fuse_layers[i][j](x[j])
+            x_fuse.append(self.relu(y))
+
+        return x_fuse
+
+
+blocks_dict = {
+    'BASIC': BasicBlock,
+    'BOTTLENECK': Bottleneck
+}
+
+
+@SPPE.register_module
+class PoseHighResolutionNet(nn.Module):
+
+    def __init__(self, **cfg):
+        self.inplanes = 64
+        super(PoseHighResolutionNet, self).__init__()
+        self._preset_cfg = cfg['PRESET']
+
+        # stem net
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
+        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1,
+                               bias=False)
+        self.bn2 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
+        self.relu = nn.ReLU(inplace=True)
+        self.layer1 = self._make_layer(Bottleneck, 64, 4)
+
+        self.stage2_cfg = cfg['STAGE2']
+        num_channels = self.stage2_cfg['NUM_CHANNELS']
+        block = blocks_dict[self.stage2_cfg['BLOCK']]
+        num_channels = [
+            num_channels[i] * block.expansion for i in range(len(num_channels))
+        ]
+        self.transition1 = self._make_transition_layer([256], num_channels)
+        self.stage2, pre_stage_channels = self._make_stage(
+            self.stage2_cfg, num_channels)
+
+        self.stage3_cfg = cfg['STAGE3']
+        num_channels = self.stage3_cfg['NUM_CHANNELS']
+        block = blocks_dict[self.stage3_cfg['BLOCK']]
+        num_channels = [
+            num_channels[i] * block.expansion for i in range(len(num_channels))
+        ]
+        self.transition2 = self._make_transition_layer(
+            pre_stage_channels, num_channels)
+        self.stage3, pre_stage_channels = self._make_stage(
+            self.stage3_cfg, num_channels)
+
+        self.stage4_cfg = cfg['STAGE4']
+        num_channels = self.stage4_cfg['NUM_CHANNELS']
+        block = blocks_dict[self.stage4_cfg['BLOCK']]
+        num_channels = [
+            num_channels[i] * block.expansion for i in range(len(num_channels))
+        ]
+        self.transition3 = self._make_transition_layer(
+            pre_stage_channels, num_channels)
+        self.stage4, pre_stage_channels = self._make_stage(
+            self.stage4_cfg, num_channels, multi_scale_output=False)
+
+        self.final_layer = nn.Conv2d(
+            in_channels=pre_stage_channels[0],
+            out_channels=self._preset_cfg['NUM_JOINTS'],
+            kernel_size=cfg['FINAL_CONV_KERNEL'],
+            stride=1,
+            padding=1 if cfg['FINAL_CONV_KERNEL'] == 3 else 0
+        )
+
+        self.pretrained_layers = cfg['PRETRAINED_LAYERS']
+
+    def _make_transition_layer(
+            self, num_channels_pre_layer, num_channels_cur_layer):
+        num_branches_cur = len(num_channels_cur_layer)
+        num_branches_pre = len(num_channels_pre_layer)
+
+        transition_layers = []
+        for i in range(num_branches_cur):
+            if i < num_branches_pre:
+                if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
+                    transition_layers.append(
+                        nn.Sequential(
+                            nn.Conv2d(
+                                num_channels_pre_layer[i],
+                                num_channels_cur_layer[i],
+                                3, 1, 1, bias=False
+                            ),
+                            nn.BatchNorm2d(num_channels_cur_layer[i]),
+                            nn.ReLU(inplace=True)
+                        )
+                    )
+                else:
+                    transition_layers.append(None)
+            else:
+                conv3x3s = []
+                for j in range(i + 1 - num_branches_pre):
+                    inchannels = num_channels_pre_layer[-1]
+                    outchannels = num_channels_cur_layer[i] \
+                        if j == i - num_branches_pre else inchannels
+                    conv3x3s.append(
+                        nn.Sequential(
+                            nn.Conv2d(
+                                inchannels, outchannels, 3, 2, 1, bias=False
+                            ),
+                            nn.BatchNorm2d(outchannels),
+                            nn.ReLU(inplace=True)
+                        )
+                    )
+                transition_layers.append(nn.Sequential(*conv3x3s))
+
+        return nn.ModuleList(transition_layers)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(
+                    self.inplanes, planes * block.expansion,
+                    kernel_size=1, stride=stride, bias=False
+                ),
+                nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def _make_stage(self, layer_config, num_inchannels,
+                    multi_scale_output=True):
+        num_modules = layer_config['NUM_MODULES']
+        num_branches = layer_config['NUM_BRANCHES']
+        num_blocks = layer_config['NUM_BLOCKS']
+        num_channels = layer_config['NUM_CHANNELS']
+        block = blocks_dict[layer_config['BLOCK']]
+        fuse_method = layer_config['FUSE_METHOD']
+
+        modules = []
+        for i in range(num_modules):
+            # multi_scale_output is only used last module
+            if not multi_scale_output and i == num_modules - 1:
+                reset_multi_scale_output = False
+            else:
+                reset_multi_scale_output = True
+
+            modules.append(
+                HighResolutionModule(
+                    num_branches,
+                    block,
+                    num_blocks,
+                    num_inchannels,
+                    num_channels,
+                    fuse_method,
+                    reset_multi_scale_output
+                )
+            )
+            num_inchannels = modules[-1].get_num_inchannels()
+
+        return nn.Sequential(*modules), num_inchannels
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+        x = self.layer1(x)
+
+        x_list = []
+        for i in range(self.stage2_cfg['NUM_BRANCHES']):
+            if self.transition1[i] is not None:
+                x_list.append(self.transition1[i](x))
+            else:
+                x_list.append(x)
+        y_list = self.stage2(x_list)
+
+        x_list = []
+        for i in range(self.stage3_cfg['NUM_BRANCHES']):
+            if self.transition2[i] is not None:
+                x_list.append(self.transition2[i](y_list[-1]))
+            else:
+                x_list.append(y_list[i])
+        y_list = self.stage3(x_list)
+
+        x_list = []
+        for i in range(self.stage4_cfg['NUM_BRANCHES']):
+            if self.transition3[i] is not None:
+                x_list.append(self.transition3[i](y_list[-1]))
+            else:
+                x_list.append(y_list[i])
+        y_list = self.stage4(x_list)
+
+        x = self.final_layer(y_list[0])
+
+        return x
+
+    def _initialize(self, pretrained=''):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                nn.init.normal_(m.weight, std=0.001)
+                for name, _ in m.named_parameters():
+                    if name in ['bias']:
+                        nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.ConvTranspose2d):
+                nn.init.normal_(m.weight, std=0.001)
+                for name, _ in m.named_parameters():
+                    if name in ['bias']:
+                        nn.init.constant_(m.bias, 0)
+
+        if os.path.isfile(pretrained):
+            pretrained_state_dict = torch.load(pretrained)
+
+            need_init_state_dict = {}
+            for name, m in pretrained_state_dict.items():
+                if name.split('.')[0] in self.pretrained_layers \
+                   or self.pretrained_layers[0] == '*':
+                    need_init_state_dict[name] = m
+            self.load_state_dict(need_init_state_dict, strict=False)
+        elif pretrained:
+            raise ValueError('{} is not exist!'.format(pretrained))
+
+
+def get_pose_net(cfg, is_train, **kwargs):
+    model = PoseHighResolutionNet(cfg, **kwargs)
+
+    if is_train and cfg.MODEL.INIT_WEIGHTS:
+        model._initialize(cfg.MODEL.PRETRAINED)
+
+    return model

alphapose/models/simplepose.py

@@ -0,0 +1,86 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+
+import torch.nn as nn
+
+from .builder import SPPE
+from .layers.Resnet import ResNet
+
+
+@SPPE.register_module
+class SimplePose(nn.Module):
+    def __init__(self, norm_layer=nn.BatchNorm2d, **cfg):
+        super(SimplePose, self).__init__()
+        self._preset_cfg = cfg['PRESET']
+        self.deconv_dim = cfg['NUM_DECONV_FILTERS']
+        self._norm_layer = norm_layer
+
+        self.preact = ResNet(f"resnet{cfg['NUM_LAYERS']}")
+
+        # Imagenet pretrain model
+        import torchvision.models as tm   # noqa: F401,F403
+        assert cfg['NUM_LAYERS'] in [18, 34, 50, 101, 152]
+        x = eval(f"tm.resnet{cfg['NUM_LAYERS']}(pretrained=True)")
+
+        model_state = self.preact.state_dict()
+        state = {k: v for k, v in x.state_dict().items()
+                 if k in self.preact.state_dict() and v.size() == self.preact.state_dict()[k].size()}
+        model_state.update(state)
+        self.preact.load_state_dict(model_state)
+
+        self.deconv_layers = self._make_deconv_layer()
+        self.final_layer = nn.Conv2d(
+            self.deconv_dim[2], self._preset_cfg['NUM_JOINTS'], kernel_size=1, stride=1, padding=0)
+
+    def _make_deconv_layer(self):
+        deconv_layers = []
+        deconv1 = nn.ConvTranspose2d(
+            2048, self.deconv_dim[0], kernel_size=4, stride=2, padding=int(4 / 2) - 1, bias=False)
+        bn1 = self._norm_layer(self.deconv_dim[0])
+        deconv2 = nn.ConvTranspose2d(
+            self.deconv_dim[0], self.deconv_dim[1], kernel_size=4, stride=2, padding=int(4 / 2) - 1, bias=False)
+        bn2 = self._norm_layer(self.deconv_dim[1])
+        deconv3 = nn.ConvTranspose2d(
+            self.deconv_dim[1], self.deconv_dim[2], kernel_size=4, stride=2, padding=int(4 / 2) - 1, bias=False)
+        bn3 = self._norm_layer(self.deconv_dim[2])
+
+        deconv_layers.append(deconv1)
+        deconv_layers.append(bn1)
+        deconv_layers.append(nn.ReLU(inplace=True))
+        deconv_layers.append(deconv2)
+        deconv_layers.append(bn2)
+        deconv_layers.append(nn.ReLU(inplace=True))
+        deconv_layers.append(deconv3)
+        deconv_layers.append(bn3)
+        deconv_layers.append(nn.ReLU(inplace=True))
+
+        return nn.Sequential(*deconv_layers)
+
+    def _initialize(self):
+        for name, m in self.deconv_layers.named_modules():
+            if isinstance(m, nn.ConvTranspose2d):
+                # logger.info('=> init {}.weight as normal(0, 0.001)'.format(name))
+                # logger.info('=> init {}.bias as 0'.format(name))
+                nn.init.normal_(m.weight, std=0.001)
+                # if self.deconv_with_bias:
+                #     nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                # logger.info('=> init {}.weight as 1'.format(name))
+                # logger.info('=> init {}.bias as 0'.format(name))
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+        for m in self.final_layer.modules():
+            if isinstance(m, nn.Conv2d):
+                # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                # logger.info('=> init {}.weight as normal(0, 0.001)'.format(name))
+                # logger.info('=> init {}.bias as 0'.format(name))
+                nn.init.normal_(m.weight, std=0.001)
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        out = self.preact(x)
+        out = self.deconv_layers(out)
+        out = self.final_layer(out)
+        return out

alphapose/opt.py

@@ -0,0 +1,87 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+import argparse
+import logging
+import os
+from types import MethodType
+
+import torch
+
+from .utils.config import update_config
+
+parser = argparse.ArgumentParser(description='AlphaPose Training')
+
+"----------------------------- Experiment options -----------------------------"
+parser.add_argument('--cfg',
+                    help='experiment configure file name',
+                    required=True,
+                    type=str)
+parser.add_argument('--exp-id', default='default', type=str,
+                    help='Experiment ID')
+
+"----------------------------- General options -----------------------------"
+parser.add_argument('--nThreads', default=60, type=int,
+                    help='Number of data loading threads')
+parser.add_argument('--snapshot', default=2, type=int,
+                    help='How often to take a snapshot of the model (0 = never)')
+
+parser.add_argument('--rank', default=-1, type=int,
+                    help='node rank for distributed training')
+parser.add_argument('--dist-url', default='tcp://192.168.1.214:23345', type=str,
+                    help='url used to set up distributed training')
+parser.add_argument('--dist-backend', default='nccl', type=str,
+                    help='distributed backend')
+parser.add_argument('--launcher', choices=['none', 'pytorch', 'slurm', 'mpi'], default='none',
+                    help='job launcher')
+
+"----------------------------- Training options -----------------------------"
+parser.add_argument('--sync', default=False, dest='sync',
+                    help='Use Sync Batchnorm', action='store_true')
+parser.add_argument('--detector', dest='detector',
+                    help='detector name', default="yolo")
+
+"----------------------------- Log options -----------------------------"
+parser.add_argument('--board', default=True, dest='board',
+                    help='Logging with tensorboard', action='store_true')
+parser.add_argument('--debug', default=False, dest='debug',
+                    help='Visualization debug', action='store_true')
+parser.add_argument('--map', default=True, dest='map',
+                    help='Evaluate mAP per epoch', action='store_true')
+
+
+opt = parser.parse_args()
+cfg_file_name = os.path.basename(opt.cfg)
+cfg = update_config(opt.cfg)
+
+cfg['FILE_NAME'] = cfg_file_name
+cfg.TRAIN.DPG_STEP = [i - cfg.TRAIN.DPG_MILESTONE for i in cfg.TRAIN.DPG_STEP]
+opt.world_size = cfg.TRAIN.WORLD_SIZE
+opt.work_dir = './exp/{}-{}/'.format(opt.exp_id, cfg_file_name)
+opt.gpus = [i for i in range(torch.cuda.device_count())]
+opt.device = torch.device("cuda:" + str(opt.gpus[0]) if opt.gpus[0] >= 0 else "cpu")
+
+if not os.path.exists("./exp/{}-{}".format(opt.exp_id, cfg_file_name)):
+    os.makedirs("./exp/{}-{}".format(opt.exp_id, cfg_file_name))
+
+filehandler = logging.FileHandler(
+    './exp/{}-{}/training.log'.format(opt.exp_id, cfg_file_name))
+streamhandler = logging.StreamHandler()
+
+logger = logging.getLogger('')
+logger.setLevel(logging.INFO)
+logger.addHandler(filehandler)
+logger.addHandler(streamhandler)
+
+
+def epochInfo(self, set, idx, loss, acc):
+    self.info('{set}-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format(
+        set=set,
+        idx=idx,
+        loss=loss,
+        acc=acc
+    ))
+
+
+logger.epochInfo = MethodType(epochInfo, logger)

alphapose/utils/__init__.py

@@ -0,0 +1,5 @@
+from .registry import Registry, build_from_cfg, retrieve_from_cfg
+
+__all__ = [
+    'Registry', 'build_from_cfg', 'retrieve_from_cfg'
+]

alphapose/utils/bbox.py

@@ -0,0 +1,251 @@
+from __future__ import division
+
+import torch
+import numpy as np
+
+
+def bbox_iou(bbox_a, bbox_b, offset=0):
+    """Calculate Intersection-Over-Union(IOU) of two bounding boxes.
+
+    Parameters
+    ----------
+    bbox_a : numpy.ndarray
+        An ndarray with shape :math:`(N, 4)`.
+    bbox_b : numpy.ndarray
+        An ndarray with shape :math:`(M, 4)`.
+    offset : float or int, default is 0
+        The ``offset`` is used to control the whether the width(or height) is computed as
+        (right - left + ``offset``).
+        Note that the offset must be 0 for normalized bboxes, whose ranges are in ``[0, 1]``.
+
+    Returns
+    -------
+    numpy.ndarray
+        An ndarray with shape :math:`(N, M)` indicates IOU between each pairs of
+        bounding boxes in `bbox_a` and `bbox_b`.
+
+    """
+    if bbox_a.shape[1] < 4 or bbox_b.shape[1] < 4:
+        raise IndexError("Bounding boxes axis 1 must have at least length 4")
+
+    tl = np.maximum(bbox_a[:, None, :2], bbox_b[:, :2])
+    br = np.minimum(bbox_a[:, None, 2:4], bbox_b[:, 2:4])
+
+    area_i = np.prod(br - tl + offset, axis=2) * (tl < br).all(axis=2)
+    area_a = np.prod(bbox_a[:, 2:4] - bbox_a[:, :2] + offset, axis=1)
+    area_b = np.prod(bbox_b[:, 2:4] - bbox_b[:, :2] + offset, axis=1)
+    return area_i / (area_a[:, None] + area_b - area_i)
+
+
+def bbox_xywh_to_xyxy(xywh):
+    """Convert bounding boxes from format (x, y, w, h) to (xmin, ymin, xmax, ymax)
+
+    Parameters
+    ----------
+    xywh : list, tuple or numpy.ndarray
+        The bbox in format (x, y, w, h).
+        If numpy.ndarray is provided, we expect multiple bounding boxes with
+        shape `(N, 4)`.
+
+    Returns
+    -------
+    tuple or numpy.ndarray
+        The converted bboxes in format (xmin, ymin, xmax, ymax).
+        If input is numpy.ndarray, return is numpy.ndarray correspondingly.
+
+    """
+    if isinstance(xywh, (tuple, list)):
+        if not len(xywh) == 4:
+            raise IndexError(
+                "Bounding boxes must have 4 elements, given {}".format(len(xywh)))
+        w, h = np.maximum(xywh[2] - 1, 0), np.maximum(xywh[3] - 1, 0)
+        return (xywh[0], xywh[1], xywh[0] + w, xywh[1] + h)
+    elif isinstance(xywh, np.ndarray):
+        if not xywh.size % 4 == 0:
+            raise IndexError(
+                "Bounding boxes must have n * 4 elements, given {}".format(xywh.shape))
+        xyxy = np.hstack((xywh[:, :2], xywh[:, :2] + np.maximum(0, xywh[:, 2:4] - 1)))
+        return xyxy
+    else:
+        raise TypeError(
+            'Expect input xywh a list, tuple or numpy.ndarray, given {}'.format(type(xywh)))
+
+
+def bbox_xyxy_to_xywh(xyxy):
+    """Convert bounding boxes from format (xmin, ymin, xmax, ymax) to (x, y, w, h).
+
+    Parameters
+    ----------
+    xyxy : list, tuple or numpy.ndarray
+        The bbox in format (xmin, ymin, xmax, ymax).
+        If numpy.ndarray is provided, we expect multiple bounding boxes with
+        shape `(N, 4)`.
+
+    Returns
+    -------
+    tuple or numpy.ndarray
+        The converted bboxes in format (x, y, w, h).
+        If input is numpy.ndarray, return is numpy.ndarray correspondingly.
+
+    """
+    if isinstance(xyxy, (tuple, list)):
+        if not len(xyxy) == 4:
+            raise IndexError(
+                "Bounding boxes must have 4 elements, given {}".format(len(xyxy)))
+        x1, y1 = xyxy[0], xyxy[1]
+        w, h = xyxy[2] - x1 + 1, xyxy[3] - y1 + 1
+        return (x1, y1, w, h)
+    elif isinstance(xyxy, np.ndarray):
+        if not xyxy.size % 4 == 0:
+            raise IndexError(
+                "Bounding boxes must have n * 4 elements, given {}".format(xyxy.shape))
+        return np.hstack((xyxy[:, :2], xyxy[:, 2:4] - xyxy[:, :2] + 1))
+    else:
+        raise TypeError(
+            'Expect input xywh a list, tuple or numpy.ndarray, given {}'.format(type(xyxy)))
+
+
+def bbox_clip_xyxy(xyxy, width, height):
+    """Clip bounding box with format (xmin, ymin, xmax, ymax) to specified boundary.
+
+    All bounding boxes will be clipped to the new region `(0, 0, width, height)`.
+
+    Parameters
+    ----------
+    xyxy : list, tuple or numpy.ndarray
+        The bbox in format (xmin, ymin, xmax, ymax).
+        If numpy.ndarray is provided, we expect multiple bounding boxes with
+        shape `(N, 4)`.
+    width : int or float
+        Boundary width.
+    height : int or float
+        Boundary height.
+
+    Returns
+    -------
+    type
+        Description of returned object.
+
+    """
+    if isinstance(xyxy, (tuple, list)):
+        if not len(xyxy) == 4:
+            raise IndexError(
+                "Bounding boxes must have 4 elements, given {}".format(len(xyxy)))
+        x1 = np.minimum(width - 1, np.maximum(0, xyxy[0]))
+        y1 = np.minimum(height - 1, np.maximum(0, xyxy[1]))
+        x2 = np.minimum(width - 1, np.maximum(0, xyxy[2]))
+        y2 = np.minimum(height - 1, np.maximum(0, xyxy[3]))
+        return (x1, y1, x2, y2)
+    elif isinstance(xyxy, np.ndarray):
+        if not xyxy.size % 4 == 0:
+            raise IndexError(
+                "Bounding boxes must have n * 4 elements, given {}".format(xyxy.shape))
+        x1 = np.minimum(width - 1, np.maximum(0, xyxy[:, 0]))
+        y1 = np.minimum(height - 1, np.maximum(0, xyxy[:, 1]))
+        x2 = np.minimum(width - 1, np.maximum(0, xyxy[:, 2]))
+        y2 = np.minimum(height - 1, np.maximum(0, xyxy[:, 3]))
+        return np.hstack((x1, y1, x2, y2))
+    else:
+        raise TypeError(
+            'Expect input xywh a list, tuple or numpy.ndarray, given {}'.format(type(xyxy)))
+
+
+def transformBox(pt, bbox, input_size, output_size):
+    inpH, inpW = input_size
+    resH, _ = output_size
+
+    center = torch.zeros(2)
+    center[0] = (bbox[2] - 1 - bbox[0]) / 2
+    center[1] = (bbox[3] - 1 - bbox[1]) / 2
+
+    lenH = max(bbox[3] - bbox[1], (bbox[2] - bbox[0]) * inpH / inpW)
+    lenW = lenH * inpW / inpH
+
+    _pt = torch.zeros(2)
+    _pt[0] = pt[0] - bbox[0]
+    _pt[1] = pt[1] - bbox[1]
+    # Move to center
+    _pt[0] = _pt[0] + max(0, (lenW - 1) / 2 - center[0])
+    _pt[1] = _pt[1] + max(0, (lenH - 1) / 2 - center[1])
+    pt = (_pt * resH) / lenH
+    pt[0] = round(float(pt[0]))
+    pt[1] = round(float(pt[1]))
+    return pt.int()
+
+
+def transformBoxInvert(pt, bbox, resH, resW):
+    center = torch.zeros(2)
+    center[0] = (bbox[2] - 1 - bbox[0]) / 2
+    center[1] = (bbox[3] - 1 - bbox[1]) / 2
+
+    lenH = max(bbox[3] - bbox[1], (bbox[2] - bbox[0]) * resH / resW)
+    lenW = lenH * resW / resH
+
+    _pt = (pt * lenH) / resH
+
+    if bool(((lenW - 1) / 2 - center[0]) > 0):
+        _pt[0] = _pt[0] - ((lenW - 1) / 2 - center[0]).item()
+    if bool(((lenH - 1) / 2 - center[1]) > 0):
+        _pt[1] = _pt[1] - ((lenH - 1) / 2 - center[1]).item()
+
+    new_point = torch.zeros(2)
+    new_point[0] = _pt[0] + bbox[0]
+    new_point[1] = _pt[1] + bbox[1]
+    return new_point
+
+
+def _box_to_center_scale(x, y, w, h, aspect_ratio=1.0, scale_mult=1.25):
+    """Convert box coordinates to center and scale.
+    adapted from https://github.com/Microsoft/human-pose-estimation.pytorch
+    """
+    pixel_std = 1
+    center = np.zeros((2), dtype=np.float32)
+    center[0] = x + w * 0.5
+    center[1] = y + h * 0.5
+
+    if w > aspect_ratio * h:
+        h = w / aspect_ratio
+    elif w < aspect_ratio * h:
+        w = h * aspect_ratio
+    scale = np.array(
+        [w * 1.0 / pixel_std, h * 1.0 / pixel_std], dtype=np.float32)
+    if center[0] != -1:
+        scale = scale * scale_mult
+    return center, scale
+
+
+def _center_scale_to_box(center, scale):
+    pixel_std = 1.0
+    w = scale[0] * pixel_std
+    h = scale[1] * pixel_std
+    xmin = center[0] - w * 0.5
+    ymin = center[1] - h * 0.5
+    xmax = xmin + w
+    ymax = ymin + h
+    bbox = [xmin, ymin, xmax, ymax]
+    return bbox
+
+
+def _clip_aspect_ratio(boxes, aspect_ratio=1.0):
+    xmin, ymin = boxes[:, 0], boxes[:, 1]
+    xmax, ymax = boxes[:, 2], boxes[:, 3]
+
+    w = xmax - xmin
+    h = ymax - ymin
+
+    c_x = xmin + w * 0.5
+    c_y = ymin + h * 0.5
+
+    idx = w > (aspect_ratio * h)
+    h[idx] = w[idx] / aspect_ratio
+
+    idx = w < (aspect_ratio * h)
+    w[idx] = h[idx] * aspect_ratio
+
+    new_boxes = torch.zeros(boxes.shape[0], 5)
+    new_boxes[:, 1] = c_x - w * 0.5
+    new_boxes[:, 2] = c_y - h * 0.5
+    new_boxes[:, 3] = c_x + w * 0.5
+    new_boxes[:, 4] = c_y + h * 0.5
+
+    return new_boxes

alphapose/utils/config.py

@@ -0,0 +1,8 @@
+import yaml
+from easydict import EasyDict as edict
+
+
+def update_config(config_file):
+    with open(config_file) as f:
+        config = edict(yaml.load(f, Loader=yaml.FullLoader))
+        return config

alphapose/utils/detector.py

@@ -0,0 +1,280 @@
+import os
+import sys
+from threading import Thread
+from queue import Queue
+
+import cv2
+import numpy as np
+
+import torch
+import torch.multiprocessing as mp
+
+from alphapose.utils.presets import SimpleTransform
+from alphapose.models import builder
+
+class DetectionLoader():
+    def __init__(self, input_source, detector, cfg, opt, mode='image', batchSize=1, queueSize=128):
+        self.cfg = cfg
+        self.opt = opt
+        self.mode = mode
+        self.device = opt.device
+
+        if mode == 'image':
+            self.img_dir = opt.inputpath
+            self.imglist = [os.path.join(self.img_dir, im_name.rstrip('\n').rstrip('\r')) for im_name in input_source]
+            self.datalen = len(input_source)
+        elif mode == 'video':
+            stream = cv2.VideoCapture(input_source)
+            assert stream.isOpened(), 'Cannot capture source'
+            self.path = input_source
+            self.datalen = int(stream.get(cv2.CAP_PROP_FRAME_COUNT))
+            self.fourcc = int(stream.get(cv2.CAP_PROP_FOURCC))
+            self.fps = stream.get(cv2.CAP_PROP_FPS)
+            self.frameSize = (int(stream.get(cv2.CAP_PROP_FRAME_WIDTH)), int(stream.get(cv2.CAP_PROP_FRAME_HEIGHT)))
+            self.videoinfo = {'fourcc': self.fourcc, 'fps': self.fps, 'frameSize': self.frameSize}
+            stream.release()
+
+        self.detector = detector
+        self.batchSize = batchSize
+        leftover = 0
+        if (self.datalen) % batchSize:
+            leftover = 1
+        self.num_batches = self.datalen // batchSize + leftover
+
+        self._input_size = cfg.DATA_PRESET.IMAGE_SIZE
+        self._output_size = cfg.DATA_PRESET.HEATMAP_SIZE
+
+        self._sigma = cfg.DATA_PRESET.SIGMA
+
+        pose_dataset = builder.retrieve_dataset(self.cfg.DATASET.TRAIN)
+        if cfg.DATA_PRESET.TYPE == 'simple':
+            self.transformation = SimpleTransform(
+                pose_dataset, scale_factor=0,
+                input_size=self._input_size,
+                output_size=self._output_size,
+                rot=0, sigma=self._sigma,
+                train=False, add_dpg=False, gpu_device=self.device)
+
+        # initialize the queue used to store data
+        """
+        image_queue: the buffer storing pre-processed images for object detection
+        det_queue: the buffer storing human detection results
+        pose_queue: the buffer storing post-processed cropped human image for pose estimation
+        """
+        if opt.sp:
+            self._stopped = False
+            self.image_queue = Queue(maxsize=queueSize)
+            self.det_queue = Queue(maxsize=10 * queueSize)
+            self.pose_queue = Queue(maxsize=10 * queueSize)
+        else:
+            self._stopped = mp.Value('b', False)
+            self.image_queue = mp.Queue(maxsize=queueSize)
+            self.det_queue = mp.Queue(maxsize=10 * queueSize)
+            self.pose_queue = mp.Queue(maxsize=10 * queueSize)
+
+    def start_worker(self, target):
+        if self.opt.sp:
+            p = Thread(target=target, args=())
+        else:
+            p = mp.Process(target=target, args=())
+        # p.daemon = True
+        p.start()
+        return p
+
+    def start(self):
+        # start a thread to pre process images for object detection
+        if self.mode == 'image':
+            image_preprocess_worker = self.start_worker(self.image_preprocess)
+        elif self.mode == 'video':
+            image_preprocess_worker = self.start_worker(self.frame_preprocess)
+        # start a thread to detect human in images
+        image_detection_worker = self.start_worker(self.image_detection)
+        # start a thread to post process cropped human image for pose estimation
+        image_postprocess_worker = self.start_worker(self.image_postprocess)
+
+        return [image_preprocess_worker, image_detection_worker, image_postprocess_worker]
+
+    def stop(self):
+        # clear queues
+        self.clear_queues()
+
+    def terminate(self):
+        if self.opt.sp:
+            self._stopped = True
+        else:
+            self._stopped.value = True
+        self.stop()
+
+    def clear_queues(self):
+        self.clear(self.image_queue)
+        self.clear(self.det_queue)
+        self.clear(self.pose_queue)
+
+    def clear(self, queue):
+        while not queue.empty():
+            queue.get()
+
+    def wait_and_put(self, queue, item):
+        queue.put(item)
+
+    def wait_and_get(self, queue):
+        return queue.get()
+
+    def image_preprocess(self):
+        for i in range(self.num_batches):
+            imgs = []
+            orig_imgs = []
+            im_names = []
+            im_dim_list = []
+            for k in range(i * self.batchSize, min((i + 1) * self.batchSize, self.datalen)):
+                if self.stopped:
+                    self.wait_and_put(self.image_queue, (None, None, None, None))
+                    return
+                im_name_k = self.imglist[k]
+
+                # expected image shape like (1,3,h,w) or (3,h,w)
+                img_k = self.detector.image_preprocess(im_name_k)
+                if isinstance(img_k, np.ndarray):
+                    img_k = torch.from_numpy(img_k)
+                # add one dimension at the front for batch if image shape (3,h,w)
+                if img_k.dim() == 3:
+                    img_k = img_k.unsqueeze(0)
+                orig_img_k = cv2.cvtColor(cv2.imread(im_name_k), cv2.COLOR_BGR2RGB) # scipy.misc.imread(im_name_k, mode='RGB') is depreciated
+                im_dim_list_k = orig_img_k.shape[1], orig_img_k.shape[0]
+
+                imgs.append(img_k)
+                orig_imgs.append(orig_img_k)
+                im_names.append(os.path.basename(im_name_k))
+                im_dim_list.append(im_dim_list_k)
+
+            with torch.no_grad():
+                # Human Detection
+                imgs = torch.cat(imgs)
+                im_dim_list = torch.FloatTensor(im_dim_list).repeat(1, 2)
+                # im_dim_list_ = im_dim_list
+
+            self.wait_and_put(self.image_queue, (imgs, orig_imgs, im_names, im_dim_list))
+
+    def frame_preprocess(self):
+        stream = cv2.VideoCapture(self.path)
+        assert stream.isOpened(), 'Cannot capture source'
+
+        for i in range(self.num_batches):
+            imgs = []
+            orig_imgs = []
+            im_names = []
+            im_dim_list = []
+            for k in range(i * self.batchSize, min((i + 1) * self.batchSize, self.datalen)):
+                (grabbed, frame) = stream.read()
+                # if the `grabbed` boolean is `False`, then we have
+                # reached the end of the video file
+                if not grabbed or self.stopped:
+                    # put the rest pre-processed data to the queue
+                    if len(imgs) > 0:
+                        with torch.no_grad():
+                            # Record original image resolution
+                            imgs = torch.cat(imgs)
+                            im_dim_list = torch.FloatTensor(im_dim_list).repeat(1, 2)
+                        self.wait_and_put(self.image_queue, (imgs, orig_imgs, im_names, im_dim_list))
+                    self.wait_and_put(self.image_queue, (None, None, None, None))
+                    print('===========================> This video get ' + str(k) + ' frames in total.')
+                    sys.stdout.flush()
+                    stream.release()
+                    return
+
+                # expected frame shape like (1,3,h,w) or (3,h,w)
+                img_k = self.detector.image_preprocess(frame)
+
+                if isinstance(img_k, np.ndarray):
+                    img_k = torch.from_numpy(img_k)
+                # add one dimension at the front for batch if image shape (3,h,w)
+                if img_k.dim() == 3:
+                    img_k = img_k.unsqueeze(0)
+
+                im_dim_list_k = frame.shape[1], frame.shape[0]
+
+                imgs.append(img_k)
+                orig_imgs.append(frame[:, :, ::-1])
+                im_names.append(str(k) + '.jpg')
+                im_dim_list.append(im_dim_list_k)
+
+            with torch.no_grad():
+                # Record original image resolution
+                imgs = torch.cat(imgs)
+                im_dim_list = torch.FloatTensor(im_dim_list).repeat(1, 2)
+                # im_dim_list_ = im_dim_list
+
+            self.wait_and_put(self.image_queue, (imgs, orig_imgs, im_names, im_dim_list))
+        stream.release()
+
+    def image_detection(self):
+        for i in range(self.num_batches):
+            imgs, orig_imgs, im_names, im_dim_list = self.wait_and_get(self.image_queue)
+            if imgs is None or self.stopped:
+                self.wait_and_put(self.det_queue, (None, None, None, None, None, None, None))
+                return
+
+            with torch.no_grad():
+                # pad useless images to fill a batch, else there will be a bug
+                for pad_i in range(self.batchSize - len(imgs)):
+                    imgs = torch.cat((imgs, torch.unsqueeze(imgs[0], dim=0)), 0)
+                    im_dim_list = torch.cat((im_dim_list, torch.unsqueeze(im_dim_list[0], dim=0)), 0)
+
+                dets = self.detector.images_detection(imgs, im_dim_list)
+                if isinstance(dets, int) or dets.shape[0] == 0:
+                    for k in range(len(orig_imgs)):
+                        self.wait_and_put(self.det_queue, (orig_imgs[k], im_names[k], None, None, None, None, None))
+                    continue
+                if isinstance(dets, np.ndarray):
+                    dets = torch.from_numpy(dets)
+                dets = dets.cpu()
+                boxes = dets[:, 1:5]
+                scores = dets[:, 5:6]
+                if self.opt.tracking:
+                    ids = dets[:, 6:7]
+                else:
+                    ids = torch.zeros(scores.shape)
+
+            for k in range(len(orig_imgs)):
+                boxes_k = boxes[dets[:, 0] == k]
+                if isinstance(boxes_k, int) or boxes_k.shape[0] == 0:
+                    self.wait_and_put(self.det_queue, (orig_imgs[k], im_names[k], None, None, None, None, None))
+                    continue
+                inps = torch.zeros(boxes_k.size(0), 3, *self._input_size)
+                cropped_boxes = torch.zeros(boxes_k.size(0), 4)
+
+                self.wait_and_put(self.det_queue, (orig_imgs[k], im_names[k], boxes_k, scores[dets[:, 0] == k], ids[dets[:, 0] == k], inps, cropped_boxes))
+
+    def image_postprocess(self):
+        for i in range(self.datalen):
+            with torch.no_grad():
+                (orig_img, im_name, boxes, scores, ids, inps, cropped_boxes) = self.wait_and_get(self.det_queue)
+                if orig_img is None or self.stopped:
+                    self.wait_and_put(self.pose_queue, (None, None, None, None, None, None, None))
+                    return
+                if boxes is None or boxes.nelement() == 0:
+                    self.wait_and_put(self.pose_queue, (None, orig_img, im_name, boxes, scores, ids, None))
+                    continue
+                # imght = orig_img.shape[0]
+                # imgwidth = orig_img.shape[1]
+                for i, box in enumerate(boxes):
+                    inps[i], cropped_box = self.transformation.test_transform(orig_img, box)
+                    cropped_boxes[i] = torch.FloatTensor(cropped_box)
+
+                # inps, cropped_boxes = self.transformation.align_transform(orig_img, boxes)
+
+                self.wait_and_put(self.pose_queue, (inps, orig_img, im_name, boxes, scores, ids, cropped_boxes))
+
+    def read(self):
+        return self.wait_and_get(self.pose_queue)
+
+    @property
+    def stopped(self):
+        if self.opt.sp:
+            return self._stopped
+        else:
+            return self._stopped.value
+
+    @property
+    def length(self):
+        return self.datalen

alphapose/utils/env.py

@@ -0,0 +1,81 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+
+import os
+import torch
+import torch.distributed as dist
+
+
+def init_dist(opt):
+    """Initialize distributed computing environment."""
+    opt.ngpus_per_node = torch.cuda.device_count()
+
+    torch.cuda.set_device(opt.gpu)
+
+    if opt.launcher == 'pytorch':
+        _init_dist_pytorch(opt)
+    elif opt.launcher == 'mpi':
+        _init_dist_mpi(opt)
+    elif opt.launcher == 'slurm':
+        _init_dist_slurm(opt)
+    else:
+        raise ValueError('Invalid launcher type: {}'.format(opt.launcher))
+
+
+def _init_dist_pytorch(opt, **kwargs):
+    """Set up environment."""
+    # TODO: use local_rank instead of rank % num_gpus
+    opt.rank = opt.rank * opt.ngpus_per_node + opt.gpu
+    opt.world_size = opt.world_size
+    dist.init_process_group(backend=opt.dist_backend, init_method=opt.dist_url,
+                            world_size=opt.world_size, rank=opt.rank)
+    print(f"{opt.dist_url}, ws:{opt.world_size}, rank:{opt.rank}")
+
+    if opt.rank % opt.ngpus_per_node == 0:
+        opt.log = True
+    else:
+        opt.log = False
+
+
+def _init_dist_slurm(opt, port=23348, **kwargs):
+    """Set up slurm environment."""
+    proc_id = int(os.environ['SLURM_PROCID'])
+    ntasks = int(os.environ['SLURM_NTASKS'])
+    node_list = os.environ['SLURM_NODELIST']
+    num_gpus = torch.cuda.device_count()
+    torch.cuda.set_device(proc_id % num_gpus)
+    if '[' in node_list:
+        beg = node_list.find('[')
+        pos1 = node_list.find('-', beg)
+        if pos1 < 0:
+            pos1 = 1000
+        pos2 = node_list.find(',', beg)
+        if pos2 < 0:
+            pos2 = 1000
+        node_list = node_list[:min(pos1, pos2)].replace('[', '')
+    addr = node_list[8:].replace('-', '.')
+    os.environ['MASTER_PORT'] = str(port)
+    os.environ['MASTER_ADDR'] = addr
+    os.environ['WORLD_SIZE'] = str(ntasks)
+    os.environ['RANK'] = str(proc_id)
+
+    opt.ngpus_per_node = num_gpus
+    opt.rank = int(proc_id)
+    opt.rank = proc_id * num_gpus + opt.gpu
+    opt.world_size = int(ntasks) * num_gpus
+
+    print(f"tcp://{node_list}:{port}, ws:{opt.world_size}, rank:{opt.rank}, proc_id:{proc_id}")
+    dist.init_process_group(backend=opt.dist_backend,
+                            init_method=f'tcp://{node_list}:{port}',
+                            world_size=opt.world_size,
+                            rank=opt.rank)
+    if opt.rank == 0:
+        opt.log = True
+    else:
+        opt.log = False
+
+
+def _init_dist_mpi(backend, **kwargs):
+    raise NotImplementedError

alphapose/utils/file_detector.py

@@ -0,0 +1,156 @@
+from itertools import count
+from threading import Thread
+from queue import Queue
+import json
+
+import cv2
+import numpy as np
+
+import torch
+import torch.multiprocessing as mp
+
+from alphapose.utils.presets import SimpleTransform
+
+
+class FileDetectionLoader():
+    def __init__(self, input_source, cfg, opt, queueSize=128):
+        self.cfg = cfg
+        self.opt = opt
+        self.bbox_file = input_source
+
+        self._input_size = cfg.DATA_PRESET.IMAGE_SIZE
+        self._output_size = cfg.DATA_PRESET.HEATMAP_SIZE
+
+        self._sigma = cfg.DATA_PRESET.SIGMA
+
+        if cfg.DATA_PRESET.TYPE == 'simple':
+            self.transformation = SimpleTransform(
+                self, scale_factor=0,
+                input_size=self._input_size,
+                output_size=self._output_size,
+                rot=0, sigma=self._sigma,
+                train=False, add_dpg=False)
+
+        # initialize the det file list        
+        boxes = None
+        if isinstance(self.bbox_file,list):
+            boxes = self.bbox_file
+        else:
+            with open(self.bbox_file, 'r') as f:
+                boxes = json.load(f)
+            assert boxes is not None, 'Load %s fail!' % self.bbox_file
+
+        self.all_imgs = []
+        self.all_boxes = {}
+        self.all_scores = {}
+        self.all_ids = {}
+        num_boxes = 0
+        for k_img in range(0, len(boxes)):
+            det_res = boxes[k_img]
+            img_name = det_res['image_id']
+            if img_name not in self.all_imgs:
+                self.all_imgs.append(img_name)
+                self.all_boxes[img_name] = []
+                self.all_scores[img_name] = []
+                self.all_ids[img_name] = []
+            x1, y1, w, h = det_res['bbox']
+            bbox = [x1, y1, x1 + w, y1 + h]
+            score = det_res['score']
+            self.all_boxes[img_name].append(bbox)
+            self.all_scores[img_name].append(score)
+            if 'idx' in det_res.keys():
+                self.all_ids[img_name].append(int(det_res['idx']))
+            else:
+                self.all_ids[img_name].append(0)
+
+        # initialize the queue used to store data
+        """
+        pose_queue: the buffer storing post-processed cropped human image for pose estimation
+        """
+        if opt.sp:
+            self._stopped = False
+            self.pose_queue = Queue(maxsize=queueSize)
+        else:
+            self._stopped = mp.Value('b', False)
+            self.pose_queue = mp.Queue(maxsize=queueSize)
+
+    def start_worker(self, target):
+        if self.opt.sp:
+            p = Thread(target=target, args=())
+        else:
+            p = mp.Process(target=target, args=())
+        # p.daemon = True
+        p.start()
+        return p
+
+    def start(self):
+        # start a thread to pre process images for object detection
+        image_preprocess_worker = self.start_worker(self.get_detection)
+        return [image_preprocess_worker]
+
+    def stop(self):
+        # clear queues
+        self.clear_queues()
+
+    def terminate(self):
+        if self.opt.sp:
+            self._stopped = True
+        else:
+            self._stopped.value = True
+        self.stop()
+
+    def clear_queues(self):
+        self.clear(self.pose_queue)
+
+    def clear(self, queue):
+        while not queue.empty():
+            queue.get()
+
+    def wait_and_put(self, queue, item):
+        if not self.stopped:
+            queue.put(item)
+
+    def wait_and_get(self, queue):
+        if not self.stopped:
+            return queue.get()
+
+    def get_detection(self):
+        
+        for im_name_k in self.all_imgs:
+            boxes = torch.from_numpy(np.array(self.all_boxes[im_name_k]))
+            scores = torch.from_numpy(np.array(self.all_scores[im_name_k]))
+            ids = torch.from_numpy(np.array(self.all_ids[im_name_k]))
+            orig_img_k = cv2.cvtColor(cv2.imread(im_name_k), cv2.COLOR_BGR2RGB) #scipy.misc.imread(im_name_k, mode='RGB') is depreciated
+
+
+            inps = torch.zeros(boxes.size(0), 3, *self._input_size)
+            cropped_boxes = torch.zeros(boxes.size(0), 4)
+            for i, box in enumerate(boxes):
+                inps[i], cropped_box = self.transformation.test_transform(orig_img_k, box)
+                cropped_boxes[i] = torch.FloatTensor(cropped_box)
+
+            
+            self.wait_and_put(self.pose_queue, (inps, orig_img_k, im_name_k, boxes, scores, ids, cropped_boxes))
+        
+        self.wait_and_put(self.pose_queue, (None, None, None, None, None, None, None))
+        return
+        
+    def read(self):
+        return self.wait_and_get(self.pose_queue)
+
+    @property
+    def stopped(self):
+        if self.opt.sp:
+            return self._stopped
+        else:
+            return self._stopped.value
+    @property
+    def length(self):
+        return len(self.all_imgs)
+
+    @property
+    def joint_pairs(self):
+        """Joint pairs which defines the pairs of joint to be swapped
+        when the image is flipped horizontally."""
+        return [[1, 2], [3, 4], [5, 6], [7, 8],
+                [9, 10], [11, 12], [13, 14], [15, 16]]

alphapose/utils/logger.py

@@ -0,0 +1,29 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+
+import torch
+import torch.nn.functional as F
+
+
+def board_writing(writer, loss, acc, iterations, dataset='Train'):
+    writer.add_scalar(
+        '{}/Loss'.format(dataset), loss, iterations)
+    writer.add_scalar(
+        '{}/acc'.format(dataset), acc, iterations)
+
+
+def debug_writing(writer, outputs, labels, inputs, iterations):
+    tmp_tar = torch.unsqueeze(labels.cpu().data[0], dim=1)
+    # tmp_out = torch.unsqueeze(outputs.cpu().data[0], dim=1)
+
+    tmp_inp = inputs.cpu().data[0]
+    tmp_inp[0] += 0.406
+    tmp_inp[1] += 0.457
+    tmp_inp[2] += 0.480
+
+    tmp_inp[0] += torch.sum(F.interpolate(tmp_tar, scale_factor=4, mode='bilinear'), dim=0)[0]
+    tmp_inp.clamp_(0, 1)
+
+    writer.add_image('Data/input', tmp_inp, iterations)

alphapose/utils/metrics.py

@@ -0,0 +1,252 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com), Haoyi Zhu
+# -----------------------------------------------------
+
+import os
+import sys
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from .transforms import get_max_pred_batch, _integral_tensor
+
+from pycocotools.coco import COCO
+from pycocotools.cocoeval import COCOeval
+
+
+class DataLogger(object):
+    """Average data logger."""
+    def __init__(self):
+        self.clear()
+
+    def clear(self):
+        self.value = 0
+        self.sum = 0
+        self.cnt = 0
+        self.avg = 0
+
+    def update(self, value, n=1):
+        self.value = value
+        self.sum += value * n
+        self.cnt += n
+        self._cal_avg()
+
+    def _cal_avg(self):
+        self.avg = self.sum / self.cnt
+
+
+def calc_iou(pred, target):
+    """Calculate mask iou"""
+    if isinstance(pred, torch.Tensor):
+        pred = pred.cpu().data.numpy()
+    if isinstance(target, torch.Tensor):
+        target = target.cpu().data.numpy()
+
+    pred = pred >= 0.5
+    target = target >= 0.5
+
+    intersect = (pred == target) * pred * target
+    union = np.maximum(pred, target)
+
+    if pred.ndim == 2:
+        iou = np.sum(intersect) / np.sum(union)
+    elif pred.ndim == 3 or pred.ndim == 4:
+        n_samples = pred.shape[0]
+        intersect = intersect.reshape(n_samples, -1)
+        union = union.reshape(n_samples, -1)
+
+        iou = np.mean(np.sum(intersect, axis=1) / np.sum(union, axis=1))
+
+    return iou
+
+
+def mask_cross_entropy(pred, target):
+    return F.binary_cross_entropy_with_logits(
+        pred, target, reduction='mean')[None]
+
+
+def evaluate_mAP(res_file, ann_type='bbox', ann_file='./data/coco/annotations/person_keypoints_val2017.json', silence=True):
+    """Evaluate mAP result for coco dataset.
+
+    Parameters
+    ----------
+    res_file: str
+        Path to result json file.
+    ann_type: str
+        annotation type, including: `bbox`, `segm`, `keypoints`.
+    ann_file: str
+        Path to groundtruth file.
+    silence: bool
+        True: disable running log.
+
+    """
+    class NullWriter(object):
+        def write(self, arg):
+            pass
+
+    # ann_file = os.path.join('./data/coco/annotations/', ann_file)
+
+    if silence:
+        nullwrite = NullWriter()
+        oldstdout = sys.stdout
+        sys.stdout = nullwrite  # disable output
+
+    cocoGt = COCO(ann_file)
+    cocoDt = cocoGt.loadRes(res_file)
+
+    cocoEval = COCOeval(cocoGt, cocoDt, ann_type)
+    cocoEval.evaluate()
+    cocoEval.accumulate()
+    cocoEval.summarize()
+
+    if silence:
+        sys.stdout = oldstdout  # enable output
+    
+    if isinstance(cocoEval.stats[0], dict):
+        stats_names = ['AP', 'Ap .5', 'AP .75', 'AP (M)', 'AP (L)',
+                       'AR', 'AR .5', 'AR .75', 'AR (M)', 'AR (L)']
+        parts = ['body', 'face', 'hand', 'fullbody']
+
+        info = {}
+        for i, part in enumerate(parts):
+            info[part] = cocoEval.stats[i][part][0]
+        return info
+    else:
+        stats_names = ['AP', 'Ap .5', 'AP .75', 'AP (M)', 'AP (L)',
+                       'AR', 'AR .5', 'AR .75', 'AR (M)', 'AR (L)']
+        info_str = {}
+        for ind, name in enumerate(stats_names):
+            info_str[name] = cocoEval.stats[ind]
+        return info_str['AP']
+
+
+def calc_accuracy(preds, labels):
+    """Calculate heatmap accuracy."""
+    preds = preds.cpu().data.numpy()
+    labels = labels.cpu().data.numpy()
+
+    num_joints = preds.shape[1]
+
+    norm = 1.0
+    hm_h = preds.shape[2]
+    hm_w = preds.shape[3]
+
+    preds, _ = get_max_pred_batch(preds)
+    labels, _ = get_max_pred_batch(labels)
+    norm = np.ones((preds.shape[0], 2)) * np.array([hm_w, hm_h]) / 10
+
+    dists = calc_dist(preds, labels, norm)
+
+    acc = 0
+    sum_acc = 0
+    cnt = 0
+    for i in range(num_joints):
+        acc = dist_acc(dists[i])
+        if acc >= 0:
+            sum_acc += acc
+            cnt += 1
+
+    if cnt > 0:
+        return sum_acc / cnt
+    else:
+        return 0
+
+
+def calc_integral_accuracy(preds, labels, label_masks, output_3d=False, norm_type='softmax'):
+    """Calculate integral coordinates accuracy."""
+    def integral_op(hm_1d):
+        hm_1d = hm_1d * torch.cuda.comm.broadcast(torch.arange(hm_1d.shape[-1]).type(
+            torch.cuda.FloatTensor), devices=[hm_1d.device.index])[0]
+        return hm_1d
+
+    preds = preds.detach()
+    hm_width = preds.shape[-1]
+    hm_height = preds.shape[-2]
+
+    if output_3d:
+        hm_depth = hm_height
+        num_joints = preds.shape[1] // hm_depth
+    else:
+        hm_depth = 1
+        num_joints = preds.shape[1]
+
+    with torch.no_grad():
+        pred_jts, _ = _integral_tensor(preds, num_joints, output_3d, hm_width, hm_height, hm_depth, integral_op, norm_type=norm_type)
+
+    coords = pred_jts.detach().cpu().numpy()
+    coords = coords.astype(float)
+    if output_3d:
+        coords = coords.reshape((coords.shape[0], int(coords.shape[1] / 3), 3))
+    else:
+        coords = coords.reshape((coords.shape[0], int(coords.shape[1] / 2), 2))
+    coords[:, :, 0] = (coords[:, :, 0] + 0.5) * hm_width
+    coords[:, :, 1] = (coords[:, :, 1] + 0.5) * hm_height
+
+    if output_3d:
+        labels = labels.cpu().data.numpy().reshape(preds.shape[0], num_joints, 3)
+        label_masks = label_masks.cpu().data.numpy().reshape(preds.shape[0], num_joints, 3)
+
+        labels[:, :, 0] = (labels[:, :, 0] + 0.5) * hm_width
+        labels[:, :, 1] = (labels[:, :, 1] + 0.5) * hm_height
+        labels[:, :, 2] = (labels[:, :, 2] + 0.5) * hm_depth
+
+        coords[:, :, 2] = (coords[:, :, 2] + 0.5) * hm_depth
+    else:
+        labels = labels.cpu().data.numpy().reshape(preds.shape[0], num_joints, 2)
+        label_masks = label_masks.cpu().data.numpy().reshape(preds.shape[0], num_joints, 2)
+
+        labels[:, :, 0] = (labels[:, :, 0] + 0.5) * hm_width
+        labels[:, :, 1] = (labels[:, :, 1] + 0.5) * hm_height
+
+    coords = coords * label_masks
+    labels = labels * label_masks
+
+    if output_3d:
+        norm = np.ones((preds.shape[0], 3)) * np.array([hm_width, hm_height, hm_depth]) / 10
+    else:
+        norm = np.ones((preds.shape[0], 2)) * np.array([hm_width, hm_height]) / 10
+
+    dists = calc_dist(coords, labels, norm)
+
+    acc = 0
+    sum_acc = 0
+    cnt = 0
+    for i in range(num_joints):
+        acc = dist_acc(dists[i])
+        if acc >= 0:
+            sum_acc += acc
+            cnt += 1
+
+    if cnt > 0:
+        return sum_acc / cnt
+    else:
+        return 0
+
+
+def calc_dist(preds, target, normalize):
+    """Calculate normalized distances"""
+    preds = preds.astype(np.float32)
+    target = target.astype(np.float32)
+    dists = np.zeros((preds.shape[1], preds.shape[0]))
+
+    for n in range(preds.shape[0]):
+        for c in range(preds.shape[1]):
+            if target[n, c, 0] > 1 and target[n, c, 1] > 1:
+                normed_preds = preds[n, c, :] / normalize[n]
+                normed_targets = target[n, c, :] / normalize[n]
+                dists[c, n] = np.linalg.norm(normed_preds - normed_targets)
+            else:
+                dists[c, n] = -1
+
+    return dists
+
+
+def dist_acc(dists, thr=0.5):
+    """Calculate accuracy with given input distance."""
+    dist_cal = np.not_equal(dists, -1)
+    num_dist_cal = dist_cal.sum()
+    if num_dist_cal > 0:
+        return np.less(dists[dist_cal], thr).sum() * 1.0 / num_dist_cal
+    else:
+        return -1

alphapose/utils/pPose_nms.py

@@ -0,0 +1,578 @@
+# -*- coding: utf-8 -*-
+import json
+import os
+import zipfile
+import time
+from multiprocessing.dummy import Pool as ThreadPool
+from collections import defaultdict
+
+import torch
+import numpy as np
+
+''' Constant Configuration '''
+delta1 = 1
+mu = 1.7
+delta2 = 2.65
+gamma = 22.48
+scoreThreds = 0.3
+matchThreds = 5
+alpha = 0.1
+vis_thr = 0.2
+oks_thr = 0.9
+#pool = ThreadPool(4)
+
+
+def oks_pose_nms(data, soft=False):
+    kpts = defaultdict(list)
+    post_data = []
+
+    for item in data:
+        img_id = item['image_id']
+        kpts[img_id].append(item)
+
+    for img_id, img_res in kpts.items():
+        for n_p in img_res:
+            box_score = n_p['score']
+            kpt_score = 0
+            valid_num = 0
+            kpt = np.array(n_p['keypoints']).reshape(-1, 3)
+            for n_jt in range(kpt.shape[0]):
+                t_s = kpt[n_jt][2]
+                if t_s > vis_thr:
+                    kpt_score += t_s
+                    valid_num += 1
+            if valid_num != 0:
+                kpt_score = kpt_score / valid_num
+            n_p['score'] = kpt_score * box_score
+
+        if soft:
+            keep = soft_oks_nms(
+                [img_res[i] for i in range(len(img_res))], oks_thr)
+        else:
+            keep = oks_nms(
+                [img_res[i] for i in range(len(img_res))], oks_thr)
+
+        if len(keep) == 0:
+            post_data += img_res
+        else:
+            post_data += [img_res[_keep] for _keep in keep]
+
+    return post_data
+
+
+def oks_nms(kpts_db, thr, sigmas=None, vis_thr=None):
+    """OKS NMS implementations.
+    Args:
+        kpts_db: keypoints.
+        thr: Retain overlap < thr.
+        sigmas: standard deviation of keypoint labelling.
+        vis_thr: threshold of the keypoint visibility.
+    Returns:
+        np.ndarray: indexes to keep.
+    """
+    if len(kpts_db) == 0:
+        return []
+
+    scores = np.array([k['score'] for k in kpts_db])
+    #kpts = np.array([k['keypoints'].flatten() for k in kpts_db])
+    kpts = np.array([k['keypoints'] for k in kpts_db])
+    areas = np.array([k['area'] for k in kpts_db])
+
+    order = scores.argsort()[::-1]
+
+    keep = []
+    while len(order) > 0:
+        i = order[0]
+        keep.append(i)
+
+        oks_ovr = oks_iou(kpts[i], kpts[order[1:]], areas[i], areas[order[1:]],
+                          sigmas, vis_thr)
+
+        inds = np.where(oks_ovr <= thr)[0]
+        order = order[inds + 1]
+
+    keep = np.array(keep)
+
+    return keep
+
+
+def soft_oks_nms(kpts_db, thr, max_dets=20, sigmas=None, vis_thr=None):
+    """Soft OKS NMS implementations.
+    Args:
+        kpts_db
+        thr: retain oks overlap < thr.
+        max_dets: max number of detections to keep.
+        sigmas: Keypoint labelling uncertainty.
+    Returns:
+        np.ndarray: indexes to keep.
+    """
+    if len(kpts_db) == 0:
+        return []
+
+    scores = np.array([k['score'] for k in kpts_db])
+    kpts = np.array([k['keypoints'].flatten() for k in kpts_db])
+    areas = np.array([k['area'] for k in kpts_db])
+
+    order = scores.argsort()[::-1]
+    scores = scores[order]
+
+    keep = np.zeros(max_dets, dtype=np.intp)
+    keep_cnt = 0
+    while len(order) > 0 and keep_cnt < max_dets:
+        i = order[0]
+
+        oks_ovr = oks_iou(kpts[i], kpts[order[1:]], areas[i], areas[order[1:]],
+                          sigmas, vis_thr)
+
+        order = order[1:]
+        scores = _rescore(oks_ovr, scores[1:], thr)
+
+        tmp = scores.argsort()[::-1]
+        order = order[tmp]
+        scores = scores[tmp]
+
+        keep[keep_cnt] = i
+        keep_cnt += 1
+
+    keep = keep[:keep_cnt]
+
+    return keep
+
+
+def oks_iou(g, d, a_g, a_d, sigmas=None, vis_thr=None):
+    """Calculate oks ious.
+    Args:
+        g: Ground truth keypoints.
+        d: Detected keypoints.
+        a_g: Area of the ground truth object.
+        a_d: Area of the detected object.
+        sigmas: standard deviation of keypoint labelling.
+        vis_thr: threshold of the keypoint visibility.
+    Returns:
+        list: The oks ious.
+    """
+    if sigmas is None:
+        if len(g) == 408:   # 136keypoints
+            sigmas = np.array([.26, .25, .25, .35, .35, .79, .79, .72, .72, .62,.62, 1.07, 1.07, .87, .87, .89, .89, .8,.8,.8,.89, .89, .89, .89, .89, .89,
+                     .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25,
+                     .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25,
+                     .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25,
+                     .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25,
+                     .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25, .25])/10.0
+        elif len(g) == 399:
+            sigmas = np.array([.026, .025, .025, .035, .035, .079, .079, .072, .072, .062, .062, 0.107, 0.107, .087, .087, .089, .089,
+                0.068, 0.066, 0.066, 0.092, 0.094, 0.094,
+                0.042, 0.043, 0.044, 0.043, 0.040, 0.035, 0.031, 0.025, 0.020, 0.023, 0.029, 0.032, 0.037, 0.038, 0.043,
+                0.041, 0.045, 0.013, 0.012, 0.011, 0.011, 0.012, 0.012, 0.011, 0.011, 0.013, 0.015, 0.009, 0.007, 0.007,
+                0.007, 0.012, 0.009, 0.008, 0.016, 0.010, 0.017, 0.011, 0.009, 0.011, 0.009, 0.007, 0.013, 0.008, 0.011,
+                0.012, 0.010, 0.034, 0.008, 0.008, 0.009, 0.008, 0.008, 0.007, 0.010, 0.008, 0.009, 0.009, 0.009, 0.007,
+                0.007, 0.008, 0.011, 0.008, 0.008, 0.008, 0.01, 0.008,
+                0.029, 0.022, 0.035, 0.037, 0.047, 0.026, 0.025, 0.024, 0.035, 0.018, 0.024, 0.022, 0.026, 0.017,
+                0.021, 0.021, 0.032, 0.02, 0.019, 0.022, 0.031,
+                0.029, 0.022, 0.035, 0.037, 0.047, 0.026, 0.025, 0.024, 0.035, 0.018, 0.024, 0.022, 0.026, 0.017,
+                0.021, 0.021, 0.032, 0.02, 0.019, 0.022, 0.031])
+        elif len(g) == 78:
+            sigmas = np.array([.26, .25, .25, .35, .35, .79, .79, .72, .72, .62,.62, 1.07, 1.07, .87, .87, .89, .89, .8,.8,.8,.89, .89, .89, .89, .89, .89])/10.0
+        else:
+            sigmas = np.array([
+                .26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, 1.07, 1.07,
+                .87, .87, .89, .89
+            ]) / 10.0
+    vars = (sigmas * 2)**2
+    xg = g[0::3]
+    yg = g[1::3]
+    vg = g[2::3]
+    ious = np.zeros(len(d))
+    for n_d in range(0, len(d)):
+        xd = d[n_d, 0::3]
+        yd = d[n_d, 1::3]
+        vd = d[n_d, 2::3]
+        dx = xd - xg
+        dy = yd - yg
+        e = (dx**2 + dy**2) / vars / ((a_g + a_d[n_d]) / 2 + np.spacing(1)) / 2
+        if vis_thr is not None:
+            ind = list(vg > vis_thr) and list(vd > vis_thr)
+            e = e[ind]
+        ious[n_d] = np.sum(np.exp(-e)) / len(e) if len(e) != 0 else 0.0
+    return ious
+
+
+def _rescore(overlap, scores, thr, type='gaussian'):
+    """Rescoring mechanism gaussian or linear.
+    Args:
+        overlap: calculated ious
+        scores: target scores.
+        thr: retain oks overlap < thr.
+        type: 'gaussian' or 'linear'
+    Returns:
+        np.ndarray: indexes to keep
+    """
+    assert len(overlap) == len(scores)
+    assert type in ['gaussian', 'linear']
+
+    if type == 'linear':
+        inds = np.where(overlap >= thr)[0]
+        scores[inds] = scores[inds] * (1 - overlap[inds])
+    else:
+        scores = scores * np.exp(-overlap**2 / thr)
+
+    return scores
+
+
+def pose_nms(bboxes, bbox_scores, bbox_ids, pose_preds, pose_scores, areaThres=0):
+    '''
+    Parametric Pose NMS algorithm
+    bboxes:         bbox locations list (n, 4)
+    bbox_scores:    bbox scores list (n, 1)
+    bbox_ids:       bbox tracking ids list (n, 1)
+    pose_preds:     pose locations list (n, kp_num, 2)
+    pose_scores:    pose scores list    (n, kp_num, 1)
+    '''
+    #global ori_pose_preds, ori_pose_scores, ref_dists
+
+    pose_scores[pose_scores == 0] = 1e-5
+    kp_nums = pose_preds.size()[1]
+    res_bboxes, res_bbox_scores, res_bbox_ids, res_pose_preds, res_pose_scores, res_pick_ids = [],[],[],[],[],[]
+    
+    ori_bboxes = bboxes.clone()
+    ori_bbox_scores = bbox_scores.clone()
+    ori_bbox_ids = bbox_ids.clone()
+    ori_pose_preds = pose_preds.clone()
+    ori_pose_scores = pose_scores.clone()
+
+    xmax = bboxes[:, 2]
+    xmin = bboxes[:, 0]
+    ymax = bboxes[:, 3]
+    ymin = bboxes[:, 1]
+
+    widths = xmax - xmin
+    heights = ymax - ymin
+    ref_dists = alpha * np.maximum(widths, heights)
+
+    nsamples = bboxes.shape[0]
+    human_scores = pose_scores.mean(dim=1)
+
+    human_ids = np.arange(nsamples)
+    mask = np.ones(len(human_ids)).astype(bool)
+    
+    # Do pPose-NMS
+    pick = []
+    merge_ids = []
+    while(mask.any()):
+        tensor_mask = torch.Tensor(mask)==True
+        # Pick the one with highest score
+        pick_id = torch.argmax(human_scores[tensor_mask])
+        pick.append(human_ids[mask][pick_id])
+
+        # Get numbers of match keypoints by calling PCK_match
+        ref_dist = ref_dists[human_ids[mask][pick_id]]
+        simi = get_parametric_distance(pick_id, pose_preds[tensor_mask], pose_scores[tensor_mask], ref_dist)
+        num_match_keypoints = PCK_match(pose_preds[tensor_mask][pick_id], pose_preds[tensor_mask], ref_dist)
+
+        # Delete humans who have more than matchThreds keypoints overlap and high similarity
+        delete_ids = torch.from_numpy(np.arange(human_scores[tensor_mask].shape[0]))[((simi > gamma) | (num_match_keypoints >= matchThreds))]
+
+        if delete_ids.shape[0] == 0:
+            delete_ids = pick_id
+
+        merge_ids.append(human_ids[mask][delete_ids])
+        newmask = mask[mask]
+        newmask[delete_ids] = False
+        mask[mask] = newmask
+
+
+    assert len(merge_ids) == len(pick)
+    preds_pick = ori_pose_preds[pick]
+    scores_pick = ori_pose_scores[pick]
+    bbox_scores_pick = ori_bbox_scores[pick]
+    bboxes_pick = ori_bboxes[pick]
+    bbox_ids_pick = ori_bbox_ids[pick]
+    #final_result = pool.map(filter_result, zip(scores_pick, merge_ids, preds_pick, pick, bbox_scores_pick))
+    #final_result = [item for item in final_result if item is not None]
+
+    for j in range(len(pick)):
+        ids = np.arange(kp_nums)
+        max_score = torch.max(scores_pick[j, ids, 0])
+
+        if max_score < scoreThreds:
+            continue
+
+        # Merge poses
+        merge_id = merge_ids[j]
+        merge_pose, merge_score = p_merge_fast(
+            preds_pick[j], ori_pose_preds[merge_id], ori_pose_scores[merge_id], ref_dists[pick[j]])
+
+        max_score = torch.max(merge_score[ids])
+        if max_score < scoreThreds:
+            continue
+
+        xmax = max(merge_pose[:, 0])
+        xmin = min(merge_pose[:, 0])
+        ymax = max(merge_pose[:, 1])
+        ymin = min(merge_pose[:, 1])
+        bbox = bboxes_pick[j].cpu().tolist()
+        bbox_score = bbox_scores_pick[j].cpu()
+
+        if (1.5 ** 2 * (xmax - xmin) * (ymax - ymin) < areaThres):
+            continue
+
+
+        res_bboxes.append(bbox)
+        res_bbox_scores.append(bbox_score)
+        res_bbox_ids.append(ori_bbox_ids[merge_id].tolist())
+        res_pose_preds.append(merge_pose)
+        res_pose_scores.append(merge_score)
+        res_pick_ids.append(pick[j])
+
+ 
+
+    return res_bboxes, res_bbox_scores, res_bbox_ids, res_pose_preds, res_pose_scores, res_pick_ids
+
+
+def filter_result(args):
+    score_pick, merge_id, pred_pick, pick, bbox_score_pick = args
+    global ori_pose_preds, ori_pose_scores, ref_dists
+    kp_nums = ori_pose_preds.size()[1]
+    ids = np.arange(kp_nums)
+    max_score = torch.max(score_pick[ids, 0])
+
+    if max_score < scoreThreds:
+        return None
+
+    # Merge poses
+    merge_pose, merge_score = p_merge_fast(
+        pred_pick, ori_pose_preds[merge_id], ori_pose_scores[merge_id], ref_dists[pick])
+
+    max_score = torch.max(merge_score[ids])
+    if max_score < scoreThreds:
+        return None
+
+    xmax = max(merge_pose[:, 0])
+    xmin = min(merge_pose[:, 0])
+    ymax = max(merge_pose[:, 1])
+    ymin = min(merge_pose[:, 1])
+
+    if (1.5 ** 2 * (xmax - xmin) * (ymax - ymin) < 40 * 40.5):
+        return None
+
+    return {
+        'keypoints': merge_pose - 0.3,
+        'kp_score': merge_score,
+        'proposal_score': torch.mean(merge_score) + bbox_score_pick + 1.25 * max(merge_score)
+    }
+
+
+def p_merge(ref_pose, cluster_preds, cluster_scores, ref_dist):
+    '''
+    Score-weighted pose merging
+    INPUT:
+        ref_pose:       reference pose          -- [kp_num, 2]
+        cluster_preds:  redundant poses         -- [n, kp_num, 2]
+        cluster_scores: redundant poses score   -- [n, kp_num, 1]
+        ref_dist:       reference scale         -- Constant
+    OUTPUT:
+        final_pose:     merged pose             -- [kp_num, 2]
+        final_score:    merged score            -- [kp_num]
+    '''
+    dist = torch.sqrt(torch.sum(
+        torch.pow(ref_pose[np.newaxis, :] - cluster_preds, 2),
+        dim=2
+    ))  # [n, kp_num]
+
+    kp_num = ref_pose.size()[0]
+    ref_dist = min(ref_dist, 15)
+
+    mask = (dist <= ref_dist)
+    final_pose = torch.zeros(kp_num, 2)
+    final_score = torch.zeros(kp_num)
+
+    if cluster_preds.dim() == 2:
+        cluster_preds.unsqueeze_(0)
+        cluster_scores.unsqueeze_(0)
+    if mask.dim() == 1:
+        mask.unsqueeze_(0)
+
+    for i in range(kp_num):
+        cluster_joint_scores = cluster_scores[:, i][mask[:, i]]  # [k, 1]
+        cluster_joint_location = cluster_preds[:, i, :][mask[:, i].unsqueeze(
+            -1).repeat(1, 2)].view((torch.sum(mask[:, i]), -1))
+
+        # Get an normalized score
+        normed_scores = cluster_joint_scores / torch.sum(cluster_joint_scores)
+
+        # Merge poses by a weighted sum
+        final_pose[i, 0] = torch.dot(cluster_joint_location[:, 0], normed_scores.squeeze(-1))
+        final_pose[i, 1] = torch.dot(cluster_joint_location[:, 1], normed_scores.squeeze(-1))
+
+        final_score[i] = torch.dot(cluster_joint_scores.transpose(0, 1).squeeze(0), normed_scores.squeeze(-1))
+
+    return final_pose, final_score
+
+
+def p_merge_fast(ref_pose, cluster_preds, cluster_scores, ref_dist):
+    '''
+    Score-weighted pose merging
+    INPUT:
+        ref_pose:       reference pose          -- [kp_num, 2]
+        cluster_preds:  redundant poses         -- [n, kp_num, 2]
+        cluster_scores: redundant poses score   -- [n, kp_num, 1]
+        ref_dist:       reference scale         -- Constant
+    OUTPUT:
+        final_pose:     merged pose             -- [kp_num, 2]
+        final_score:    merged score            -- [kp_num]
+    '''
+    dist = torch.sqrt(torch.sum(
+        torch.pow(ref_pose[np.newaxis, :] - cluster_preds, 2),
+        dim=2
+    ))
+
+    kp_num = ref_pose.size()[0]
+    ref_dist = min(ref_dist, 15)
+
+    mask = (dist <= ref_dist)
+    final_pose = torch.zeros(kp_num, 2)
+    final_score = torch.zeros(kp_num)
+
+    if cluster_preds.dim() == 2:
+        cluster_preds.unsqueeze_(0)
+        cluster_scores.unsqueeze_(0)
+    if mask.dim() == 1:
+        mask.unsqueeze_(0)
+
+    # Weighted Merge
+    masked_scores = cluster_scores.mul(mask.float().unsqueeze(-1))
+    normed_scores = masked_scores / torch.sum(masked_scores, dim=0)
+
+    final_pose = torch.mul(cluster_preds, normed_scores.repeat(1, 1, 2)).sum(dim=0)
+    final_score = torch.mul(masked_scores, normed_scores).sum(dim=0)
+    return final_pose, final_score
+
+
+def get_parametric_distance(i, all_preds, keypoint_scores, ref_dist):
+    pick_preds = all_preds[i]
+    pred_scores = keypoint_scores[i]
+    dist = torch.sqrt(torch.sum(
+        torch.pow(pick_preds[np.newaxis, :] - all_preds, 2),
+        dim=2
+    ))
+    mask = (dist <= 1)
+
+    kp_nums = all_preds.size()[1]
+    # Define a keypoints distance
+    score_dists = torch.zeros(all_preds.shape[0], kp_nums)
+    keypoint_scores.squeeze_()
+    if keypoint_scores.dim() == 1:
+        keypoint_scores.unsqueeze_(0)
+    if pred_scores.dim() == 1:
+        pred_scores.unsqueeze_(1)
+    # The predicted scores are repeated up to do broadcast
+    pred_scores = pred_scores.repeat(1, all_preds.shape[0]).transpose(0, 1)
+
+    score_dists[mask] = torch.tanh(pred_scores[mask] / delta1) * torch.tanh(keypoint_scores[mask] / delta1)
+
+    point_dist = torch.exp((-1) * dist / delta2)
+    final_dist = torch.sum(score_dists, dim=1) + mu * torch.sum(point_dist, dim=1)
+
+    return final_dist
+
+
+def PCK_match(pick_pred, all_preds, ref_dist):
+    dist = torch.sqrt(torch.sum(
+        torch.pow(pick_pred[np.newaxis, :] - all_preds, 2),
+        dim=2
+    ))
+    ref_dist = min(ref_dist, 7)
+    num_match_keypoints = torch.sum(
+        dist / ref_dist <= 1,
+        dim=1
+    )
+
+    return num_match_keypoints
+
+
+def write_json(all_results, outputpath, save_file_name='alphapose-results', form=None, for_eval=False):
+    '''
+    all_result: result dict of predictions
+    outputpath: output directory
+    '''
+    json_results = []
+    json_results_cmu = {}
+    for im_res in all_results:
+        im_name = im_res['imgname']
+        for human in im_res['result']:
+            keypoints = []
+            result = {}
+            if for_eval:
+                result['image_id'] = int(os.path.basename(im_name).split('.')[0].split('_')[-1])
+            else:
+                result['image_id'] = os.path.basename(im_name)
+            result['category_id'] = 1
+
+            kp_preds = human['keypoints']
+            kp_scores = human['kp_score']
+            pro_scores = human['proposal_score']
+            for n in range(kp_scores.shape[0]):
+                keypoints.append(float(kp_preds[n, 0]))
+                keypoints.append(float(kp_preds[n, 1]))
+                keypoints.append(float(kp_scores[n]))
+            result['keypoints'] = keypoints
+            result['score'] = float(pro_scores)
+            if 'box' in human.keys():
+                result['box'] = human['box']
+            #pose track results by PoseFlow
+            if 'idx' in human.keys():
+                result['idx'] = human['idx']
+
+            if form == 'cmu': # the form of CMU-Pose
+                if result['image_id'] not in json_results_cmu.keys():
+                    json_results_cmu[result['image_id']]={}
+                    json_results_cmu[result['image_id']]['version']="AlphaPose v0.3"
+                    json_results_cmu[result['image_id']]['bodies']=[]
+                tmp={'joints':[]}
+                result['keypoints'].append((result['keypoints'][15]+result['keypoints'][18])/2)
+                result['keypoints'].append((result['keypoints'][16]+result['keypoints'][19])/2)
+                result['keypoints'].append((result['keypoints'][17]+result['keypoints'][20])/2)
+                indexarr=[0,51,18,24,30,15,21,27,36,42,48,33,39,45,6,3,12,9]
+                for i in indexarr:
+                    tmp['joints'].append(result['keypoints'][i])
+                    tmp['joints'].append(result['keypoints'][i+1])
+                    tmp['joints'].append(result['keypoints'][i+2])
+                json_results_cmu[result['image_id']]['bodies'].append(tmp)
+            elif form == 'open': # the form of OpenPose
+                if result['image_id'] not in json_results_cmu.keys():
+                    json_results_cmu[result['image_id']]={}
+                    json_results_cmu[result['image_id']]['version']="AlphaPose v0.3"
+                    json_results_cmu[result['image_id']]['people']=[]
+                tmp={'pose_keypoints_2d':[]}
+                result['keypoints'].append((result['keypoints'][15]+result['keypoints'][18])/2)
+                result['keypoints'].append((result['keypoints'][16]+result['keypoints'][19])/2)
+                result['keypoints'].append((result['keypoints'][17]+result['keypoints'][20])/2)
+                indexarr=[0,51,18,24,30,15,21,27,36,42,48,33,39,45,6,3,12,9]
+                for i in indexarr:
+                    tmp['pose_keypoints_2d'].append(result['keypoints'][i])
+                    tmp['pose_keypoints_2d'].append(result['keypoints'][i+1])
+                    tmp['pose_keypoints_2d'].append(result['keypoints'][i+2])
+                json_results_cmu[result['image_id']]['people'].append(tmp)
+            else:
+                json_results.append(result)
+
+    if form == 'cmu': # the form of CMU-Pose
+        with open(os.path.join(outputpath,'{}.json'.format(save_file_name)), 'w') as json_file:
+            json_file.write(json.dumps(json_results_cmu))
+            if not os.path.exists(os.path.join(outputpath,'sep-json')):
+                os.mkdir(os.path.join(outputpath,'sep-json'))
+            for name in json_results_cmu.keys():
+                with open(os.path.join(outputpath,'sep-json',name.split('.')[0]+'.json'),'w') as json_file:
+                    json_file.write(json.dumps(json_results_cmu[name]))
+    elif form == 'open': # the form of OpenPose
+        with open(os.path.join(outputpath,'{}.json'.format(save_file_name)), 'w') as json_file:
+            json_file.write(json.dumps(json_results_cmu))
+            if not os.path.exists(os.path.join(outputpath,'sep-json')):
+                os.mkdir(os.path.join(outputpath,'sep-json'))
+            for name in json_results_cmu.keys():
+                with open(os.path.join(outputpath,'sep-json',name.split('.')[0]+'.json'),'w') as json_file:
+                    json_file.write(json.dumps(json_results_cmu[name]))
+    else:
+        with open(os.path.join(outputpath,'{}.json'.format(save_file_name)), 'w') as json_file:
+            json_file.write(json.dumps(json_results))
+

alphapose/utils/registry.py

@@ -0,0 +1,103 @@
+import inspect
+
+
+class Registry(object):
+
+    def __init__(self, name):
+        self._name = name
+        self._module_dict = dict()
+
+    def __repr__(self):
+        format_str = self.__class__.__name__ + '(name={}, items={})'.format(
+            self._name, list(self._module_dict.keys()))
+        return format_str
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def module_dict(self):
+        return self._module_dict
+
+    def get(self, key):
+        return self._module_dict.get(key, None)
+
+    def _register_module(self, module_class):
+        """Register a module.
+
+        Args:
+            module (:obj:`nn.Module`): Module to be registered.
+        """
+        if not inspect.isclass(module_class):
+            raise TypeError('module must be a class, but got {}'.format(
+                type(module_class)))
+        module_name = module_class.__name__
+        if module_name in self._module_dict:
+            raise KeyError('{} is already registered in {}'.format(
+                module_name, self.name))
+        self._module_dict[module_name] = module_class
+
+    def register_module(self, cls):
+        self._register_module(cls)
+        return cls
+
+
+def build_from_cfg(cfg, registry, default_args=None):
+    """Build a module from config dict.
+
+    Args:
+        cfg (dict): Config dict. It should at least contain the key "type".
+        registry (:obj:`Registry`): The registry to search the type from.
+        default_args (dict, optional): Default initialization arguments.
+
+    Returns:
+        obj: The constructed object.
+    """
+    assert isinstance(cfg, dict) and 'TYPE' in cfg
+    assert isinstance(default_args, dict) or default_args is None
+    args = cfg.copy()
+    obj_type = args.pop('TYPE')
+
+    if isinstance(obj_type, str):
+        obj_cls = registry.get(obj_type)
+        if obj_cls is None:
+            raise KeyError('{} is not in the {} registry'.format(
+                obj_type, registry.name))
+    elif inspect.isclass(obj_type):
+        obj_cls = obj_type
+    else:
+        raise TypeError('type must be a str or valid type, but got {}'.format(
+            type(obj_type)))
+    if default_args is not None:
+        for name, value in default_args.items():
+            args.setdefault(name, value)
+    return obj_cls(**args)
+
+
+def retrieve_from_cfg(cfg, registry):
+    """Retrieve a module class from config dict.
+
+    Args:
+        cfg (dict): Config dict. It should at least contain the key "type".
+        registry (:obj:`Registry`): The registry to search the type from.
+
+    Returns:
+        class: The class.
+    """
+    assert isinstance(cfg, dict) and 'TYPE' in cfg
+    args = cfg.copy()
+    obj_type = args.pop('TYPE')
+
+    if isinstance(obj_type, str):
+        obj_cls = registry.get(obj_type)
+        if obj_cls is None:
+            raise KeyError('{} is not in the {} registry'.format(
+                obj_type, registry.name))
+    elif inspect.isclass(obj_type):
+        obj_cls = obj_type
+    else:
+        raise TypeError('type must be a str or valid type, but got {}'.format(
+            type(obj_type)))
+
+    return obj_cls

alphapose/utils/transforms.py

@@ -0,0 +1,810 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Jiefeng Li (jeff.lee.sjtu@gmail.com)
+# -----------------------------------------------------
+
+"""Pose related transforrmation functions."""
+
+import random
+
+import cv2
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+
+def rnd(x):
+    return max(-2 * x, min(2 * x, np.random.randn(1)[0] * x))
+
+
+def box_transform(bbox, sf, imgwidth, imght, train):
+    """Random scaling."""
+    width = bbox[2] - bbox[0]
+    ht = bbox[3] - bbox[1]
+    if train:
+        scaleRate = 0.25 * np.clip(np.random.randn() * sf, - sf, sf)
+
+        bbox[0] = max(0, bbox[0] - width * scaleRate / 2)
+        bbox[1] = max(0, bbox[1] - ht * scaleRate / 2)
+        bbox[2] = min(imgwidth, bbox[2] + width * scaleRate / 2)
+        bbox[3] = min(imght, bbox[3] + ht * scaleRate / 2)
+    else:
+        scaleRate = 0.25
+
+        bbox[0] = max(0, bbox[0] - width * scaleRate / 2)
+        bbox[1] = max(0, bbox[1] - ht * scaleRate / 2)
+        bbox[2] = min(imgwidth, max(bbox[2] + width * scaleRate / 2, bbox[0] + 5))
+        bbox[3] = min(imght, max(bbox[3] + ht * scaleRate / 2, bbox[1] + 5))
+
+    return bbox
+
+
+def addDPG(bbox, imgwidth, imght):
+    """Add dpg for data augmentation, including random crop and random sample."""
+    PatchScale = random.uniform(0, 1)
+    width = bbox[2] - bbox[0]
+    ht = bbox[3] - bbox[1]
+
+    if PatchScale > 0.85:
+        ratio = ht / width
+        if (width < ht):
+            patchWidth = PatchScale * width
+            patchHt = patchWidth * ratio
+        else:
+            patchHt = PatchScale * ht
+            patchWidth = patchHt / ratio
+
+        xmin = bbox[0] + random.uniform(0, 1) * (width - patchWidth)
+        ymin = bbox[1] + random.uniform(0, 1) * (ht - patchHt)
+        xmax = xmin + patchWidth + 1
+        ymax = ymin + patchHt + 1
+    else:
+        xmin = max(1, min(bbox[0] + np.random.normal(-0.0142, 0.1158) * width, imgwidth - 3))
+        ymin = max(1, min(bbox[1] + np.random.normal(0.0043, 0.068) * ht, imght - 3))
+        xmax = min(max(xmin + 2, bbox[2] + np.random.normal(0.0154, 0.1337) * width), imgwidth - 3)
+        ymax = min(max(ymin + 2, bbox[3] + np.random.normal(-0.0013, 0.0711) * ht), imght - 3)
+
+    bbox[0] = xmin
+    bbox[1] = ymin
+    bbox[2] = xmax
+    bbox[3] = ymax
+
+    return bbox
+
+
+def im_to_torch(img):
+    """Transform ndarray image to torch tensor.
+    Parameters
+    ----------
+    img: numpy.ndarray
+        An ndarray with shape: `(H, W, 3)`.
+    Returns
+    -------
+    torch.Tensor
+        A tensor with shape: `(3, H, W)`.
+    """
+    img = np.transpose(img, (2, 0, 1))  # C*H*W
+    img = to_torch(img).float()
+    if img.max() > 1:
+        img /= 255
+    return img
+
+
+def torch_to_im(img):
+    """Transform torch tensor to ndarray image.
+    Parameters
+    ----------
+    img: torch.Tensor
+        A tensor with shape: `(3, H, W)`.
+    Returns
+    -------
+    numpy.ndarray
+        An ndarray with shape: `(H, W, 3)`.
+    """
+    img = to_numpy(img)
+    img = np.transpose(img, (1, 2, 0))  # C*H*W
+    return img
+
+
+def load_image(img_path):
+    # H x W x C => C x H x W
+    return im_to_torch(cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB))#scipy.misc.imread(img_path, mode='RGB'))
+
+
+def to_numpy(tensor):
+    # torch.Tensor => numpy.ndarray
+    if torch.is_tensor(tensor):
+        return tensor.cpu().numpy()
+    elif type(tensor).__module__ != 'numpy':
+        raise ValueError("Cannot convert {} to numpy array"
+                         .format(type(tensor)))
+    return tensor
+
+
+def to_torch(ndarray):
+    # numpy.ndarray => torch.Tensor
+    if type(ndarray).__module__ == 'numpy':
+        return torch.from_numpy(ndarray)
+    elif not torch.is_tensor(ndarray):
+        raise ValueError("Cannot convert {} to torch tensor"
+                         .format(type(ndarray)))
+    return ndarray
+
+
+def cv_cropBox(img, bbox, input_size):
+    """Crop bbox from image by Affinetransform.
+    Parameters
+    ----------
+    img: torch.Tensor
+        A tensor with shape: `(3, H, W)`.
+    bbox: list or tuple
+        [xmin, ymin, xmax, ymax].
+    input_size: tuple
+        Resulting image size, as (height, width).
+    Returns
+    -------
+    torch.Tensor
+        A tensor with shape: `(3, height, width)`.
+    """
+    xmin, ymin, xmax, ymax = bbox
+    xmax -= 1
+    ymax -= 1
+    resH, resW = input_size
+
+    lenH = max((ymax - ymin), (xmax - xmin) * resH / resW)
+    lenW = lenH * resW / resH
+    if img.dim() == 2:
+        img = img[np.newaxis, :, :]
+
+    box_shape = [ymax - ymin, xmax - xmin]
+    pad_size = [(lenH - box_shape[0]) // 2, (lenW - box_shape[1]) // 2]
+    # Padding Zeros
+    img[:, :ymin, :], img[:, :, :xmin] = 0, 0
+    img[:, ymax + 1:, :], img[:, :, xmax + 1:] = 0, 0
+
+    src = np.zeros((3, 2), dtype=np.float32)
+    dst = np.zeros((3, 2), dtype=np.float32)
+
+    src[0, :] = np.array([xmin - pad_size[1], ymin - pad_size[0]], np.float32)
+    src[1, :] = np.array([xmax + pad_size[1], ymax + pad_size[0]], np.float32)
+    dst[0, :] = 0
+    dst[1, :] = np.array([resW - 1, resH - 1], np.float32)
+
+    src[2:, :] = get_3rd_point(src[0, :], src[1, :])
+    dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
+
+    trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
+    dst_img = cv2.warpAffine(torch_to_im(img), trans,
+                             (resW, resH), flags=cv2.INTER_LINEAR)
+    if dst_img.ndim == 2:
+        dst_img = dst_img[:, :, np.newaxis]
+
+    return im_to_torch(torch.Tensor(dst_img))
+
+
+def cv_cropBox_rot(img, bbox, input_size, rot):
+    """Crop bbox from image by Affinetransform.
+    Parameters
+    ----------
+    img: torch.Tensor
+        A tensor with shape: `(3, H, W)`.
+    bbox: list or tuple
+        [xmin, ymin, xmax, ymax].
+    input_size: tuple
+        Resulting image size, as (height, width).
+    Returns
+    -------
+    torch.Tensor
+        A tensor with shape: `(3, height, width)`.
+    """
+    xmin, ymin, xmax, ymax = bbox
+    xmax -= 1
+    ymax -= 1
+    resH, resW = input_size
+    rot_rad = np.pi * rot / 180
+
+    if img.dim() == 2:
+        img = img[np.newaxis, :, :]
+
+    src = np.zeros((3, 2), dtype=np.float32)
+    dst = np.zeros((3, 2), dtype=np.float32)
+    center = np.array([(xmax + xmin) / 2, (ymax + ymin) / 2])
+
+    src_dir = get_dir([0, (ymax - ymin) * -0.5], rot_rad)
+    dst_dir = np.array([0, (resH - 1) * -0.5], np.float32)
+
+    src = np.zeros((3, 2), dtype=np.float32)
+    dst = np.zeros((3, 2), dtype=np.float32)
+
+    src[0, :] = center
+    src[1, :] = center + src_dir
+    dst[0, :] = [(resW - 1) * 0.5, (resH - 1) * 0.5]
+    dst[1, :] = np.array([(resW - 1) * 0.5, (resH - 1) * 0.5]) + dst_dir
+
+    src[2:, :] = get_3rd_point(src[0, :], src[1, :])
+    dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
+
+    trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
+    dst_img = cv2.warpAffine(torch_to_im(img), trans,
+                             (resW, resH), flags=cv2.INTER_LINEAR)
+    if dst_img.ndim == 2:
+        dst_img = dst_img[:, :, np.newaxis]
+
+    return im_to_torch(torch.Tensor(dst_img))
+
+
+def fix_cropBox(img, bbox, input_size):
+    """Crop bbox from image by Affinetransform.
+    Parameters
+    ----------
+    img: torch.Tensor
+        A tensor with shape: `(3, H, W)`.
+    bbox: list or tuple
+        [xmin, ymin, xmax, ymax].
+    input_size: tuple
+        Resulting image size, as (height, width).
+    Returns
+    -------
+    torch.Tensor
+        A tensor with shape: `(3, height, width)`.
+    """
+    xmin, ymin, xmax, ymax = bbox
+    input_ratio = input_size[0] / input_size[1]
+    bbox_ratio = (ymax - ymin) / (xmax - xmin)
+    if bbox_ratio > input_ratio:
+        # expand width
+        cx = (xmax + xmin) / 2
+        h = ymax - ymin
+        w = h / input_ratio
+        xmin = cx - w / 2
+        xmax = cx + w / 2
+    elif bbox_ratio < input_ratio:
+        # expand height
+        cy = (ymax + ymin) / 2
+        w = xmax - xmin
+        h = w * input_ratio
+        ymin = cy - h / 2
+        ymax = cy + h / 2
+    bbox = [int(x) for x in [xmin, ymin, xmax, ymax]]
+
+    return cv_cropBox(img, bbox, input_size), bbox
+
+
+def fix_cropBox_rot(img, bbox, input_size, rot):
+    """Crop bbox from image by Affinetransform.
+    Parameters
+    ----------
+    img: torch.Tensor
+        A tensor with shape: `(3, H, W)`.
+    bbox: list or tuple
+        [xmin, ymin, xmax, ymax].
+    input_size: tuple
+        Resulting image size, as (height, width).
+    Returns
+    -------
+    torch.Tensor
+        A tensor with shape: `(3, height, width)`.
+    """
+    xmin, ymin, xmax, ymax = bbox
+    input_ratio = input_size[0] / input_size[1]
+    bbox_ratio = (ymax - ymin) / (xmax - xmin)
+    if bbox_ratio > input_ratio:
+        # expand width
+        cx = (xmax + xmin) / 2
+        h = ymax - ymin
+        w = h / input_ratio
+        xmin = cx - w / 2
+        xmax = cx + w / 2
+    elif bbox_ratio < input_ratio:
+        # expand height
+        cy = (ymax + ymin) / 2
+        w = xmax - xmin
+        h = w * input_ratio
+        ymin = cy - h / 2
+        ymax = cy + h / 2
+    bbox = [int(x) for x in [xmin, ymin, xmax, ymax]]
+
+    return cv_cropBox_rot(img, bbox, input_size, rot), bbox
+
+
+def get_3rd_point(a, b):
+    """Return vector c that perpendicular to (a - b)."""
+    direct = a - b
+    return b + np.array([-direct[1], direct[0]], dtype=np.float32)
+
+
+def get_dir(src_point, rot_rad):
+    """Rotate the point by `rot_rad` degree."""
+    sn, cs = np.sin(rot_rad), np.cos(rot_rad)
+
+    src_result = [0, 0]
+    src_result[0] = src_point[0] * cs - src_point[1] * sn
+    src_result[1] = src_point[0] * sn + src_point[1] * cs
+
+    return src_result
+
+
+def cv_cropBoxInverse(inp, bbox, img_size, output_size):
+    """Paste the cropped bbox to the original image.
+    Parameters
+    ----------
+    inp: torch.Tensor
+        A tensor with shape: `(3, height, width)`.
+    bbox: list or tuple
+        [xmin, ymin, xmax, ymax].
+    img_size: tuple
+        Original image size, as (img_H, img_W).
+    output_size: tuple
+        Cropped input size, as (height, width).
+    Returns
+    -------
+    torch.Tensor
+        A tensor with shape: `(3, img_H, img_W)`.
+    """
+    xmin, ymin, xmax, ymax = bbox
+    xmax -= 1
+    ymax -= 1
+    resH, resW = output_size
+    imgH, imgW = img_size
+
+    lenH = max((ymax - ymin), (xmax - xmin) * resH / resW)
+    lenW = lenH * resW / resH
+    if inp.dim() == 2:
+        inp = inp[np.newaxis, :, :]
+
+    box_shape = [ymax - ymin, xmax - xmin]
+    pad_size = [(lenH - box_shape[0]) // 2, (lenW - box_shape[1]) // 2]
+
+    src = np.zeros((3, 2), dtype=np.float32)
+    dst = np.zeros((3, 2), dtype=np.float32)
+
+    src[0, :] = 0
+    src[1, :] = np.array([resW - 1, resH - 1], np.float32)
+    dst[0, :] = np.array([xmin - pad_size[1], ymin - pad_size[0]], np.float32)
+    dst[1, :] = np.array([xmax + pad_size[1], ymax + pad_size[0]], np.float32)
+
+    src[2:, :] = get_3rd_point(src[0, :], src[1, :])
+    dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
+
+    trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
+    dst_img = cv2.warpAffine(torch_to_im(inp), trans,
+                             (imgW, imgH), flags=cv2.INTER_LINEAR)
+    if dst_img.ndim == 3 and dst_img.shape[2] == 1:
+        dst_img = dst_img[:, :, 0]
+        return dst_img
+    elif dst_img.ndim == 2:
+        return dst_img
+    else:
+        return im_to_torch(torch.Tensor(dst_img))
+
+
+def cv_rotate(img, rot, input_size):
+    """Rotate image by Affinetransform.
+    Parameters
+    ----------
+    img: torch.Tensor
+        A tensor with shape: `(3, H, W)`.
+    rot: int
+        Rotation degree.
+    input_size: tuple
+        Resulting image size, as (height, width).
+    Returns
+    -------
+    torch.Tensor
+        A tensor with shape: `(3, height, width)`.
+    """
+    resH, resW = input_size
+    center = np.array((resW - 1, resH - 1)) / 2
+    rot_rad = np.pi * rot / 180
+
+    src_dir = get_dir([0, (resH - 1) * -0.5], rot_rad)
+    dst_dir = np.array([0, (resH - 1) * -0.5], np.float32)
+
+    src = np.zeros((3, 2), dtype=np.float32)
+    dst = np.zeros((3, 2), dtype=np.float32)
+
+    src[0, :] = center
+    src[1, :] = center + src_dir
+    dst[0, :] = [(resW - 1) * 0.5, (resH - 1) * 0.5]
+    dst[1, :] = np.array([(resW - 1) * 0.5, (resH - 1) * 0.5]) + dst_dir
+
+    src[2:, :] = get_3rd_point(src[0, :], src[1, :])
+    dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
+
+    trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
+
+    dst_img = cv2.warpAffine(torch_to_im(img), trans,
+                             (resW, resH), flags=cv2.INTER_LINEAR)
+    if dst_img.ndim == 2:
+        dst_img = dst_img[:, :, np.newaxis]
+
+    return im_to_torch(torch.Tensor(dst_img))
+
+
+def count_visible(bbox, joints_3d):
+    """Count number of visible joints given bound box."""
+    vis = np.logical_and.reduce((
+        joints_3d[:, 0, 0] > 0,
+        joints_3d[:, 0, 0] > bbox[0],
+        joints_3d[:, 0, 0] < bbox[2],
+        joints_3d[:, 1, 0] > 0,
+        joints_3d[:, 1, 0] > bbox[1],
+        joints_3d[:, 1, 0] < bbox[3],
+        joints_3d[:, 0, 1] > 0,
+        joints_3d[:, 1, 1] > 0
+    ))
+    return np.sum(vis), vis
+
+
+def drawGaussian(img, pt, sigma):
+    """Draw 2d gaussian on input image.
+    Parameters
+    ----------
+    img: torch.Tensor
+        A tensor with shape: `(3, H, W)`.
+    pt: list or tuple
+        A point: (x, y).
+    sigma: int
+        Sigma of gaussian distribution.
+    Returns
+    -------
+    torch.Tensor
+        A tensor with shape: `(3, H, W)`.
+    """
+    img = to_numpy(img)
+    tmpSize = 3 * sigma
+    # Check that any part of the gaussian is in-bounds
+    ul = [int(pt[0] - tmpSize), int(pt[1] - tmpSize)]
+    br = [int(pt[0] + tmpSize + 1), int(pt[1] + tmpSize + 1)]
+
+    if (ul[0] >= img.shape[1] or ul[1] >= img.shape[0] or br[0] < 0 or br[1] < 0):
+        # If not, just return the image as is
+        return to_torch(img)
+
+    # Generate gaussian
+    size = 2 * tmpSize + 1
+    x = np.arange(0, size, 1, float)
+    y = x[:, np.newaxis]
+    x0 = y0 = size // 2
+    # The gaussian is not normalized, we want the center value to equal 1
+    g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2))
+
+    # Usable gaussian range
+    g_x = max(0, -ul[0]), min(br[0], img.shape[1]) - ul[0]
+    g_y = max(0, -ul[1]), min(br[1], img.shape[0]) - ul[1]
+    # Image range
+    img_x = max(0, ul[0]), min(br[0], img.shape[1])
+    img_y = max(0, ul[1]), min(br[1], img.shape[0])
+
+    img[img_y[0]:img_y[1], img_x[0]:img_x[1]] = g[g_y[0]:g_y[1], g_x[0]:g_x[1]]
+    return to_torch(img)
+
+
+def flip(x):
+    assert (x.dim() == 3 or x.dim() == 4)
+    dim = x.dim() - 1
+
+    return x.flip(dims=(dim,))
+
+
+def flip_heatmap(heatmap, joint_pairs, shift=False):
+    """Flip pose heatmap according to joint pairs.
+    Parameters
+    ----------
+    heatmap : numpy.ndarray
+        Heatmap of joints.
+    joint_pairs : list
+        List of joint pairs.
+    shift : bool
+        Whether to shift the output.
+    Returns
+    -------
+    numpy.ndarray
+        Flipped heatmap.
+    """
+    assert (heatmap.dim() == 3 or heatmap.dim() == 4)
+    out = flip(heatmap)
+
+    for pair in joint_pairs:
+        dim0, dim1 = pair
+        idx = torch.Tensor((dim0, dim1)).long()
+        inv_idx = torch.Tensor((dim1, dim0)).long()
+        if out.dim() == 4:
+            out[:, idx] = out[:, inv_idx]
+        else:
+            out[idx] = out[inv_idx]
+
+    if shift:
+        if out.dim() == 3:
+            out[:, :, 1:] = out[:, :, 0:-1]
+        else:
+            out[:, :, :, 1:] = out[:, :, :, 0:-1]
+    return out
+
+
+def flip_joints_3d(joints_3d, width, joint_pairs):
+    """Flip 3d joints.
+    Parameters
+    ----------
+    joints_3d : numpy.ndarray
+        Joints in shape (num_joints, 3, 2)
+    width : int
+        Image width.
+    joint_pairs : list
+        List of joint pairs.
+    Returns
+    -------
+    numpy.ndarray
+        Flipped 3d joints with shape (num_joints, 3, 2)
+    """
+    joints = joints_3d.copy()
+    # flip horizontally
+    joints[:, 0, 0] = width - joints[:, 0, 0] - 1
+    # change left-right parts
+    for pair in joint_pairs:
+        joints[pair[0], :, 0], joints[pair[1], :, 0] = \
+            joints[pair[1], :, 0], joints[pair[0], :, 0].copy()
+        joints[pair[0], :, 1], joints[pair[1], :, 1] = \
+            joints[pair[1], :, 1], joints[pair[0], :, 1].copy()
+
+    joints[:, :, 0] *= joints[:, :, 1]
+    return joints
+
+
+def heatmap_to_coord_simple(hms, bbox, hms_flip=None, **kwargs):
+    if hms_flip is not None:
+        hms = (hms + hms_flip) / 2
+    if not isinstance(hms,np.ndarray):
+        hms = hms.cpu().data.numpy()
+    coords, maxvals = get_max_pred(hms)
+
+    hm_h = hms.shape[1]
+    hm_w = hms.shape[2]
+
+    # post-processing
+    for p in range(coords.shape[0]):
+        hm = hms[p]
+        px = int(round(float(coords[p][0])))
+        py = int(round(float(coords[p][1])))
+        if 1 < px < hm_w - 1 and 1 < py < hm_h - 1:
+            diff = np.array((hm[py][px + 1] - hm[py][px - 1],
+                             hm[py + 1][px] - hm[py - 1][px]))
+            coords[p] += np.sign(diff) * .25
+
+    preds = np.zeros_like(coords)
+
+    # transform bbox to scale
+    xmin, ymin, xmax, ymax = bbox
+    w = xmax - xmin
+    h = ymax - ymin
+    center = np.array([xmin + w * 0.5, ymin + h * 0.5])
+    scale = np.array([w, h])
+    # Transform back
+    for i in range(coords.shape[0]):
+        preds[i] = transform_preds(coords[i], center, scale,
+                                   [hm_w, hm_h])
+
+    return preds, maxvals
+
+
+def heatmap_to_coord_simple_regress(preds, bbox, hm_shape, norm_type, hms_flip=None):
+    def integral_op(hm_1d):
+        if hm_1d.device.index is not None:
+            hm_1d = hm_1d * torch.cuda.comm.broadcast(torch.arange(hm_1d.shape[-1]).type(
+                torch.cuda.FloatTensor), devices=[hm_1d.device.index])[0]
+        else:
+            hm_1d = hm_1d * torch.arange(hm_1d.shape[-1]).type(torch.FloatTensor)
+        return hm_1d
+
+    if preds.dim() == 3:
+        preds = preds.unsqueeze(0)
+    hm_height, hm_width = hm_shape
+    num_joints = preds.shape[1]
+
+    pred_jts, pred_scores = _integral_tensor(preds, num_joints, False, hm_width, hm_height, 1, integral_op, norm_type)
+    pred_jts = pred_jts.reshape(pred_jts.shape[0], num_joints, 2)
+
+    if hms_flip is not None:
+        if hms_flip.dim() == 3:
+            hms_flip = hms_flip.unsqueeze(0)
+        pred_jts_flip, pred_scores_flip = _integral_tensor(hms_flip, num_joints, False, hm_width, hm_height, 1, integral_op, norm_type)
+        pred_jts_flip = pred_jts_flip.reshape(pred_jts_flip.shape[0], num_joints, 2)
+
+        pred_jts = (pred_jts + pred_jts_flip) / 2
+        pred_scores = (pred_scores + pred_scores_flip) / 2
+
+    ndims = pred_jts.dim()
+    assert ndims in [2, 3], "Dimensions of input heatmap should be 3 or 4"
+    if ndims == 2:
+        pred_jts = pred_jts.unsqueeze(0)
+        pred_scores = pred_scores.unsqueeze(0)
+
+    coords = pred_jts.cpu().numpy()
+    coords = coords.astype(np.float32)
+    pred_scores = pred_scores.cpu().numpy()
+    pred_scores = pred_scores.astype(np.float32)
+
+    coords[:, :, 0] = (coords[:, :, 0] + 0.5) * hm_width
+    coords[:, :, 1] = (coords[:, :, 1] + 0.5) * hm_height
+
+    preds = np.zeros_like(coords)
+    # transform bbox to scale
+    xmin, ymin, xmax, ymax = bbox
+    w = xmax - xmin
+    h = ymax - ymin
+    center = np.array([xmin + w * 0.5, ymin + h * 0.5])
+    scale = np.array([w, h])
+    # Transform back
+    for i in range(coords.shape[0]):
+        for j in range(coords.shape[1]):
+            preds[i, j, 0:2] = transform_preds(coords[i, j, 0:2], center, scale,
+                                               [hm_width, hm_height])
+
+    if preds.shape[0] == 1:
+        preds = preds[0]
+        pred_scores = pred_scores[0]
+    return preds, pred_scores
+
+
+def _integral_tensor(preds, num_joints, output_3d, hm_width, hm_height, hm_depth, integral_operation, norm_type='softmax'):
+    # normalization
+    preds = preds.reshape((preds.shape[0], num_joints, -1))
+    preds = norm_heatmap(norm_type, preds)
+
+    # get heatmap confidence
+    if norm_type == 'sigmoid':
+        maxvals, _ = torch.max(preds, dim=2, keepdim=True)
+    else:
+        maxvals = torch.ones(
+            (*preds.shape[:2], 1), dtype=torch.float, device=preds.device)
+
+    # normalized to probability
+    heatmaps = preds / preds.sum(dim=2, keepdim=True)
+    heatmaps = heatmaps.reshape(
+        (heatmaps.shape[0], num_joints, hm_depth, hm_height, hm_width))
+
+    # The edge probability
+    hm_x = heatmaps.sum((2, 3))
+    hm_y = heatmaps.sum((2, 4))
+    hm_z = heatmaps.sum((3, 4))
+
+    hm_x = integral_operation(hm_x)
+    hm_y = integral_operation(hm_y)
+    hm_z = integral_operation(hm_z)
+
+    coord_x = hm_x.sum(dim=2, keepdim=True)
+    coord_y = hm_y.sum(dim=2, keepdim=True)
+    coord_z = hm_z.sum(dim=2, keepdim=True)
+
+    coord_x = coord_x / float(hm_width) - 0.5
+    coord_y = coord_y / float(hm_height) - 0.5
+    if output_3d:
+        coord_z = coord_z / float(hm_depth) - 0.5
+        pred_jts = torch.cat((coord_x, coord_y, coord_z), dim=2)
+        pred_jts = pred_jts.reshape((pred_jts.shape[0], num_joints * 3))
+    else:
+        pred_jts = torch.cat((coord_x, coord_y), dim=2)
+        pred_jts = pred_jts.reshape((pred_jts.shape[0], num_joints * 2))
+    return pred_jts, maxvals.float()
+
+
+def norm_heatmap(norm_type, heatmap):
+    # Input tensor shape: [N,C,...]
+    shape = heatmap.shape
+    if norm_type == 'softmax':
+        heatmap = heatmap.reshape(*shape[:2], -1)
+        # global soft max
+        heatmap = F.softmax(heatmap, 2)
+        return heatmap.reshape(*shape)
+    elif norm_type == 'sigmoid':
+        return heatmap.sigmoid()
+    elif norm_type == 'divide_sum':
+        heatmap = heatmap.reshape(*shape[:2], -1)
+        heatmap = heatmap / heatmap.sum(dim=2, keepdim=True)
+        return heatmap.reshape(*shape)
+    else:
+        raise NotImplementedError
+
+
+def transform_preds(coords, center, scale, output_size):
+    target_coords = np.zeros(coords.shape)
+    trans = get_affine_transform(center, scale, 0, output_size, inv=1)
+    target_coords[0:2] = affine_transform(coords[0:2], trans)
+    return target_coords
+
+
+def get_max_pred(heatmaps):
+    num_joints = heatmaps.shape[0]
+    width = heatmaps.shape[2]
+    heatmaps_reshaped = heatmaps.reshape((num_joints, -1))
+    idx = np.argmax(heatmaps_reshaped, 1)
+    maxvals = np.max(heatmaps_reshaped, 1)
+
+    maxvals = maxvals.reshape((num_joints, 1))
+    idx = idx.reshape((num_joints, 1))
+
+    preds = np.tile(idx, (1, 2)).astype(np.float32)
+
+    preds[:, 0] = (preds[:, 0]) % width
+    preds[:, 1] = np.floor((preds[:, 1]) / width)
+
+    pred_mask = np.tile(np.greater(maxvals, 0.0), (1, 2))
+    pred_mask = pred_mask.astype(np.float32)
+
+    preds *= pred_mask
+    return preds, maxvals
+
+
+def get_max_pred_batch(batch_heatmaps):
+    batch_size = batch_heatmaps.shape[0]
+    num_joints = batch_heatmaps.shape[1]
+    width = batch_heatmaps.shape[3]
+    heatmaps_reshaped = batch_heatmaps.reshape((batch_size, num_joints, -1))
+    idx = np.argmax(heatmaps_reshaped, 2)
+    maxvals = np.max(heatmaps_reshaped, 2)
+
+    maxvals = maxvals.reshape((batch_size, num_joints, 1))
+    idx = idx.reshape((batch_size, num_joints, 1))
+
+    preds = np.tile(idx, (1, 1, 2)).astype(np.float32)
+
+    preds[:, :, 0] = (preds[:, :, 0]) % width
+    preds[:, :, 1] = np.floor((preds[:, :, 1]) / width)
+
+    pred_mask = np.tile(np.greater(maxvals, 0.0), (1, 1, 2))
+    pred_mask = pred_mask.astype(np.float32)
+
+    preds *= pred_mask
+    return preds, maxvals
+
+
+def get_affine_transform(center,
+                         scale,
+                         rot,
+                         output_size,
+                         shift=np.array([0, 0], dtype=np.float32),
+                         inv=0):
+    if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
+        scale = np.array([scale, scale])
+
+    scale_tmp = scale
+    src_w = scale_tmp[0]
+    dst_w = output_size[0]
+    dst_h = output_size[1]
+
+    rot_rad = np.pi * rot / 180
+    src_dir = get_dir([0, src_w * -0.5], rot_rad)
+    dst_dir = np.array([0, dst_w * -0.5], np.float32)
+
+    src = np.zeros((3, 2), dtype=np.float32)
+    dst = np.zeros((3, 2), dtype=np.float32)
+    src[0, :] = center + scale_tmp * shift
+    src[1, :] = center + src_dir + scale_tmp * shift
+    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
+    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
+
+    src[2:, :] = get_3rd_point(src[0, :], src[1, :])
+    dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
+
+    if inv:
+        trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
+    else:
+        trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
+
+    return trans
+
+
+def affine_transform(pt, t):
+    new_pt = np.array([pt[0], pt[1], 1.]).T
+    new_pt = np.dot(t, new_pt)
+    return new_pt[:2]
+
+
+def get_func_heatmap_to_coord(cfg):
+    if cfg.DATA_PRESET.TYPE == 'simple':
+        if cfg.LOSS.TYPE == 'MSELoss':
+            return heatmap_to_coord_simple
+        elif cfg.LOSS.TYPE == 'L1JointRegression':
+            return heatmap_to_coord_simple_regress
+    else:
+        raise NotImplementedError

alphapose/utils/vis.py

@@ -0,0 +1,359 @@
+import math
+import time
+
+import cv2
+import numpy as np
+import torch
+
+RED = (0, 0, 255)
+GREEN = (0, 255, 0)
+BLUE = (255, 0, 0)
+CYAN = (255, 255, 0)
+YELLOW = (0, 255, 255)
+ORANGE = (0, 165, 255)
+PURPLE = (255, 0, 255)
+WHITE = (255, 255, 255)
+BLACK = (0, 0, 0)
+
+DEFAULT_FONT = cv2.FONT_HERSHEY_SIMPLEX
+
+
+def get_color(idx):
+    idx = idx * 3
+    color = ((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255)
+
+    return color
+
+
+def get_color_fast(idx):
+    color_pool = [RED, GREEN, BLUE, CYAN, YELLOW, ORANGE, PURPLE, WHITE]
+    color = color_pool[idx % 8]
+
+    return color
+
+
+def vis_frame_fast(frame, im_res, opt, format='coco'):
+    '''
+    frame: frame image
+    im_res: im_res of predictions
+    format: coco or mpii
+
+    return rendered image
+    '''
+    kp_num = 17
+    if len(im_res['result']) > 0:
+    	kp_num = len(im_res['result'][0]['keypoints'])
+    if kp_num == 17:
+        if format == 'coco':
+            l_pair = [
+                (0, 1), (0, 2), (1, 3), (2, 4),  # Head
+                (5, 6), (5, 7), (7, 9), (6, 8), (8, 10),
+                (17, 11), (17, 12),  # Body
+                (11, 13), (12, 14), (13, 15), (14, 16)
+            ]
+            p_color = [(0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0),  # Nose, LEye, REye, LEar, REar
+                       (77, 255, 255), (77, 255, 204), (77, 204, 255), (191, 255, 77), (77, 191, 255), (191, 255, 77),  # LShoulder, RShoulder, LElbow, RElbow, LWrist, RWrist
+                       (204, 77, 255), (77, 255, 204), (191, 77, 255), (77, 255, 191), (127, 77, 255), (77, 255, 127), (0, 255, 255)]  # LHip, RHip, LKnee, Rknee, LAnkle, RAnkle, Neck
+            line_color = [(0, 215, 255), (0, 255, 204), (0, 134, 255), (0, 255, 50),
+                          (77, 255, 222), (77, 196, 255), (77, 135, 255), (191, 255, 77), (77, 255, 77),
+                          (77, 222, 255), (255, 156, 127),
+                          (0, 127, 255), (255, 127, 77), (0, 77, 255), (255, 77, 36)]
+        elif format == 'mpii':
+            l_pair = [
+                (8, 9), (11, 12), (11, 10), (2, 1), (1, 0),
+                (13, 14), (14, 15), (3, 4), (4, 5),
+                (8, 7), (7, 6), (6, 2), (6, 3), (8, 12), (8, 13)
+            ]
+            p_color = [PURPLE, BLUE, BLUE, RED, RED, BLUE, BLUE, RED, RED, PURPLE, PURPLE, PURPLE, RED, RED, BLUE, BLUE]
+        else:
+            raise NotImplementedError
+    elif kp_num == 136:
+        l_pair = [
+            (0, 1), (0, 2), (1, 3), (2, 4),  # Head
+            (5, 18), (6, 18), (5, 7), (7, 9), (6, 8), (8, 10),# Body
+            (17, 18), (18, 19), (19, 11), (19, 12),
+            (11, 13), (12, 14), (13, 15), (14, 16),
+            (20, 24), (21, 25), (23, 25), (22, 24), (15, 24), (16, 25),# Foot
+            (26, 27),(27, 28),(28, 29),(29, 30),(30, 31),(31, 32),(32, 33),(33, 34),(34, 35),(35, 36),(36, 37),(37, 38),#Face
+            (38, 39),(39, 40),(40, 41),(41, 42),(43, 44),(44, 45),(45, 46),(46, 47),(48, 49),(49, 50),(50, 51),(51, 52),#Face
+            (53, 54),(54, 55),(55, 56),(57, 58),(58, 59),(59, 60),(60, 61),(62, 63),(63, 64),(64, 65),(65, 66),(66, 67),#Face
+            (68, 69),(69, 70),(70, 71),(71, 72),(72, 73),(74, 75),(75, 76),(76, 77),(77, 78),(78, 79),(79, 80),(80, 81),#Face
+            (81, 82),(82, 83),(83, 84),(84, 85),(85, 86),(86, 87),(87, 88),(88, 89),(89, 90),(90, 91),(91, 92),(92, 93),#Face
+            (94,95),(95,96),(96,97),(97,98),(94,99),(99,100),(100,101),(101,102),(94,103),(103,104),(104,105),#LeftHand
+            (105,106),(94,107),(107,108),(108,109),(109,110),(94,111),(111,112),(112,113),(113,114),#LeftHand
+            (115,116),(116,117),(117,118),(118,119),(115,120),(120,121),(121,122),(122,123),(115,124),(124,125),#RightHand
+            (125,126),(126,127),(115,128),(128,129),(129,130),(130,131),(115,132),(132,133),(133,134),(134,135)#RightHand
+        ]
+        p_color = [(0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0),  # Nose, LEye, REye, LEar, REar
+                   (77, 255, 255), (77, 255, 204), (77, 204, 255), (191, 255, 77), (77, 191, 255), (191, 255, 77),  # LShoulder, RShoulder, LElbow, RElbow, LWrist, RWrist
+                   (204, 77, 255), (77, 255, 204), (191, 77, 255), (77, 255, 191), (127, 77, 255), (77, 255, 127),  # LHip, RHip, LKnee, Rknee, LAnkle, RAnkle, Neck
+                   (77, 255, 255), (0, 255, 255), (77, 204, 255),  # head, neck, shoulder
+                   (0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0), (77, 255, 255)] # foot
+    
+        line_color = [(0, 215, 255), (0, 255, 204), (0, 134, 255), (0, 255, 50),
+                      (0, 255, 102), (77, 255, 222), (77, 196, 255), (77, 135, 255), (191, 255, 77), (77, 255, 77),
+                      (77, 191, 255), (204, 77, 255), (77, 222, 255), (255, 156, 127),
+                      (0, 127, 255), (255, 127, 77), (0, 77, 255), (255, 77, 36), 
+                      (0, 77, 255), (0, 77, 255), (0, 77, 255), (0, 77, 255), (255, 156, 127), (255, 156, 127)]
+    elif kp_num == 26:
+        l_pair = [
+            (0, 1), (0, 2), (1, 3), (2, 4),  # Head
+            (5, 18), (6, 18), (5, 7), (7, 9), (6, 8), (8, 10),# Body
+            (17, 18), (18, 19), (19, 11), (19, 12),
+            (11, 13), (12, 14), (13, 15), (14, 16),
+            (20, 24), (21, 25), (23, 25), (22, 24), (15, 24), (16, 25),# Foot
+        ]
+        p_color = [(0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0),  # Nose, LEye, REye, LEar, REar
+                   (77, 255, 255), (77, 255, 204), (77, 204, 255), (191, 255, 77), (77, 191, 255), (191, 255, 77),  # LShoulder, RShoulder, LElbow, RElbow, LWrist, RWrist
+                   (204, 77, 255), (77, 255, 204), (191, 77, 255), (77, 255, 191), (127, 77, 255), (77, 255, 127),  # LHip, RHip, LKnee, Rknee, LAnkle, RAnkle, Neck
+                   (77, 255, 255), (0, 255, 255), (77, 204, 255),  # head, neck, shoulder
+                   (0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0), (77, 255, 255)] # foot
+    
+        line_color = [(0, 215, 255), (0, 255, 204), (0, 134, 255), (0, 255, 50),
+                      (0, 255, 102), (77, 255, 222), (77, 196, 255), (77, 135, 255), (191, 255, 77), (77, 255, 77),
+                      (77, 191, 255), (204, 77, 255), (77, 222, 255), (255, 156, 127),
+                      (0, 127, 255), (255, 127, 77), (0, 77, 255), (255, 77, 36), 
+                      (0, 77, 255), (0, 77, 255), (0, 77, 255), (0, 77, 255), (255, 156, 127), (255, 156, 127)]
+    else:
+        raise NotImplementedError
+    # im_name = os.path.basename(im_res['imgname'])
+    img = frame.copy()
+    height, width = img.shape[:2]
+    for human in im_res['result']:
+        part_line = {}
+        kp_preds = human['keypoints']
+        kp_scores = human['kp_score']
+        if kp_num == 17:
+            kp_preds = torch.cat((kp_preds, torch.unsqueeze((kp_preds[5, :] + kp_preds[6, :]) / 2, 0)))
+            kp_scores = torch.cat((kp_scores, torch.unsqueeze((kp_scores[5, :] + kp_scores[6, :]) / 2, 0)))
+        if opt.pose_track or opt.tracking:
+            color = get_color_fast(int(abs(human['idx'])))
+        else:
+            color = BLUE
+
+        # Draw bboxes
+        if opt.showbox:
+            if 'box' in human.keys():
+                bbox = human['box']
+                bbox = [bbox[0], bbox[0]+bbox[2], bbox[1], bbox[1]+bbox[3]]#xmin,xmax,ymin,ymax
+            else:
+                from trackers.PoseFlow.poseflow_infer import get_box
+                keypoints = []
+                for n in range(kp_scores.shape[0]):
+                    keypoints.append(float(kp_preds[n, 0]))
+                    keypoints.append(float(kp_preds[n, 1]))
+                    keypoints.append(float(kp_scores[n]))
+                bbox = get_box(keypoints, height, width)
+            
+            cv2.rectangle(img, (int(bbox[0]), int(bbox[2])), (int(bbox[1]), int(bbox[3])), color, 2)
+            if opt.tracking:
+                cv2.putText(img, str(human['idx']), (int(bbox[0]), int((bbox[2] + 26))), DEFAULT_FONT, 1, BLACK, 2)
+        # Draw keypoints
+        vis_thres = 0.05 if kp_num == 136 else 0.4
+        for n in range(kp_scores.shape[0]):
+            if kp_scores[n] <= vis_thres:
+                continue
+            cor_x, cor_y = int(kp_preds[n, 0]), int(kp_preds[n, 1])
+            part_line[n] = (cor_x, cor_y)
+            if n < len(p_color):
+                if opt.tracking:
+                    cv2.circle(img, (cor_x, cor_y), 3, color, -1)
+                else:
+                    cv2.circle(img, (cor_x, cor_y), 3, p_color[n], -1)
+            else:
+                cv2.circle(img, (cor_x, cor_y), 1, (255,255,255), 2)
+        # Draw limbs
+        for i, (start_p, end_p) in enumerate(l_pair):
+            if start_p in part_line and end_p in part_line:
+                start_xy = part_line[start_p]
+                end_xy = part_line[end_p]
+                if i < len(line_color):
+                    if opt.tracking:
+                        cv2.line(img, start_xy, end_xy, color, 2 * int(kp_scores[start_p] + kp_scores[end_p]) + 1)
+                    else:
+                        cv2.line(img, start_xy, end_xy, line_color[i], 2 * int(kp_scores[start_p] + kp_scores[end_p]) + 1)
+                else:
+                    cv2.line(img, start_xy, end_xy, (255,255,255), 1)  
+
+    return img
+
+
+def vis_frame(frame, im_res, opt, format='coco'):
+    '''
+    frame: frame image
+    im_res: im_res of predictions
+    format: coco or mpii
+
+    return rendered image
+    '''
+    kp_num = 17
+    if len(im_res['result']) > 0:
+    	kp_num = len(im_res['result'][0]['keypoints'])
+
+    if kp_num == 17:
+        if format == 'coco':
+            l_pair = [
+                (0, 1), (0, 2), (1, 3), (2, 4),  # Head
+                (5, 6), (5, 7), (7, 9), (6, 8), (8, 10),
+                (17, 11), (17, 12),  # Body
+                (11, 13), (12, 14), (13, 15), (14, 16)
+            ]
+
+            p_color = [(0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0),  # Nose, LEye, REye, LEar, REar
+                       (77, 255, 255), (77, 255, 204), (77, 204, 255), (191, 255, 77), (77, 191, 255), (191, 255, 77),  # LShoulder, RShoulder, LElbow, RElbow, LWrist, RWrist
+                       (204, 77, 255), (77, 255, 204), (191, 77, 255), (77, 255, 191), (127, 77, 255), (77, 255, 127), (0, 255, 255)]  # LHip, RHip, LKnee, Rknee, LAnkle, RAnkle, Neck
+            line_color = [(0, 215, 255), (0, 255, 204), (0, 134, 255), (0, 255, 50),
+                          (77, 255, 222), (77, 196, 255), (77, 135, 255), (191, 255, 77), (77, 255, 77),
+                          (77, 222, 255), (255, 156, 127),
+                          (0, 127, 255), (255, 127, 77), (0, 77, 255), (255, 77, 36)]
+        elif format == 'mpii':
+            l_pair = [
+                (8, 9), (11, 12), (11, 10), (2, 1), (1, 0),
+                (13, 14), (14, 15), (3, 4), (4, 5),
+                (8, 7), (7, 6), (6, 2), (6, 3), (8, 12), (8, 13)
+            ]
+            p_color = [PURPLE, BLUE, BLUE, RED, RED, BLUE, BLUE, RED, RED, PURPLE, PURPLE, PURPLE, RED, RED, BLUE, BLUE]
+            line_color = [PURPLE, BLUE, BLUE, RED, RED, BLUE, BLUE, RED, RED, PURPLE, PURPLE, RED, RED, BLUE, BLUE]
+        else:
+            raise NotImplementedError
+    elif kp_num == 136:
+        l_pair = [
+            (0, 1), (0, 2), (1, 3), (2, 4),  # Head
+            (5, 18), (6, 18), (5, 7), (7, 9), (6, 8), (8, 10),# Body
+            (17, 18), (18, 19), (19, 11), (19, 12),
+            (11, 13), (12, 14), (13, 15), (14, 16),
+            (20, 24), (21, 25), (23, 25), (22, 24), (15, 24), (16, 25),# Foot
+            (26, 27),(27, 28),(28, 29),(29, 30),(30, 31),(31, 32),(32, 33),(33, 34),(34, 35),(35, 36),(36, 37),(37, 38),#Face
+            (38, 39),(39, 40),(40, 41),(41, 42),(43, 44),(44, 45),(45, 46),(46, 47),(48, 49),(49, 50),(50, 51),(51, 52),#Face
+            (53, 54),(54, 55),(55, 56),(57, 58),(58, 59),(59, 60),(60, 61),(62, 63),(63, 64),(64, 65),(65, 66),(66, 67),#Face
+            (68, 69),(69, 70),(70, 71),(71, 72),(72, 73),(74, 75),(75, 76),(76, 77),(77, 78),(78, 79),(79, 80),(80, 81),#Face
+            (81, 82),(82, 83),(83, 84),(84, 85),(85, 86),(86, 87),(87, 88),(88, 89),(89, 90),(90, 91),(91, 92),(92, 93),#Face
+            (94,95),(95,96),(96,97),(97,98),(94,99),(99,100),(100,101),(101,102),(94,103),(103,104),(104,105),#LeftHand
+            (105,106),(94,107),(107,108),(108,109),(109,110),(94,111),(111,112),(112,113),(113,114),#LeftHand
+            (115,116),(116,117),(117,118),(118,119),(115,120),(120,121),(121,122),(122,123),(115,124),(124,125),#RightHand
+            (125,126),(126,127),(115,128),(128,129),(129,130),(130,131),(115,132),(132,133),(133,134),(134,135)#RightHand
+        ]
+        p_color = [(0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0),  # Nose, LEye, REye, LEar, REar
+                   (77, 255, 255), (77, 255, 204), (77, 204, 255), (191, 255, 77), (77, 191, 255), (191, 255, 77),  # LShoulder, RShoulder, LElbow, RElbow, LWrist, RWrist
+                   (204, 77, 255), (77, 255, 204), (191, 77, 255), (77, 255, 191), (127, 77, 255), (77, 255, 127),  # LHip, RHip, LKnee, Rknee, LAnkle, RAnkle, Neck
+                   (77, 255, 255), (0, 255, 255), (77, 204, 255),  # head, neck, shoulder
+                   (0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0), (77, 255, 255)] # foot
+    
+        line_color = [(0, 215, 255), (0, 255, 204), (0, 134, 255), (0, 255, 50),
+                      (0, 255, 102), (77, 255, 222), (77, 196, 255), (77, 135, 255), (191, 255, 77), (77, 255, 77),
+                      (77, 191, 255), (204, 77, 255), (77, 222, 255), (255, 156, 127),
+                      (0, 127, 255), (255, 127, 77), (0, 77, 255), (255, 77, 36), 
+                      (0, 77, 255), (0, 77, 255), (0, 77, 255), (0, 77, 255), (255, 156, 127), (255, 156, 127)]
+    elif kp_num == 26:
+        l_pair = [
+            (0, 1), (0, 2), (1, 3), (2, 4),  # Head
+            (5, 18), (6, 18), (5, 7), (7, 9), (6, 8), (8, 10),# Body
+            (17, 18), (18, 19), (19, 11), (19, 12),
+            (11, 13), (12, 14), (13, 15), (14, 16),
+            (20, 24), (21, 25), (23, 25), (22, 24), (15, 24), (16, 25),# Foot
+        ]
+        p_color = [(0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0),  # Nose, LEye, REye, LEar, REar
+                   (77, 255, 255), (77, 255, 204), (77, 204, 255), (191, 255, 77), (77, 191, 255), (191, 255, 77),  # LShoulder, RShoulder, LElbow, RElbow, LWrist, RWrist
+                   (204, 77, 255), (77, 255, 204), (191, 77, 255), (77, 255, 191), (127, 77, 255), (77, 255, 127),  # LHip, RHip, LKnee, Rknee, LAnkle, RAnkle, Neck
+                   (77, 255, 255), (0, 255, 255), (77, 204, 255),  # head, neck, shoulder
+                   (0, 255, 255), (0, 191, 255), (0, 255, 102), (0, 77, 255), (0, 255, 0), (77, 255, 255)] # foot
+    
+        line_color = [(0, 215, 255), (0, 255, 204), (0, 134, 255), (0, 255, 50),
+                      (0, 255, 102), (77, 255, 222), (77, 196, 255), (77, 135, 255), (191, 255, 77), (77, 255, 77),
+                      (77, 191, 255), (204, 77, 255), (77, 222, 255), (255, 156, 127),
+                      (0, 127, 255), (255, 127, 77), (0, 77, 255), (255, 77, 36), 
+                      (0, 77, 255), (0, 77, 255), (0, 77, 255), (0, 77, 255), (255, 156, 127), (255, 156, 127)]
+    else:
+        raise NotImplementedError
+    # im_name = os.path.basename(im_res['imgname'])
+    img = frame.copy()
+    height, width = img.shape[:2]
+    for human in im_res['result']:
+        part_line = {}
+        kp_preds = human['keypoints']
+        kp_scores = human['kp_score']
+        if kp_num == 17:
+            kp_preds = torch.cat((kp_preds, torch.unsqueeze((kp_preds[5, :] + kp_preds[6, :]) / 2, 0)))
+            kp_scores = torch.cat((kp_scores, torch.unsqueeze((kp_scores[5, :] + kp_scores[6, :]) / 2, 0)))
+        if opt.tracking:
+            color = get_color_fast(int(abs(human['idx'])))
+        else:
+            color = BLUE
+
+        # Draw bboxes
+        if opt.showbox:
+            if 'box' in human.keys():
+                bbox = human['box']
+                bbox = [bbox[0], bbox[0]+bbox[2], bbox[1], bbox[1]+bbox[3]]#xmin,xmax,ymin,ymax
+            else:
+                from trackers.PoseFlow.poseflow_infer import get_box
+                keypoints = []
+                for n in range(kp_scores.shape[0]):
+                    keypoints.append(float(kp_preds[n, 0]))
+                    keypoints.append(float(kp_preds[n, 1]))
+                    keypoints.append(float(kp_scores[n]))
+                bbox = get_box(keypoints, height, width)
+            # color = get_color_fast(int(abs(human['idx'][0][0])))
+            cv2.rectangle(img, (int(bbox[0]), int(bbox[2])), (int(bbox[1]),int(bbox[3])), color, 1)
+            if opt.tracking:
+                cv2.putText(img, str(human['idx']), (int(bbox[0]), int((bbox[2] + 26))), DEFAULT_FONT, 1, BLACK, 2)
+
+        # Draw keypoints
+        vis_thres = 0.05 if kp_num == 136 else 0.4
+        for n in range(kp_scores.shape[0]):
+            if kp_scores[n] <= vis_thres:
+                continue
+            cor_x, cor_y = int(kp_preds[n, 0]), int(kp_preds[n, 1])
+            part_line[n] = (int(cor_x), int(cor_y))
+            bg = img.copy()
+            if n < len(p_color):
+                if opt.tracking:
+                    cv2.circle(bg, (int(cor_x), int(cor_y)), 2, color, -1)
+                else:
+                    cv2.circle(bg, (int(cor_x), int(cor_y)), 2, p_color[n], -1)
+            else:
+                cv2.circle(bg, (int(cor_x), int(cor_y)), 1, (255,255,255), 2)
+            # Now create a mask of logo and create its inverse mask also
+            if n < len(p_color):
+                transparency = float(max(0, min(1, kp_scores[n])))
+            else:
+                transparency = float(max(0, min(1, kp_scores[n]*2)))
+            img = cv2.addWeighted(bg, transparency, img, 1 - transparency, 0)
+        # Draw limbs
+        for i, (start_p, end_p) in enumerate(l_pair):
+            if start_p in part_line and end_p in part_line:
+                start_xy = part_line[start_p]
+                end_xy = part_line[end_p]
+                bg = img.copy()
+
+                X = (start_xy[0], end_xy[0])
+                Y = (start_xy[1], end_xy[1])
+                mX = np.mean(X)
+                mY = np.mean(Y)
+                length = ((Y[0] - Y[1]) ** 2 + (X[0] - X[1]) ** 2) ** 0.5
+                angle = math.degrees(math.atan2(Y[0] - Y[1], X[0] - X[1]))
+                stickwidth = (kp_scores[start_p] + kp_scores[end_p]) + 1
+                polygon = cv2.ellipse2Poly((int(mX), int(mY)), (int(length/2), int(stickwidth)), int(angle), 0, 360, 1)
+                if i < len(line_color):
+                    if opt.tracking:
+                        cv2.fillConvexPoly(bg, polygon, color)
+                    else:
+                        cv2.fillConvexPoly(bg, polygon, line_color[i])
+                else:
+                    cv2.line(bg, start_xy, end_xy, (255,255,255), 1)
+                if n < len(p_color):
+                    transparency = float(max(0, min(1, 0.5 * (kp_scores[start_p] + kp_scores[end_p])-0.1)))
+                else:
+                    transparency = float(max(0, min(1, (kp_scores[start_p] + kp_scores[end_p]))))
+
+                #transparency = float(max(0, min(1, 0.5 * (kp_scores[start_p] + kp_scores[end_p])-0.1)))
+                img = cv2.addWeighted(bg, transparency, img, 1 - transparency, 0)
+    return img
+
+
+def getTime(time1=0):
+    if not time1:
+        return time.time()
+    else:
+        interval = time.time() - time1
+        return time.time(), interval

alphapose/utils/webcam_detector.py

@@ -0,0 +1,194 @@
+from itertools import count
+from threading import Thread
+from queue import Queue
+
+import cv2
+import numpy as np
+
+import torch
+import torch.multiprocessing as mp
+
+from alphapose.utils.presets import SimpleTransform
+
+
+class WebCamDetectionLoader():
+    def __init__(self, input_source, detector, cfg, opt, queueSize=1):
+        self.cfg = cfg
+        self.opt = opt
+
+        stream = cv2.VideoCapture(int(input_source))
+        assert stream.isOpened(), 'Cannot capture source'
+        self.path = input_source
+        self.fourcc = int(stream.get(cv2.CAP_PROP_FOURCC))
+        self.fps = stream.get(cv2.CAP_PROP_FPS)
+        self.frameSize = (int(stream.get(cv2.CAP_PROP_FRAME_WIDTH)), int(stream.get(cv2.CAP_PROP_FRAME_HEIGHT)))
+        self.videoinfo = {'fourcc': self.fourcc, 'fps': self.fps, 'frameSize': self.frameSize}
+        stream.release()
+
+        self.detector = detector
+
+        self._input_size = cfg.DATA_PRESET.IMAGE_SIZE
+        self._output_size = cfg.DATA_PRESET.HEATMAP_SIZE
+
+        self._sigma = cfg.DATA_PRESET.SIGMA
+
+        if cfg.DATA_PRESET.TYPE == 'simple':
+            self.transformation = SimpleTransform(
+                self, scale_factor=0,
+                input_size=self._input_size,
+                output_size=self._output_size,
+                rot=0, sigma=self._sigma,
+                train=False, add_dpg=False)
+
+        # initialize the queue used to store data
+        """
+        pose_queue: the buffer storing post-processed cropped human image for pose estimation
+        """
+        if opt.sp:
+            self._stopped = False
+            self.pose_queue = Queue(maxsize=queueSize)
+        else:
+            self._stopped = mp.Value('b', False)
+            self.pose_queue = mp.Queue(maxsize=queueSize)
+
+    def start_worker(self, target):
+        if self.opt.sp:
+            p = Thread(target=target, args=())
+        else:
+            p = mp.Process(target=target, args=())
+        # p.daemon = True
+        p.start()
+        return p
+
+    def start(self):
+        # start a thread to pre process images for object detection
+        image_preprocess_worker = self.start_worker(self.frame_preprocess)
+        return [image_preprocess_worker]
+
+    def stop(self):
+        # clear queues
+        self.clear_queues()
+
+    def terminate(self):
+        if self.opt.sp:
+            self._stopped = True
+        else:
+            self._stopped.value = True
+        self.stop()
+
+    def clear_queues(self):
+        self.clear(self.pose_queue)
+
+    def clear(self, queue):
+        while not queue.empty():
+            queue.get()
+
+    def wait_and_put(self, queue, item):
+        if not self.stopped:
+            queue.put(item)
+
+    def wait_and_get(self, queue):
+        if not self.stopped:
+            return queue.get()
+
+    def frame_preprocess(self):
+        stream = cv2.VideoCapture(self.path)
+        assert stream.isOpened(), 'Cannot capture source'
+
+        # keep looping infinitely
+        for i in count():
+            if self.stopped:
+                stream.release()
+                return
+            if not self.pose_queue.full():
+                # otherwise, ensure the queue has room in it
+                (grabbed, frame) = stream.read()
+                # if the `grabbed` boolean is `False`, then we have
+                # reached the end of the video file
+                if not grabbed:
+                    self.wait_and_put(self.pose_queue, (None, None, None, None, None, None, None))
+                    stream.release()
+                    return
+
+                # expected frame shape like (1,3,h,w) or (3,h,w)
+                img_k = self.detector.image_preprocess(frame)
+
+                if isinstance(img_k, np.ndarray):
+                    img_k = torch.from_numpy(img_k)
+                # add one dimension at the front for batch if image shape (3,h,w)
+                if img_k.dim() == 3:
+                    img_k = img_k.unsqueeze(0)
+
+                im_dim_list_k = frame.shape[1], frame.shape[0]
+
+                orig_img = frame[:, :, ::-1]
+                im_name = str(i) + '.jpg'
+                # im_dim_list = im_dim_list_k
+
+                with torch.no_grad():
+                    # Record original image resolution
+                    im_dim_list_k = torch.FloatTensor(im_dim_list_k).repeat(1, 2)
+                img_det = self.image_detection((img_k, orig_img, im_name, im_dim_list_k))
+                self.image_postprocess(img_det)
+
+    def image_detection(self, inputs):
+        img, orig_img, im_name, im_dim_list = inputs
+        if img is None or self.stopped:
+            return (None, None, None, None, None, None, None)
+
+        with torch.no_grad():
+            dets = self.detector.images_detection(img, im_dim_list)
+            if isinstance(dets, int) or dets.shape[0] == 0:
+                return (orig_img, im_name, None, None, None, None, None)
+            if isinstance(dets, np.ndarray):
+                dets = torch.from_numpy(dets)
+            dets = dets.cpu()
+            boxes = dets[:, 1:5]
+            scores = dets[:, 5:6]
+            if self.opt.tracking:
+                ids = dets[:, 6:7]
+            else:
+                ids = torch.zeros(scores.shape)
+
+        boxes_k = boxes[dets[:, 0] == 0]
+        if isinstance(boxes_k, int) or boxes_k.shape[0] == 0:
+            return (orig_img, im_name, None, None, None, None, None)
+        inps = torch.zeros(boxes_k.size(0), 3, *self._input_size)
+        cropped_boxes = torch.zeros(boxes_k.size(0), 4)
+        return (orig_img, im_name, boxes_k, scores[dets[:, 0] == 0], ids[dets[:, 0] == 0], inps, cropped_boxes)
+
+    def image_postprocess(self, inputs):
+        with torch.no_grad():
+            (orig_img, im_name, boxes, scores, ids, inps, cropped_boxes) = inputs
+            if orig_img is None or self.stopped:
+                self.wait_and_put(self.pose_queue, (None, None, None, None, None, None, None))
+                return
+            if boxes is None or boxes.nelement() == 0:
+                self.wait_and_put(self.pose_queue, (None, orig_img, im_name, boxes, scores, ids, None))
+                return
+            # imght = orig_img.shape[0]
+            # imgwidth = orig_img.shape[1]
+            for i, box in enumerate(boxes):
+                inps[i], cropped_box = self.transformation.test_transform(orig_img, box)
+                cropped_boxes[i] = torch.FloatTensor(cropped_box)
+
+            # inps, cropped_boxes = self.transformation.align_transform(orig_img, boxes)
+
+            self.wait_and_put(self.pose_queue, (inps, orig_img, im_name, boxes, scores, ids, cropped_boxes))
+
+    def read(self):
+        return self.wait_and_get(self.pose_queue)
+
+    @property
+    def stopped(self):
+        if self.opt.sp:
+            return self._stopped
+        else:
+            return self._stopped.value
+
+    @property
+    def joint_pairs(self):
+        """Joint pairs which defines the pairs of joint to be swapped
+        when the image is flipped horizontally."""
+        return [[1, 2], [3, 4], [5, 6], [7, 8],
+                [9, 10], [11, 12], [13, 14], [15, 16]]

alphapose/utils/writer.py

@@ -0,0 +1,208 @@
+import os
+import time
+from threading import Thread
+from queue import Queue
+
+import cv2
+import numpy as np
+import torch
+import torch.multiprocessing as mp
+
+from alphapose.utils.transforms import get_func_heatmap_to_coord
+from alphapose.utils.pPose_nms import pose_nms, write_json
+
+DEFAULT_VIDEO_SAVE_OPT = {
+    'savepath': 'examples/res/1.mp4',
+    'fourcc': cv2.VideoWriter_fourcc(*'mp4v'),
+    'fps': 25,
+    'frameSize': (640, 480)
+}
+
+EVAL_JOINTS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
+
+
+class DataWriter():
+    def __init__(self, cfg, opt, save_video=False,
+                 video_save_opt=DEFAULT_VIDEO_SAVE_OPT,
+                 queueSize=1024):
+        self.cfg = cfg
+        self.opt = opt
+        self.video_save_opt = video_save_opt
+
+        self.eval_joints = EVAL_JOINTS
+        self.save_video = save_video
+        self.heatmap_to_coord = get_func_heatmap_to_coord(cfg)
+        # initialize the queue used to store frames read from
+        # the video file
+        if opt.sp:
+            self.result_queue = Queue(maxsize=queueSize)
+        else:
+            self.result_queue = mp.Queue(maxsize=queueSize)
+
+        if opt.save_img:
+            if not os.path.exists(opt.outputpath + '/vis'):
+                os.mkdir(opt.outputpath + '/vis')
+
+        if opt.pose_flow:
+            from trackers.PoseFlow.poseflow_infer import PoseFlowWrapper
+            self.pose_flow_wrapper = PoseFlowWrapper(save_path=os.path.join(opt.outputpath, 'poseflow'))
+
+    def start_worker(self, target):
+        if self.opt.sp:
+            p = Thread(target=target, args=())
+        else:
+            p = mp.Process(target=target, args=())
+        # p.daemon = True
+        p.start()
+        return p
+
+    def start(self):
+        # start a thread to read pose estimation results per frame
+        self.result_worker = self.start_worker(self.update)
+        return self
+
+    def update(self):
+        final_result = []
+        norm_type = self.cfg.LOSS.get('NORM_TYPE', None)
+        hm_size = self.cfg.DATA_PRESET.HEATMAP_SIZE
+        if self.save_video:
+            # initialize the file video stream, adapt ouput video resolution to original video
+            stream = cv2.VideoWriter(*[self.video_save_opt[k] for k in ['savepath', 'fourcc', 'fps', 'frameSize']])
+            if not stream.isOpened():
+                print("Try to use other video encoders...")
+                ext = self.video_save_opt['savepath'].split('.')[-1]
+                fourcc, _ext = self.recognize_video_ext(ext)
+                self.video_save_opt['fourcc'] = fourcc
+                self.video_save_opt['savepath'] = self.video_save_opt['savepath'][:-4] + _ext
+                stream = cv2.VideoWriter(*[self.video_save_opt[k] for k in ['savepath', 'fourcc', 'fps', 'frameSize']])
+            assert stream.isOpened(), 'Cannot open video for writing'
+        # keep looping infinitelyd
+        while True:
+            # ensure the queue is not empty and get item
+            (boxes, scores, ids, hm_data, cropped_boxes, orig_img, im_name) = self.wait_and_get(self.result_queue)
+            if orig_img is None:
+                # if the thread indicator variable is set (img is None), stop the thread
+                if self.save_video:
+                    stream.release()
+                write_json(final_result, self.opt.outputpath, save_file_name=self.opt.video.split('.')[0].split('/')[-1], form=self.opt.format, for_eval=self.opt.eval)
+                print("Results have been written to json.")
+                return
+            # image channel RGB->BGR
+            orig_img = np.array(orig_img, dtype=np.uint8)[:, :, ::-1]
+            if boxes is None or len(boxes) == 0:
+                if self.opt.save_img or self.save_video or self.opt.vis:
+                    self.write_image(orig_img, im_name, stream=stream if self.save_video else None)
+            else:
+                # location prediction (n, kp, 2) | score prediction (n, kp, 1)
+                assert hm_data.dim() == 4
+                #pred = hm_data.cpu().data.numpy()
+
+                if hm_data.size()[1] == 136:
+                    self.eval_joints = [*range(0,136)]
+                elif hm_data.size()[1] == 26:
+                    self.eval_joints = [*range(0,26)]
+                pose_coords = []
+                pose_scores = []
+                for i in range(hm_data.shape[0]):
+                    bbox = cropped_boxes[i].tolist()
+                    pose_coord, pose_score = self.heatmap_to_coord(hm_data[i][self.eval_joints], bbox, hm_shape=hm_size, norm_type=norm_type)
+                    pose_coords.append(torch.from_numpy(pose_coord).unsqueeze(0))
+                    pose_scores.append(torch.from_numpy(pose_score).unsqueeze(0))
+                preds_img = torch.cat(pose_coords)
+                preds_scores = torch.cat(pose_scores)
+                if not self.opt.pose_track:
+                    boxes, scores, ids, preds_img, preds_scores, pick_ids = \
+                        pose_nms(boxes, scores, ids, preds_img, preds_scores, self.opt.min_box_area)
+
+                _result = []
+                for k in range(len(scores)):
+                    _result.append(
+                        {
+                            'keypoints':preds_img[k],
+                            'kp_score':preds_scores[k],
+                            'proposal_score': torch.mean(preds_scores[k]) + scores[k] + 1.25 * max(preds_scores[k]),
+                            'idx':ids[k],
+                            'box':[boxes[k][0], boxes[k][1], boxes[k][2]-boxes[k][0],boxes[k][3]-boxes[k][1]] 
+                        }
+                    )
+
+                result = {
+                    'imgname': im_name,
+                    'result': _result
+                }
+
+
+                if self.opt.pose_flow:
+                    poseflow_result = self.pose_flow_wrapper.step(orig_img, result)
+                    for i in range(len(poseflow_result)):
+                        result['result'][i]['idx'] = poseflow_result[i]['idx']
+
+                final_result.append(result)
+                if self.opt.save_img or self.save_video or self.opt.vis:
+                    if hm_data.size()[1] == 49:
+                        from alphapose.utils.vis import vis_frame_dense as vis_frame
+                    elif self.opt.vis_fast:
+                        from alphapose.utils.vis import vis_frame_fast as vis_frame
+                    else:
+                        from alphapose.utils.vis import vis_frame
+                    img = vis_frame(orig_img, result, self.opt)
+                    self.write_image(img, im_name, stream=stream if self.save_video else None)
+
+    def write_image(self, img, im_name, stream=None):
+        if self.opt.vis:
+            cv2.imshow("AlphaPose Demo", img)
+            cv2.waitKey(30)
+        if self.opt.save_img:
+            cv2.imwrite(os.path.join(self.opt.outputpath, 'vis', im_name), img)
+        if self.save_video:
+            stream.write(img)
+
+    def wait_and_put(self, queue, item):
+        queue.put(item)
+
+    def wait_and_get(self, queue):
+        return queue.get()
+
+    def save(self, boxes, scores, ids, hm_data, cropped_boxes, orig_img, im_name):
+        # save next frame in the queue
+        self.wait_and_put(self.result_queue, (boxes, scores, ids, hm_data, cropped_boxes, orig_img, im_name))
+
+    def running(self):
+        # indicate that the thread is still running
+        return not self.result_queue.empty()
+
+    def count(self):
+        # indicate the remaining images
+        return self.result_queue.qsize()
+
+    def stop(self):
+        # indicate that the thread should be stopped
+        self.save(None, None, None, None, None, None, None)
+        self.result_worker.join()
+
+    def terminate(self):
+        # directly terminate
+        self.result_worker.terminate()
+
+    def clear_queues(self):
+        self.clear(self.result_queue)
+        
+    def clear(self, queue):
+        while not queue.empty():
+            queue.get()
+
+    def results(self):
+        # return final result
+        print(self.final_result)
+        return self.final_result
+
+    def recognize_video_ext(self, ext=''):
+        if ext == 'mp4':
+            return cv2.VideoWriter_fourcc(*'mp4v'), '.' + ext
+        elif ext == 'avi':
+            return cv2.VideoWriter_fourcc(*'XVID'), '.' + ext
+        elif ext == 'mov':
+            return cv2.VideoWriter_fourcc(*'XVID'), '.' + ext
+        else:
+            print("Unknow video format {}, will use .mp4 instead of it".format(ext))
+            return cv2.VideoWriter_fourcc(*'mp4v'), '.mp4'

alphapose/version.py

@@ -0,0 +1,5 @@
+# GENERATED VERSION FILE
+# TIME: Thu Jul 29 15:24:22 2021
+
+__version__ = '0.3.0+cbc364f'
+short_version = '0.3.0'

configs/1.txt

@@ -0,0 +1 @@
+1

configs/coco/1.txt

@@ -0,0 +1 @@
+1

configs/coco/resnet/256x192_res152_lr1e-3_1x-duc.yaml

@@ -0,0 +1,73 @@
+DATASET:
+  TRAIN:
+    TYPE: 'Mscoco'
+    ROOT: './data/coco/'
+    IMG_PREFIX: 'train2017'
+    ANN: 'annotations/person_keypoints_train2017.json'
+    AUG:
+      FLIP: true
+      ROT_FACTOR: 40
+      SCALE_FACTOR: 0.3
+      NUM_JOINTS_HALF_BODY: 8
+      PROB_HALF_BODY: -1
+  VAL:
+    TYPE: 'Mscoco'
+    ROOT: './data/coco/'
+    IMG_PREFIX: 'val2017'
+    ANN: 'annotations/person_keypoints_val2017.json'
+  TEST:
+    TYPE: 'Mscoco_det'
+    ROOT: './data/coco/'
+    IMG_PREFIX: 'val2017'
+    DET_FILE: './exp/json/test_det_yolo.json'
+    ANN: 'annotations/person_keypoints_val2017.json'
+DATA_PRESET:
+  TYPE: 'simple'
+  SIGMA: 2
+  NUM_JOINTS: 17
+  IMAGE_SIZE:
+  - 256
+  - 192
+  HEATMAP_SIZE:
+  - 64
+  - 48
+MODEL:
+  TYPE: 'FastPose_DUC'
+  BACKBONE: 'se-resnet'
+  PRETRAINED: ''
+  TRY_LOAD: ''
+  NUM_DECONV_FILTERS:
+  - 256
+  - 256
+  - 256
+  NUM_LAYERS: 152
+  FINAL_CONV_KERNEL: 1
+  STAGE1:
+    NUM_CONV: 4
+  STAGE2:
+    NUM_CONV: 2
+  STAGE3:
+    NUM_CONV: 1
+LOSS:
+  TYPE: 'MSELoss'
+DETECTOR:
+  NAME: 'yolo'
+  CONFIG: 'detector/yolo/cfg/yolov3-spp.cfg'
+  WEIGHTS: 'detector/yolo/data/yolov3-spp.weights'
+  NMS_THRES: 0.6
+  CONFIDENCE: 0.05
+TRAIN:
+  WORLD_SIZE: 4
+  BATCH_SIZE: 32
+  BEGIN_EPOCH: 0
+  END_EPOCH: 200
+  OPTIMIZER: 'adam'
+  LR: 0.001
+  LR_FACTOR: 0.1
+  LR_STEP:
+  - 90
+  - 120
+  DPG_MILESTONE: 140
+  DPG_STEP:
+  - 160
+  - 190

data/video.txt

@@ -0,0 +1 @@
+这个文件夹用来装待处理的视频。

detector/apis.py

@@ -0,0 +1,41 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Chao Xu (xuchao.19962007@sjtu.edu.cn)
+# -----------------------------------------------------
+
+"""API of detector"""
+from abc import ABC, abstractmethod
+
+
+def get_detector(opt=None):
+    if opt.detector == 'yolo':
+        from detector.yolo_api import YOLODetector
+        from detector.yolo_cfg import cfg
+        return YOLODetector(cfg, opt)
+    elif opt.detector == 'tracker':
+        from detector.tracker_api import Tracker
+        from detector.tracker_cfg import cfg
+        return Tracker(cfg, opt)
+    elif opt.detector.startswith('efficientdet_d'):
+        from detector.effdet_api import EffDetDetector
+        from detector.effdet_cfg import cfg
+        return EffDetDetector(cfg, opt)
+    else:
+        raise NotImplementedError
+
+
+class BaseDetector(ABC):
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def image_preprocess(self, img_name):
+        pass
+
+    @abstractmethod
+    def images_detection(self, imgs, orig_dim_list):
+        pass
+
+    @abstractmethod
+    def detect_one_img(self, img_name):
+        pass

detector/effdet_api.py

@@ -0,0 +1,289 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Haoshu Fang (fhaoshu@gmail.com)
+# -----------------------------------------------------
+
+"""API of efficientdet detector"""
+import os
+import sys
+sys.path.insert(0, os.path.dirname(__file__))
+from abc import ABC, abstractmethod
+import platform
+
+import torch
+import numpy as np
+
+from efficientdet.utils import unique, prep_image, prep_frame, bbox_iou
+from efficientdet.effdet import EfficientDet, get_efficientdet_config, DetBenchEval, load_checkpoint
+
+from detector.apis import BaseDetector
+
+try:
+    from apex import amp
+    has_amp = True
+except ImportError:
+    has_amp = False
+
+#only windows visual studio 2013 ~2017 support compile c/cuda extensions
+#If you force to compile extension on Windows and ensure appropriate visual studio
+#is intalled, you can try to use these ext_modules.
+if platform.system() != 'Windows':
+    from detector.nms import nms_wrapper
+
+
+class EffDetDetector(BaseDetector):
+    def __init__(self, cfg, opt=None):
+        super(EffDetDetector, self).__init__()
+
+        self.detector_cfg = cfg
+        self.detector_opt = opt
+        self.model_cfg = get_efficientdet_config(opt.detector)
+        self.model_weights = 'detector/efficientdet/weights/'+opt.detector+'.pth'
+        #Input image dimension, uses model default if empty
+        self.inp_dim = cfg.get('INP_DIM', None) if cfg.get('INP_DIM', None) is not None else self.model_cfg.image_size
+        self.nms_thres = cfg.get('NMS_THRES', 0.6)
+        self.confidence = cfg.get('CONFIDENCE', 0.05)
+        self.num_classes = cfg.get('NUM_CLASSES', 80)
+        self.max_dets = cfg.get('MAX_DETECTIONS', 100)
+        self.model = None
+
+
+    def load_model(self):
+        args = self.detector_opt
+
+        net = EfficientDet(self.model_cfg)
+        load_checkpoint(net, self.model_weights)
+        self.model = DetBenchEval(net, self.model_cfg, nms_thres=self.nms_thres, max_dets=self.max_dets)
+
+        if args:
+            if len(args.gpus) > 1:
+                if has_amp:
+                    print('Using AMP mixed precision.')
+                    self.model = amp.initialize(self.model, opt_level='O1')
+                else:
+                    print('AMP not installed, running network in FP32.')
+
+                self.model = torch.nn.DataParallel(self.model, device_ids=args.gpus).to(args.device)
+            else:
+                self.model.to(args.device)
+        else:
+            if has_amp:
+                print('Using AMP mixed precision.')
+                self.model = amp.initialize(self.model, opt_level='O1')
+            else:
+                print('AMP not installed, running network in FP32.')
+            self.model.cuda()
+
+        net.eval()
+
+    def image_preprocess(self, img_source):
+        """
+        Pre-process the img before fed to the object detection network
+        Input: image name(str) or raw image data(ndarray or torch.Tensor,channel GBR)
+        Output: pre-processed image data(torch.FloatTensor,(1,3,h,w))
+        """
+        if isinstance(img_source, str):
+            img, orig_img, im_dim_list = prep_image(img_source, self.inp_dim)
+        elif isinstance(img_source, torch.Tensor) or isinstance(img_source, np.ndarray):
+            img, orig_img, im_dim_list = prep_frame(img_source, self.inp_dim)
+        else:
+            raise IOError('Unknown image source type: {}'.format(type(img_source)))
+
+        return img
+
+    def images_detection(self, imgs, orig_dim_list):
+        """
+        Feed the img data into object detection network and 
+        collect bbox w.r.t original image size
+        Input: imgs(torch.FloatTensor,(b,3,h,w)): pre-processed mini-batch image input
+               orig_dim_list(torch.FloatTensor, (b,(w,h,w,h))): original mini-batch image size
+        Output: dets(torch.cuda.FloatTensor,(n,(batch_idx,x1,y1,x2,y2,c,s,idx of cls))): human detection results
+        """
+        args = self.detector_opt
+        if not self.model:
+            self.load_model()
+        with torch.no_grad():
+            imgs = imgs.to(args.device) if args else imgs.cuda()
+            scaling_factors = torch.FloatTensor([1./min(self.inp_dim / orig_dim[0], self.inp_dim / orig_dim[1]) for orig_dim in orig_dim_list]).view(-1, 1)
+            scaling_factors = scaling_factors.to(args.device) if args else scaling_factors.cuda()
+            prediction = self.model(imgs, scaling_factors) 
+            #change the pred format to alphapose (nms has already been done in effdeteval model)
+            prediction = prediction.cpu()
+            write = False
+            for index, sample in enumerate(prediction):
+                for det in sample:
+                    score = float(det[4])
+                    if score < .001:  # stop when below this threshold, scores in descending order
+                        break
+                    if int(det[5]) != 1 or score < self.confidence:
+                        continue
+                    det_new = prediction.new(1,8)
+                    det_new[0,0] = index    #index of img
+                    det_new[0,1:3] = det[0:2]  # bbox x1,y1
+                    det_new[0,3:5] = det[0:2] + det[2:4] # bbox x2,y2
+                    det_new[0,6:7] = det[4]  # cls conf
+                    det_new[0,7] = det[5]   # cls idx
+                    if not write:
+                        dets = det_new
+                        write = True
+                    else:
+                        dets = torch.cat((dets, det_new))    
+            if not write:
+                return 0
+
+            orig_dim_list = torch.index_select(orig_dim_list, 0, dets[:, 0].long())
+            for i in range(dets.shape[0]):
+                dets[i, [1, 3]] = torch.clamp(dets[i, [1, 3]], 0.0, orig_dim_list[i, 0])
+                dets[i, [2, 4]] = torch.clamp(dets[i, [2, 4]], 0.0, orig_dim_list[i, 1])
+
+            return dets
+
+    def detect_one_img(self, img_name):
+        """
+        Detect bboxs in one image
+        Input: 'str', full path of image
+        Output: '[{"category_id":1,"score":float,"bbox":[x,y,w,h],"image_id":str},...]',
+        The output results are similar with coco results type, except that image_id uses full path str
+        instead of coco %012d id for generalization. 
+        """
+        args = self.detector_opt
+        _CUDA = True
+        if args:
+            if args.gpus[0] < 0:
+                _CUDA = False
+        if not self.model:
+            self.load_model()
+        if isinstance(self.model, torch.nn.DataParallel):
+            self.model = self.model.module
+        dets_results = []
+        #pre-process(scale, normalize, ...) the image
+        img, orig_img, img_dim_list = prep_image(img_name, self.inp_dim)
+        with torch.no_grad():
+            img_dim_list = torch.FloatTensor([img_dim_list]).repeat(1, 2)
+            img = img.to(args.device) if args else img.cuda()
+            scaling_factor = torch.FloatTensor([1/min(self.inp_dim / orig_dim[0], self.inp_dim / orig_dim[1]) for orig_dim in img_dim_list]).view(-1, 1)
+            scaling_factor = scaling_factor.to(args.device) if args else scaling_factor.cuda()
+            prediction = self.model(img, scaling_factor) 
+            #change the pred format to alphapose (nms has already been done in effdeteval model)
+            prediction = prediction.cpu()
+            write = False
+            for index, sample in enumerate(prediction):
+                for det in sample:
+                    score = float(det[4])
+                    if score < .001:  # stop when below this threshold, scores in descending order
+                        break
+                    if int(det[5]) != 1 or score < self.confidence:
+                        continue
+                    det_new = prediction.new(1,8)
+                    det_new[0,0] = index    #index of img
+                    det_new[0,1:3] = det[0:2]  # bbox x1,y1
+                    det_new[0,3:5] = det[0:2] + det[2:4] # bbox x2,y2
+                    det_new[0,6:7] = det[4]  # cls conf
+                    det_new[0,7] = det[5]   # cls idx
+                    if not write:
+                        dets = det_new
+                        write = True
+                    else:
+                        dets = torch.cat((dets, det_new))          
+            if not write:
+                return None
+
+            img_dim_list = torch.index_select(img_dim_list, 0, dets[:, 0].long())
+            for i in range(dets.shape[0]):
+                dets[i, [1, 3]] = torch.clamp(dets[i, [1, 3]], 0.0, img_dim_list[i, 0])
+                dets[i, [2, 4]] = torch.clamp(dets[i, [2, 4]], 0.0, img_dim_list[i, 1])
+
+                #write results
+                det_dict = {}
+                x = float(dets[i, 1])
+                y = float(dets[i, 2])
+                w = float(dets[i, 3] - dets[i, 1])
+                h = float(dets[i, 4] - dets[i, 2])
+                det_dict["category_id"] = 1
+                det_dict["score"] = float(dets[i, 5])
+                det_dict["bbox"] = [x, y, w, h]
+                det_dict["image_id"] = int(os.path.basename(img_name).split('.')[0])
+                dets_results.append(det_dict)
+
+            return dets_results
+
+
+    def check_detector(self, img_name):
+        """
+        Detect bboxs in one image
+        Input: 'str', full path of image
+        Output: '[{"category_id":1,"score":float,"bbox":[x,y,w,h],"image_id":str},...]',
+        The output results are similar with coco results type, except that image_id uses full path str
+        instead of coco %012d id for generalization. 
+        """
+        args = self.detector_opt
+        _CUDA = True
+        if args:
+            if args.gpus[0] < 0:
+                _CUDA = False
+        if not self.model:
+            self.load_model()
+        if isinstance(self.model, torch.nn.DataParallel):
+            self.model = self.model.module
+        dets_results = []
+        #pre-process(scale, normalize, ...) the image
+        img, orig_img, img_dim_list = prep_image(img_name, self.inp_dim)
+        with torch.no_grad():
+            img_dim_list = torch.FloatTensor([img_dim_list]).repeat(1, 2)
+            img = img.to(args.device) if args else img.cuda()
+            scaling_factor = torch.FloatTensor([1/min(self.inp_dim / orig_dim[0], self.inp_dim / orig_dim[1]) for orig_dim in img_dim_list]).view(-1, 1)
+            scaling_factor = scaling_factor.to(args.device) if args else scaling_factor.cuda()
+            output = self.model(img, scaling_factor) 
+
+            output = output.cpu()
+            for index, sample in enumerate(output):
+                image_id = int(os.path.basename(img_name).split('.')[0])
+                for det in sample:
+                    score = float(det[4])
+                    if score < .001:  # stop when below this threshold, scores in descending order
+                        break
+                    #### uncomment it for only human detection
+                    # if int(det[5]) != 1 or score < self.confidence:
+                    #     continue
+                    coco_det = dict(
+                        image_id=image_id,
+                        bbox=det[0:4].tolist(),
+                        score=score,
+                        category_id=int(det[5]))
+                    dets_results.append(coco_det)
+
+        return dets_results
+
+if __name__ == "__main__":
+#run with python detector/effdet_api.py /DATA1/Benchmark/coco/ efficientdet_d0
+        from pycocotools.coco import COCO
+        from pycocotools.cocoeval import COCOeval
+        from easydict import EasyDict as edict
+        from apis import get_detector
+        from tqdm import tqdm
+        import json
+
+        opt = edict()
+        _coco = COCO(sys.argv[1]+'/annotations/instances_val2017.json')
+        # _coco = COCO(sys.argv[1]+'/annotations/person_keypoints_val2017.json')
+        opt.detector = sys.argv[2]
+        opt.gpus = [0] if torch.cuda.device_count() >= 1 else [-1]
+        opt.device = torch.device("cuda:" + str(opt.gpus[0]) if opt.gpus[0] >= 0 else "cpu")
+        image_ids = sorted(_coco.getImgIds())
+        det_model = get_detector(opt)
+        dets = []
+        for entry in tqdm(_coco.loadImgs(image_ids)):
+            abs_path = os.path.join(
+                sys.argv[1], 'val2017', entry['file_name'])
+            det = det_model.check_detector(abs_path)
+            if det:
+                dets += det
+        result_file = 'results.json'
+        json.dump(dets, open(result_file, 'w'))
+
+        coco_results = _coco.loadRes(result_file)
+        coco_eval = COCOeval(_coco, coco_results, 'bbox')
+        coco_eval.params.imgIds = image_ids  # score only ids we've used
+        coco_eval.evaluate()
+        coco_eval.accumulate()
+        coco_eval.summarize()

detector/effdet_cfg.py

@@ -0,0 +1,8 @@
+from easydict import EasyDict as edict
+
+cfg = edict()
+
+cfg.NMS_THRES =  0.6  # 0.6(0.713) 0.5(0.707)
+cfg.CONFIDENCE = 0.2  # 0.15       0.1
+cfg.NUM_CLASSES = 80
+cfg.MAX_DETECTIONS = 200  # 100

detector/tracker_api.py

@@ -0,0 +1,238 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Haoshu Fang (fhaoshu@gmail.com)
+# -----------------------------------------------------
+
+"""API of yolo tracker"""
+import os
+import sys
+sys.path.insert(0, os.path.dirname(__file__))
+from abc import ABC, abstractmethod
+import platform
+
+import torch
+import numpy as np
+
+from tracker.tracker.multitracker import STrack, joint_stracks, sub_stracks, remove_duplicate_stracks
+
+from tracker.preprocess import prep_image, prep_frame
+from tracker.utils.kalman_filter import KalmanFilter
+from tracker.utils.utils import non_max_suppression, scale_coords
+from tracker.utils.log import logger
+from tracker.tracker import matching
+from tracker.tracker.basetrack import BaseTrack, TrackState
+from tracker.models import Darknet
+
+from detector.apis import BaseDetector
+
+
+class Tracker(BaseDetector):
+    def __init__(self, cfg, opt=None):
+        super(Tracker, self).__init__()
+
+        self.tracker_opt = opt
+        self.model_cfg = cfg.get('CONFIG', 'detector/tracker/cfg/yolov3.cfg')
+        self.model_weights = cfg.get('WEIGHTS', 'detector/tracker/data/jde.1088x608.uncertainty.pt')
+        self.img_size = cfg.get('IMG_SIZE', (1088, 608))
+        self.nms_thres = cfg.get('NMS_THRES', 0.6)
+        self.confidence = cfg.get('CONFIDENCE', 0.05)
+        self.max_time_lost = cfg.get('BUFFER_SIZE', 30) # buffer
+        self.model = None
+        
+        self.tracked_stracks = []  # type: list[STrack]
+        self.lost_stracks = []  # type: list[STrack]
+        self.removed_stracks = []  # type: list[STrack]
+
+        self.frame_id = 0
+        self.emb_dim = None
+        
+
+        self.kalman_filter = KalmanFilter()
+
+    def load_model(self):
+        print('Loading tracking model..')
+        self.model = Darknet(self.model_cfg, self.img_size, nID=14455)
+        # load_darknet_weights(self.model, args.weights)
+        self.model.load_state_dict(torch.load(self.model_weights, map_location='cpu')['model'], strict=False)
+        self.emb_dim = self.model.emb_dim
+
+        if self.tracker_opt:
+            if len(self.tracker_opt.gpus) > 1:
+                self.model = torch.nn.DataParallel(self.model, device_ids=self.tracker_opt.gpus).to(self.tracker_opt.device)
+            else:
+                self.model.to(self.tracker_opt.device)
+        else:
+            self.model.cuda()
+        self.model.eval()
+        print("Network successfully loaded")
+
+        
+
+    def image_preprocess(self, img_source):
+        """
+        Pre-process the img before fed to the object detection network
+        Input: image name(str) or raw image data(ndarray or torch.Tensor,channel GBR)
+        Output: pre-processed image data(torch.FloatTensor,(1,3,h,w))
+        """
+        if isinstance(img_source, str):
+            img, orig_img, im_dim_list = prep_image(img_source, self.img_size)
+        elif isinstance(img_source, torch.Tensor) or isinstance(img_source, np.ndarray):
+            img, orig_img, im_dim_list = prep_frame(img_source, self.img_size)
+        else:
+            raise IOError('Unknown image source type: {}'.format(type(img_source)))
+
+        return img
+
+    def images_detection(self, imgs, orig_dim_list):
+        """
+        Feed the img data into object detection network and 
+        collect bbox w.r.t original image size
+        Input: imgs(torch.FloatTensor,(b,3,h,w)): pre-processed mini-batch image input
+               orig_dim_list(torch.FloatTensor, (b,(w,h,w,h))): original mini-batch image size
+        Output: dets(torch.cuda.FloatTensor,(n,(batch_idx,x1,y1,x2,y2,c,s,idx of cls))): human detection results
+        """
+        args = self.tracker_opt
+        _CUDA = True
+        if args:
+            if args.gpus[0] < 0:
+                _CUDA = False
+        if not self.model:
+            self.load_model()
+        
+        
+        activated_starcks = []
+        refind_stracks = []
+        lost_stracks = []
+        removed_stracks = []
+
+        ''' Step 1: Network forward, get detections & embeddings'''
+        with torch.no_grad():
+            imgs = imgs.to(args.device) if args else imgs.cuda()
+            pred = self.model(imgs)
+
+        if len(pred) > 0:
+            dets = non_max_suppression(pred, self.confidence, self.nms_thres)
+
+        output_stracks = []
+        for image_i in range(len(imgs)):
+            self.frame_id += 1
+            if dets[image_i] is not None:
+                det_i = scale_coords(self.img_size, dets[image_i], orig_dim_list[image_i])
+                '''Detections'''
+                detections = [STrack(STrack.tlbr_to_tlwh(tlbrs[:4]), tlbrs[4], f.numpy(), 30) for
+                              (tlbrs, f) in zip(det_i[:, :5], det_i[:, -self.emb_dim:])]
+            else:
+                detections = []
+
+
+            ''' Add newly detected tracklets to tracked_stracks'''
+            unconfirmed = []
+            tracked_stracks = []  # type: list[STrack]
+            for track in self.tracked_stracks:
+                if not track.is_activated:
+                    unconfirmed.append(track)
+                else:
+                    tracked_stracks.append(track)
+
+            ''' Step 2: First association, with embedding'''
+            strack_pool = joint_stracks(tracked_stracks, self.lost_stracks)
+            # Predict the current location with KF
+            for strack in strack_pool:
+                strack.predict()
+
+            dists = matching.embedding_distance(strack_pool, detections)
+            dists = matching.gate_cost_matrix(self.kalman_filter, dists, strack_pool, detections)
+            matches, u_track, u_detection = matching.linear_assignment(dists, thresh=0.7)
+
+            for itracked, idet in matches:
+                track = strack_pool[itracked]
+                det = detections[idet]
+                if track.state == TrackState.Tracked:
+                    track.update(detections[idet], self.frame_id)
+                    activated_starcks.append(track)
+                else:
+                    track.re_activate(det, self.frame_id, new_id=False)
+                    refind_stracks.append(track)
+
+            ''' Step 3: Second association, with IOU'''
+            detections = [detections[i] for i in u_detection]
+            r_tracked_stracks = [strack_pool[i] for i in u_track if strack_pool[i].state==TrackState.Tracked ]
+            dists = matching.iou_distance(r_tracked_stracks, detections)
+            matches, u_track, u_detection = matching.linear_assignment(dists, thresh=0.5)
+            
+            for itracked, idet in matches:
+                track = r_tracked_stracks[itracked]
+                det = detections[idet]
+                if track.state == TrackState.Tracked:
+                    track.update(det, self.frame_id)
+                    activated_starcks.append(track)
+                else:
+                    track.re_activate(det, self.frame_id, new_id=False)
+                    refind_stracks.append(track)
+
+            for it in u_track:
+                track = r_tracked_stracks[it]
+                if not track.state == TrackState.Lost:
+                    track.mark_lost()
+                    lost_stracks.append(track)
+
+            '''Deal with unconfirmed tracks, usually tracks with only one beginning frame'''
+            detections = [detections[i] for i in u_detection]
+            dists = matching.iou_distance(unconfirmed, detections)
+            matches, u_unconfirmed, u_detection = matching.linear_assignment(dists, thresh=0.7)
+            for itracked, idet in matches:
+                unconfirmed[itracked].update(detections[idet], self.frame_id)
+                activated_starcks.append(unconfirmed[itracked])
+            for it in u_unconfirmed:
+                track = unconfirmed[it]
+                track.mark_removed()
+                removed_stracks.append(track)
+
+            """ Step 4: Init new stracks"""
+            for inew in u_detection:
+                track = detections[inew]
+                if track.score < self.confidence:
+                    continue
+                track.activate(self.kalman_filter, self.frame_id)
+                activated_starcks.append(track)
+
+            """ Step 5: Update state"""
+            for track in self.lost_stracks:
+                if self.frame_id - track.end_frame > self.max_time_lost:
+                    track.mark_removed()
+                    removed_stracks.append(track)
+
+            self.tracked_stracks = [t for t in self.tracked_stracks if t.state == TrackState.Tracked]
+            self.tracked_stracks = joint_stracks(self.tracked_stracks, activated_starcks)
+            self.tracked_stracks = joint_stracks(self.tracked_stracks, refind_stracks)
+            # self.lost_stracks = [t for t in self.lost_stracks if t.state == TrackState.Lost]  # type: list[STrack]
+            self.lost_stracks = sub_stracks(self.lost_stracks, self.tracked_stracks)
+            self.lost_stracks.extend(lost_stracks)
+            self.lost_stracks = sub_stracks(self.lost_stracks, self.removed_stracks)
+            self.removed_stracks.extend(removed_stracks)
+            self.tracked_stracks, self.lost_stracks = remove_duplicate_stracks(self.tracked_stracks, self.lost_stracks)
+
+            if self.tracker_opt.debug:
+                logger.debug('===========Frame {}=========='.format(self.frame_id))
+                logger.debug('Activated: {}'.format([track.track_id for track in activated_starcks]))
+                logger.debug('Refind: {}'.format([track.track_id for track in refind_stracks]))
+                logger.debug('Lost: {}'.format([track.track_id for track in lost_stracks]))
+                logger.debug('Removed: {}'.format([track.track_id for track in removed_stracks]))
+
+            # Add tracks to outputs
+            for t in self.tracked_stracks:
+                tlwh = t.tlwh
+                tid = t.track_id
+                tlbr = t.tlbr
+                ts = t.score
+                if tlwh[2] * tlwh[3] > self.tracker_opt.min_box_area:
+                    res = torch.tensor([image_i, tlbr[0], tlbr[1], tlbr[2], tlbr[3], ts, tid])
+                output_stracks.append(res)
+
+        if len(output_stracks) == 0:
+            return 0
+            
+        return torch.stack(output_stracks)
+
+    def detect_one_img(self, img_name):
+        pass

detector/tracker_cfg.py

@@ -0,0 +1,9 @@
+from easydict import EasyDict as edict
+
+cfg = edict()
+cfg.CONFIG = 'detector/tracker/cfg/yolov3.cfg'
+cfg.WEIGHTS = 'detector/tracker/data/jde.1088x608.uncertainty.pt'
+cfg.IMG_SIZE =  (1088, 608)
+cfg.NMS_THRES =  0.6
+cfg.CONFIDENCE = 0.4
+cfg.BUFFER_SIZE = 30 # frame buffer

detector/yolo_api.py

@@ -0,0 +1,292 @@
+# -----------------------------------------------------
+# Copyright (c) Shanghai Jiao Tong University. All rights reserved.
+# Written by Chao Xu (xuchao.19962007@sjtu.edu.cn)
+# -----------------------------------------------------
+
+"""API of yolo detector"""
+import os
+import sys
+sys.path.insert(0, os.path.dirname(__file__))
+from abc import ABC, abstractmethod
+import platform
+
+import torch
+import numpy as np
+
+from yolo.preprocess import prep_image, prep_frame
+from yolo.darknet import Darknet
+from yolo.util import unique
+from yolo.bbox import bbox_iou
+
+from detector.apis import BaseDetector
+
+#only windows visual studio 2013 ~2017 support compile c/cuda extensions
+#If you force to compile extension on Windows and ensure appropriate visual studio
+#is intalled, you can try to use these ext_modules.
+if platform.system() != 'Windows':
+    from detector.nms import nms_wrapper
+
+
+class YOLODetector(BaseDetector):
+    def __init__(self, cfg, opt=None):
+        super(YOLODetector, self).__init__()
+
+        self.detector_cfg = cfg
+        self.detector_opt = opt
+        self.model_cfg = cfg.get('CONFIG', 'detector/yolo/cfg/yolov3-spp.cfg')
+        self.model_weights = cfg.get('WEIGHTS', 'detector/yolo/data/yolov3-spp.weights')
+        self.inp_dim = cfg.get('INP_DIM', 608)
+        self.nms_thres = cfg.get('NMS_THRES', 0.6)
+        self.confidence = 0.3 if (False if not hasattr(opt, 'tracking') else opt.tracking) else cfg.get('CONFIDENCE', 0.05)
+        self.num_classes = cfg.get('NUM_CLASSES', 80)
+        self.model = None
+
+    def load_model(self):
+        args = self.detector_opt
+
+        print('Loading YOLO model..')
+        self.model = Darknet(self.model_cfg)
+        self.model.load_weights(self.model_weights)
+        self.model.net_info['height'] = self.inp_dim
+        
+
+        if args:
+            if len(args.gpus) > 1:
+                self.model = torch.nn.DataParallel(self.model, device_ids=args.gpus).to(args.device)
+            else:
+                self.model.to(args.device)
+        else:
+            self.model.cuda()
+        self.model.eval()
+
+    def image_preprocess(self, img_source):
+        """
+        Pre-process the img before fed to the object detection network
+        Input: image name(str) or raw image data(ndarray or torch.Tensor,channel GBR)
+        Output: pre-processed image data(torch.FloatTensor,(1,3,h,w))
+        """
+        if isinstance(img_source, str):
+            img, orig_img, im_dim_list = prep_image(img_source, self.inp_dim)
+        elif isinstance(img_source, torch.Tensor) or isinstance(img_source, np.ndarray):
+            img, orig_img, im_dim_list = prep_frame(img_source, self.inp_dim)
+        else:
+            raise IOError('Unknown image source type: {}'.format(type(img_source)))
+
+        return img
+
+    def images_detection(self, imgs, orig_dim_list):
+        """
+        Feed the img data into object detection network and 
+        collect bbox w.r.t original image size
+        Input: imgs(torch.FloatTensor,(b,3,h,w)): pre-processed mini-batch image input
+               orig_dim_list(torch.FloatTensor, (b,(w,h,w,h))): original mini-batch image size
+        Output: dets(torch.cuda.FloatTensor,(n,(batch_idx,x1,y1,x2,y2,c,s,idx of cls))): human detection results
+        """
+        args = self.detector_opt
+        _CUDA = True
+        if args:
+            if args.gpus[0] < 0:
+                _CUDA = False
+        if not self.model:
+            self.load_model()
+        with torch.no_grad():
+            imgs = imgs.to(args.device) if args else imgs.cuda()
+            prediction = self.model(imgs, args=args) 
+            #do nms to the detection results, only human category is left
+            dets = self.dynamic_write_results(prediction, self.confidence, 
+                                              self.num_classes, nms=True, 
+                                              nms_conf=self.nms_thres)
+            if isinstance(dets, int) or dets.shape[0] == 0:
+                return 0
+            dets = dets.cpu()
+
+            orig_dim_list = torch.index_select(orig_dim_list, 0, dets[:, 0].long())
+            scaling_factor = torch.min(self.inp_dim / orig_dim_list, 1)[0].view(-1, 1)
+            dets[:, [1, 3]] -= (self.inp_dim - scaling_factor * orig_dim_list[:, 0].view(-1, 1)) / 2
+            dets[:, [2, 4]] -= (self.inp_dim - scaling_factor * orig_dim_list[:, 1].view(-1, 1)) / 2
+            dets[:, 1:5] /= scaling_factor
+            for i in range(dets.shape[0]):
+                dets[i, [1, 3]] = torch.clamp(dets[i, [1, 3]], 0.0, orig_dim_list[i, 0])
+                dets[i, [2, 4]] = torch.clamp(dets[i, [2, 4]], 0.0, orig_dim_list[i, 1])
+
+            return dets
+
+    def dynamic_write_results(self, prediction, confidence, num_classes, nms=True, nms_conf=0.4):
+        prediction_bak = prediction.clone()
+        dets = self.write_results(prediction.clone(), confidence, num_classes, nms, nms_conf)
+        if isinstance(dets, int):
+            return dets
+
+        if dets.shape[0] > 100:
+            nms_conf -= 0.05
+            dets = self.write_results(prediction_bak.clone(), confidence, num_classes, nms, nms_conf)
+
+        return dets
+
+    def write_results(self, prediction, confidence, num_classes, nms=True, nms_conf=0.4):
+        args = self.detector_opt
+        #prediction: (batchsize, num of objects, (xc,yc,w,h,box confidence, 80 class scores))
+        conf_mask = (prediction[:, :, 4] > confidence).float().float().unsqueeze(2)
+        prediction = prediction * conf_mask
+
+        try:
+            ind_nz = torch.nonzero(prediction[:,:,4]).transpose(0,1).contiguous()
+        except:
+            return 0
+
+        #the 3rd channel of prediction: (xc,yc,w,h)->(x1,y1,x2,y2)
+        box_a = prediction.new(prediction.shape)
+        box_a[:,:,0] = (prediction[:,:,0] - prediction[:,:,2]/2)
+        box_a[:,:,1] = (prediction[:,:,1] - prediction[:,:,3]/2)
+        box_a[:,:,2] = (prediction[:,:,0] + prediction[:,:,2]/2) 
+        box_a[:,:,3] = (prediction[:,:,1] + prediction[:,:,3]/2)
+        prediction[:,:,:4] = box_a[:,:,:4]
+
+        batch_size = prediction.size(0)
+
+        output = prediction.new(1, prediction.size(2) + 1)
+        write = False
+        num = 0
+        for ind in range(batch_size):
+            #select the image from the batch
+            image_pred = prediction[ind]
+
+            #Get the class having maximum score, and the index of that class
+            #Get rid of num_classes softmax scores 
+            #Add the class index and the class score of class having maximum score
+            max_conf, max_conf_score = torch.max(image_pred[:,5:5+ num_classes], 1)
+            max_conf = max_conf.float().unsqueeze(1)
+            max_conf_score = max_conf_score.float().unsqueeze(1)
+            seq = (image_pred[:,:5], max_conf, max_conf_score)
+            #image_pred:(n,(x1,y1,x2,y2,c,s,idx of cls))
+            image_pred = torch.cat(seq, 1)
+
+            #Get rid of the zero entries
+            non_zero_ind =  (torch.nonzero(image_pred[:,4]))
+
+            image_pred_ = image_pred[non_zero_ind.squeeze(),:].view(-1,7)
+
+            #Get the various classes detected in the image
+            try:
+                img_classes = unique(image_pred_[:,-1])
+            except:
+                continue
+
+            #WE will do NMS classwise
+            #print(img_classes)
+            for cls in img_classes:
+                if cls != 0:
+                    continue
+                #get the detections with one particular class
+                cls_mask = image_pred_*(image_pred_[:,-1] == cls).float().unsqueeze(1)
+                class_mask_ind = torch.nonzero(cls_mask[:,-2]).squeeze()
+
+                image_pred_class = image_pred_[class_mask_ind].view(-1,7)
+
+                #sort the detections such that the entry with the maximum objectness
+                #confidence is at the top
+                conf_sort_index = torch.sort(image_pred_class[:,4], descending = True )[1]
+                image_pred_class = image_pred_class[conf_sort_index]
+                idx = image_pred_class.size(0)
+
+                #if nms has to be done
+                if nms:
+                    if platform.system() != 'Windows':
+                        #We use faster rcnn implementation of nms (soft nms is optional)
+                        nms_op = getattr(nms_wrapper, 'nms')
+                        #nms_op input:(n,(x1,y1,x2,y2,c))
+                        #nms_op output: input[inds,:], inds
+                        _, inds = nms_op(image_pred_class[:,:5], nms_conf)
+
+                        image_pred_class = image_pred_class[inds]
+                    else:
+                        # Perform non-maximum suppression
+                        max_detections = []
+                        while image_pred_class.size(0):
+                            # Get detection with highest confidence and save as max detection
+                            max_detections.append(image_pred_class[0].unsqueeze(0))
+                            # Stop if we're at the last detection
+                            if len(image_pred_class) == 1:
+                                break
+                            # Get the IOUs for all boxes with lower confidence
+                            ious = bbox_iou(max_detections[-1], image_pred_class[1:], args)
+                            # Remove detections with IoU >= NMS threshold
+                            image_pred_class = image_pred_class[1:][ious < nms_conf]
+
+                        image_pred_class = torch.cat(max_detections).data
+
+                #Concatenate the batch_id of the image to the detection
+                #this helps us identify which image does the detection correspond to 
+                #We use a linear straucture to hold ALL the detections from the batch
+                #the batch_dim is flattened
+                #batch is identified by extra batch column
+
+                batch_ind = image_pred_class.new(image_pred_class.size(0), 1).fill_(ind)
+                seq = batch_ind, image_pred_class
+                if not write:
+                    output = torch.cat(seq,1)
+                    write = True
+                else:
+                    out = torch.cat(seq,1)
+                    output = torch.cat((output,out))
+                num += 1
+    
+        if not num:
+            return 0
+        #output:(n,(batch_ind,x1,y1,x2,y2,c,s,idx of cls))
+        return output
+
+    def detect_one_img(self, img_name):
+        """
+        Detect bboxs in one image
+        Input: 'str', full path of image
+        Output: '[{"category_id":1,"score":float,"bbox":[x,y,w,h],"image_id":str},...]',
+        The output results are similar with coco results type, except that image_id uses full path str
+        instead of coco %012d id for generalization. 
+        """
+        args = self.detector_opt
+        _CUDA = True
+        if args:
+            if args.gpus[0] < 0:
+                _CUDA = False
+        if not self.model:
+            self.load_model()
+        if isinstance(self.model, torch.nn.DataParallel):
+            self.model = self.model.module
+        dets_results = []
+        #pre-process(scale, normalize, ...) the image
+        img, orig_img, img_dim_list = prep_image(img_name, self.inp_dim)
+        with torch.no_grad():
+            img_dim_list = torch.FloatTensor([img_dim_list]).repeat(1, 2)
+            img = img.to(args.device) if args else img.cuda()
+            prediction = self.model(img, args=args)
+            #do nms to the detection results, only human category is left
+            dets = self.dynamic_write_results(prediction, self.confidence,
+                                              self.num_classes, nms=True,
+                                              nms_conf=self.nms_thres)
+            if isinstance(dets, int) or dets.shape[0] == 0:
+                return None
+            dets = dets.cpu()
+
+            img_dim_list = torch.index_select(img_dim_list, 0, dets[:, 0].long())
+            scaling_factor = torch.min(self.inp_dim / img_dim_list, 1)[0].view(-1, 1)
+            dets[:, [1, 3]] -= (self.inp_dim - scaling_factor * img_dim_list[:, 0].view(-1, 1)) / 2
+            dets[:, [2, 4]] -= (self.inp_dim - scaling_factor * img_dim_list[:, 1].view(-1, 1)) / 2
+            dets[:, 1:5] /= scaling_factor
+            for i in range(dets.shape[0]):
+                dets[i, [1, 3]] = torch.clamp(dets[i, [1, 3]], 0.0, img_dim_list[i, 0])
+                dets[i, [2, 4]] = torch.clamp(dets[i, [2, 4]], 0.0, img_dim_list[i, 1])
+
+                #write results
+                det_dict = {}
+                x = float(dets[i, 1])
+                y = float(dets[i, 2])
+                w = float(dets[i, 3] - dets[i, 1])
+                h = float(dets[i, 4] - dets[i, 2])
+                det_dict["category_id"] = 1
+                det_dict["score"] = float(dets[i, 5])
+                det_dict["bbox"] = [x, y, w, h]
+                det_dict["image_id"] = int(os.path.basename(img_name).split('.')[0])
+                dets_results.append(det_dict)
+
+            return dets_results

detector/yolo_cfg.py

@@ -0,0 +1,9 @@
+from easydict import EasyDict as edict
+
+cfg = edict()
+cfg.CONFIG = 'detector/yolo/cfg/yolov3-spp.cfg'
+cfg.WEIGHTS = 'detector/yolo/data/yolov3-spp.weights'
+cfg.INP_DIM =  608
+cfg.NMS_THRES =  0.6
+cfg.CONFIDENCE = 0.1
+cfg.NUM_CLASSES = 80

pretrained_models/fast_421_res152_256x192.txt

@@ -0,0 +1 @@
+这里应该有fast_421_res152_256x192.pth的模型文件，大小为318MB

scripts/demo_inference.py

@@ -0,0 +1,298 @@
+"""Script for single-gpu/multi-gpu demo."""
+import argparse
+import os
+import platform
+import sys
+import time
+
+import numpy as np
+import torch
+from tqdm import tqdm
+import natsort
+
+from detector.apis import get_detector
+from trackers.tracker_api import Tracker
+from trackers.tracker_cfg import cfg as tcfg
+from trackers import track
+from alphapose.models import builder
+from alphapose.utils.config import update_config
+from alphapose.utils.detector import DetectionLoader
+from alphapose.utils.transforms import flip, flip_heatmap
+from alphapose.utils.vis import getTime
+from alphapose.utils.webcam_detector import WebCamDetectionLoader
+from alphapose.utils.writer import DataWriter
+
+"""----------------------------- Demo options -----------------------------"""
+parser = argparse.ArgumentParser(description='AlphaPose Demo')
+parser.add_argument('--cfg', type=str, required=True,
+                    help='experiment configure file name')
+parser.add_argument('--checkpoint', type=str, required=True,
+                    help='checkpoint file name')
+parser.add_argument('--sp', default=False, action='store_true',
+                    help='Use single process for pytorch')
+parser.add_argument('--detector', dest='detector',
+                    help='detector name', default="yolo")
+parser.add_argument('--detfile', dest='detfile',
+                    help='detection result file', default="")
+parser.add_argument('--indir', dest='inputpath',
+                    help='image-directory', default="")
+parser.add_argument('--list', dest='inputlist',
+                    help='image-list', default="")
+parser.add_argument('--image', dest='inputimg',
+                    help='image-name', default="")
+parser.add_argument('--outdir', dest='outputpath',
+                    help='output-directory', default="examples/res/")
+parser.add_argument('--save_img', default=False, action='store_true',
+                    help='save result as image')
+parser.add_argument('--vis', default=False, action='store_true',
+                    help='visualize image')
+parser.add_argument('--showbox', default=False, action='store_true',
+                    help='visualize human bbox')
+parser.add_argument('--profile', default=False, action='store_true',
+                    help='add speed profiling at screen output')
+parser.add_argument('--format', type=str,
+                    help='save in the format of cmu or coco or openpose, option: coco/cmu/open')
+parser.add_argument('--min_box_area', type=int, default=0,
+                    help='min box area to filter out')
+parser.add_argument('--detbatch', type=int, default=5,
+                    help='detection batch size PER GPU')
+parser.add_argument('--posebatch', type=int, default=80,
+                    help='pose estimation maximum batch size PER GPU')
+parser.add_argument('--eval', dest='eval', default=False, action='store_true',
+                    help='save the result json as coco format, using image index(int) instead of image name(str)')
+parser.add_argument('--gpus', type=str, dest='gpus', default="0",
+                    help='choose which cuda device to use by index and input comma to use multi gpus, e.g. 0,1,2,3. (input -1 for cpu only)')
+parser.add_argument('--qsize', type=int, dest='qsize', default=1024,
+                    help='the length of result buffer, where reducing it will lower requirement of cpu memory')
+parser.add_argument('--flip', default=False, action='store_true',
+                    help='enable flip testing')
+parser.add_argument('--debug', default=False, action='store_true',
+                    help='print detail information')
+"""----------------------------- Video options -----------------------------"""
+parser.add_argument('--video', dest='video',
+                    help='video-name', default="")
+parser.add_argument('--webcam', dest='webcam', type=int,
+                    help='webcam number', default=-1)
+parser.add_argument('--save_video', dest='save_video',
+                    help='whether to save rendered video', default=False, action='store_true')
+parser.add_argument('--vis_fast', dest='vis_fast',
+                    help='use fast rendering', action='store_true', default=False)
+"""----------------------------- Tracking options -----------------------------"""
+parser.add_argument('--pose_flow', dest='pose_flow',
+                    help='track humans in video with PoseFlow', action='store_true', default=False)
+parser.add_argument('--pose_track', dest='pose_track',
+                    help='track humans in video with reid', action='store_true', default=False)
+
+args = parser.parse_args()
+cfg = update_config(args.cfg)
+print(args)
+if platform.system() == 'Windows':
+    args.sp = True
+
+args.gpus = [int(i) for i in args.gpus.split(',')] if torch.cuda.device_count() >= 1 else [-1]
+args.device = torch.device("cuda:" + str(args.gpus[0]) if args.gpus[0] >= 0 else "cpu")
+args.detbatch = args.detbatch * len(args.gpus)
+args.posebatch = args.posebatch * len(args.gpus)
+args.tracking = args.pose_track or args.pose_flow or args.detector=='tracker'
+
+if not args.sp:
+    torch.multiprocessing.set_start_method('forkserver', force=True)
+    torch.multiprocessing.set_sharing_strategy('file_system')
+
+
+def check_input():
+    # for wecam
+    if args.webcam != -1:
+        args.detbatch = 1
+        return 'webcam', int(args.webcam)
+
+    # for video
+    if len(args.video):
+        if os.path.isfile(args.video):
+            videofile = args.video
+            return 'video', videofile
+        else:
+            raise IOError('Error: --video must refer to a video file, not directory.')
+
+    # for detection results
+    if len(args.detfile):
+        if os.path.isfile(args.detfile):
+            detfile = args.detfile
+            return 'detfile', detfile
+        else:
+            raise IOError('Error: --detfile must refer to a detection json file, not directory.')
+
+    # for images
+    if len(args.inputpath) or len(args.inputlist) or len(args.inputimg):
+        inputpath = args.inputpath
+        inputlist = args.inputlist
+        inputimg = args.inputimg
+
+        if len(inputlist):
+            im_names = open(inputlist, 'r').readlines()
+        elif len(inputpath) and inputpath != '/':
+            for root, dirs, files in os.walk(inputpath):
+                im_names = files
+            im_names = natsort.natsorted(im_names)
+        elif len(inputimg):
+            args.inputpath = os.path.split(inputimg)[0]
+            im_names = [os.path.split(inputimg)[1]]
+
+        return 'image', im_names
+
+    else:
+        raise NotImplementedError
+
+
+def print_finish_info():
+    print('===========================> Finish Model Running.')
+    if (args.save_img or args.save_video) and not args.vis_fast:
+        print('===========================> Rendering remaining images in the queue...')
+        print('===========================> If this step takes too long, you can enable the --vis_fast flag to use fast rendering (real-time).')
+
+
+def loop():
+    n = 0
+    while True:
+        yield n
+        n += 1
+
+
+if __name__ == "__main__":
+    if args.debug:
+        import pdb;pdb.set_trace()
+    mode, input_source = check_input()
+
+    if not os.path.exists(args.outputpath):
+        os.makedirs(args.outputpath)
+
+    # Load detection loader
+    if mode == 'webcam':
+        det_loader = WebCamDetectionLoader(input_source, get_detector(args), cfg, args)
+        det_worker = det_loader.start()
+    elif mode == 'detfile':
+        det_loader = FileDetectionLoader(input_source, cfg, args)
+        det_worker = det_loader.start()
+    else:
+        det_loader = DetectionLoader(input_source, get_detector(args), cfg, args, batchSize=args.detbatch, mode=mode, queueSize=args.qsize)
+        det_worker = det_loader.start()
+
+    # Load pose model
+    pose_model = builder.build_sppe(cfg.MODEL, preset_cfg=cfg.DATA_PRESET)
+
+    print('Loading pose model from %s...' % (args.checkpoint,))
+    pose_model.load_state_dict(torch.load(args.checkpoint, map_location=args.device))
+    pose_dataset = builder.retrieve_dataset(cfg.DATASET.TRAIN)
+    if args.pose_track:
+        tracker = Tracker(tcfg, args)
+    if len(args.gpus) > 1:
+        pose_model = torch.nn.DataParallel(pose_model, device_ids=args.gpus).to(args.device)
+    else:
+        pose_model.to(args.device)
+    pose_model.eval()
+
+    runtime_profile = {
+        'dt': [],
+        'pt': [],
+        'pn': []
+    }
+
+    # Init data writer
+    queueSize = 2 if mode == 'webcam' else args.qsize
+    if args.save_video and mode != 'image':
+        from alphapose.utils.writer import DEFAULT_VIDEO_SAVE_OPT as video_save_opt
+        if mode == 'video':
+            video_save_opt['savepath'] = os.path.join(args.outputpath, 'AlphaPose_' + os.path.basename(input_source))
+        else:
+            video_save_opt['savepath'] = os.path.join(args.outputpath, 'AlphaPose_webcam' + str(input_source) + '.mp4')
+        video_save_opt.update(det_loader.videoinfo)
+        writer = DataWriter(cfg, args, save_video=True, video_save_opt=video_save_opt, queueSize=queueSize).start()
+    else:
+        writer = DataWriter(cfg, args, save_video=False, queueSize=queueSize).start()
+
+    if mode == 'webcam':
+        print('Starting webcam demo, press Ctrl + C to terminate...')
+        sys.stdout.flush()
+        im_names_desc = tqdm(loop())
+    else:
+        data_len = det_loader.length
+        im_names_desc = tqdm(range(data_len), dynamic_ncols=True)
+
+    batchSize = args.posebatch
+    if args.flip:
+        batchSize = int(batchSize / 2)
+    try:
+        for i in im_names_desc:
+            start_time = getTime()
+            with torch.no_grad():
+                (inps, orig_img, im_name, boxes, scores, ids, cropped_boxes) = det_loader.read()
+                if orig_img is None:
+                    break
+                if boxes is None or boxes.nelement() == 0:
+                    writer.save(None, None, None, None, None, orig_img, im_name)
+                    continue
+                if args.profile:
+                    ckpt_time, det_time = getTime(start_time)
+                    runtime_profile['dt'].append(det_time)
+                # Pose Estimation
+                inps = inps.to(args.device)
+                datalen = inps.size(0)
+                leftover = 0
+                if (datalen) % batchSize:
+                    leftover = 1
+                num_batches = datalen // batchSize + leftover
+                hm = []
+                for j in range(num_batches):
+                    inps_j = inps[j * batchSize:min((j + 1) * batchSize, datalen)]
+                    if args.flip:
+                        inps_j = torch.cat((inps_j, flip(inps_j)))
+                    hm_j = pose_model(inps_j)
+                    if args.flip:
+                        hm_j_flip = flip_heatmap(hm_j[int(len(hm_j) / 2):], pose_dataset.joint_pairs, shift=True)
+                        hm_j = (hm_j[0:int(len(hm_j) / 2)] + hm_j_flip) / 2
+                    hm.append(hm_j)
+                hm = torch.cat(hm)
+                if args.profile:
+                    ckpt_time, pose_time = getTime(ckpt_time)
+                    runtime_profile['pt'].append(pose_time)
+                if args.pose_track:
+                    boxes,scores,ids,hm,cropped_boxes = track(tracker,args,orig_img,inps,boxes,hm,cropped_boxes,im_name,scores)
+                hm = hm.cpu()
+                writer.save(boxes, scores, ids, hm, cropped_boxes, orig_img, im_name)
+                if args.profile:
+                    ckpt_time, post_time = getTime(ckpt_time)
+                    runtime_profile['pn'].append(post_time)
+            
+            if args.profile:
+                # TQDM
+                im_names_desc.set_description(
+                    'det time: {dt:.4f} | pose time: {pt:.4f} | post processing: {pn:.4f}'.format(
+                        dt=np.mean(runtime_profile['dt']), pt=np.mean(runtime_profile['pt']), pn=np.mean(runtime_profile['pn']))
+                )
+        print_finish_info()
+        while(writer.running()):
+            time.sleep(1)
+            print('===========================> Rendering remaining ' + str(writer.count()) + ' images in the queue...')
+        writer.stop()
+        det_loader.stop()
+    except Exception as e:
+        print(repr(e))
+        print('An error as above occurs when processing the images, please check it')
+        pass
+    except KeyboardInterrupt:
+        print_finish_info()
+        # Thread won't be killed when press Ctrl+C
+        if args.sp:
+            det_loader.terminate()
+            while(writer.running()):
+                time.sleep(1)
+                print('===========================> Rendering remaining ' + str(writer.count()) + ' images in the queue...')
+            writer.stop()
+        else:
+            # subprocesses are killed, manually clear queues
+
+            det_loader.terminate()
+            writer.terminate()
+            writer.clear_queues()
+            det_loader.clear_queues()
+

scripts/inference.sh

@@ -0,0 +1,13 @@
+set -x
+
+CONFIG=$1
+CKPT=$2
+VIDEO=$3
+OUTDIR=${4:-"./examples/res"}
+
+python scripts/demo_inference.py \
+    --cfg ${CONFIG} \
+    --checkpoint ${CKPT} \
+    --video ${VIDEO} \
+    --outdir ${OUTDIR} \
+    --detector yolo  --sp  --posebatch 1 #--save_video  #--save_img 

setup.py

@@ -0,0 +1,222 @@
+import os
+import platform
+import subprocess
+import time
+
+import numpy as np
+from Cython.Build import cythonize
+from setuptools import Extension, find_packages, setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+MAJOR = 0
+MINOR = 3
+PATCH = 0
+SUFFIX = ''
+SHORT_VERSION = '{}.{}.{}{}'.format(MAJOR, MINOR, PATCH, SUFFIX)
+
+version_file = 'alphapose/version.py'
+
+
+def readme():
+    with open('README.md') as f:
+        content = f.read()
+    return content
+
+
+def get_git_hash():
+
+    def _minimal_ext_cmd(cmd):
+        # construct minimal environment
+        env = {}
+        for k in ['SYSTEMROOT', 'PATH', 'HOME']:
+            v = os.environ.get(k)
+            if v is not None:
+                env[k] = v
+        # LANGUAGE is used on win32
+        env['LANGUAGE'] = 'C'
+        env['LANG'] = 'C'
+        env['LC_ALL'] = 'C'
+        out = subprocess.Popen(
+            cmd, stdout=subprocess.PIPE, env=env).communicate()[0]
+        return out
+
+    try:
+        out = _minimal_ext_cmd(['git', 'rev-parse', 'HEAD'])
+        sha = out.strip().decode('ascii')
+    except OSError:
+        sha = 'unknown'
+
+    return sha
+
+
+def get_hash():
+    if os.path.exists('.git'):
+        sha = get_git_hash()[:7]
+    elif os.path.exists(version_file):
+        try:
+            from alphapose.version import __version__
+            sha = __version__.split('+')[-1]
+        except ImportError:
+            raise ImportError('Unable to get git version')
+    else:
+        sha = 'unknown'
+
+    return sha
+
+
+def write_version_py():
+    content = """# GENERATED VERSION FILE
+# TIME: {}
+
+__version__ = '{}'
+short_version = '{}'
+"""
+    sha = get_hash()
+    VERSION = SHORT_VERSION + '+' + sha
+
+    with open(version_file, 'w') as f:
+        f.write(content.format(time.asctime(), VERSION, SHORT_VERSION))
+
+
+def get_version():
+    with open(version_file, 'r') as f:
+        exec(compile(f.read(), version_file, 'exec'))
+    return locals()['__version__']
+
+
+def make_cython_ext(name, module, sources):
+    extra_compile_args = None
+    if platform.system() != 'Windows':
+        extra_compile_args = {
+            'cxx': ['-Wno-unused-function', '-Wno-write-strings']
+        }
+
+    extension = Extension(
+        '{}.{}'.format(module, name),
+        [os.path.join(*module.split('.'), p) for p in sources],
+        include_dirs=[np.get_include()],
+        language='c++',
+        extra_compile_args=extra_compile_args)
+    extension, = cythonize(extension)
+    return extension
+
+
+def make_cuda_ext(name, module, sources):
+
+    return CUDAExtension(
+        name='{}.{}'.format(module, name),
+        sources=[os.path.join(*module.split('.'), p) for p in sources],
+        extra_compile_args={
+            'cxx': [],
+            'nvcc': [
+                '-D__CUDA_NO_HALF_OPERATORS__',
+                '-D__CUDA_NO_HALF_CONVERSIONS__',
+                '-D__CUDA_NO_HALF2_OPERATORS__',
+            ]
+        })
+
+
+def get_ext_modules():
+    ext_modules = []
+    # only windows visual studio 2013+ support compile c/cuda extensions
+    # If you force to compile extension on Windows and ensure appropriate visual studio
+    # is intalled, you can try to use these ext_modules.
+    force_compile = False
+    if platform.system() != 'Windows' or force_compile:
+        ext_modules = [
+            make_cython_ext(
+                name='soft_nms_cpu',
+                module='detector.nms',
+                sources=['src/soft_nms_cpu.pyx']),
+            make_cuda_ext(
+                name='nms_cpu',
+                module='detector.nms',
+                sources=['src/nms_cpu.cpp']),
+            make_cuda_ext(
+                name='nms_cuda',
+                module='detector.nms',
+                sources=['src/nms_cuda.cpp', 'src/nms_kernel.cu']),
+            make_cuda_ext(
+                name='roi_align_cuda',
+                module='alphapose.utils.roi_align',
+                sources=['src/roi_align_cuda.cpp', 'src/roi_align_kernel.cu']),
+            make_cuda_ext(
+                name='deform_conv_cuda',
+                module='alphapose.models.layers.dcn',
+                sources=[
+                    'src/deform_conv_cuda.cpp',
+                    'src/deform_conv_cuda_kernel.cu'
+                ]),
+            make_cuda_ext(
+                name='deform_pool_cuda',
+                module='alphapose.models.layers.dcn',
+                sources=[
+                    'src/deform_pool_cuda.cpp',
+                    'src/deform_pool_cuda_kernel.cu'
+                ]),
+        ]
+    return ext_modules
+
+
+def get_install_requires():
+    install_requires = [
+        'six', 'terminaltables', 'scipy==1.1.0',
+        'opencv-python', 'matplotlib', 'visdom',
+        'tqdm', 'tensorboardx', 'easydict',
+        'pyyaml',
+        'torch>=1.1.0', 'torchvision>=0.3.0',
+        'munkres', 'timm==0.1.20', 'natsort'
+    ]
+    # official pycocotools doesn't support Windows, we will install it by third-party git repository later
+    if platform.system() != 'Windows':
+        install_requires.append('pycocotools==2.0.0')
+    return install_requires
+
+
+def is_installed(package_name):
+    from pip._internal.utils.misc import get_installed_distributions
+    for p in get_installed_distributions():
+        if package_name in p.egg_name():
+            return True
+    return False
+
+
+if __name__ == '__main__':
+    write_version_py()
+    setup(
+        name='alphapose',
+        version=get_version(),
+        description='Code for AlphaPose',
+        long_description=readme(),
+        keywords='computer vision, human pose estimation',
+        url='https://github.com/MVIG-SJTU/AlphaPose',
+        packages=find_packages(exclude=('data', 'exp',)),
+        package_data={'': ['*.json', '*.txt']},
+        classifiers=[
+            'Development Status :: 4 - Beta',
+            'License :: OSI Approved :: Apache Software License',
+            'Operating System :: OS Independent',
+            'Programming Language :: Python :: 2',
+            'Programming Language :: Python :: 2.7',
+            'Programming Language :: Python :: 3',
+            'Programming Language :: Python :: 3.4',
+            'Programming Language :: Python :: 3.5',
+            'Programming Language :: Python :: 3.6',
+        ],
+        license='GPLv3',
+        python_requires=">=3",
+        setup_requires=['pytest-runner', 'numpy', 'cython'],
+        tests_require=['pytest'],
+        install_requires=get_install_requires(),
+        ext_modules=get_ext_modules(),
+        cmdclass={'build_ext': BuildExtension},
+        zip_safe=False)
+    # Windows need pycocotools here: https://github.com/philferriere/cocoapi#subdirectory=PythonAPI
+    if platform.system() == 'Windows' and not is_installed('pycocotools'):
+        print("\nInstall third-party pycocotools for Windows...")
+        cmd = 'python -m pip install git+https://github.com/philferriere/cocoapi.git#subdirectory=PythonAPI'
+        os.system(cmd)
+    if not is_installed('cython_bbox'):
+        print("\nInstall `cython_bbox`...")
+        cmd = 'python -m pip install git+https://github.com/yanfengliu/cython_bbox.git'
+        os.system(cmd)