Testing onnx runtime with easypose.

extras/easypose/README.md

@@ -0,0 +1,9 @@
+# Test ONNX with EasyPose
+
+Code files originally from: https://github.com/Dominic23331/EasyPose.git
+
+```bash
+docker build --progress=plain -f extras/easypose/dockerfile -t rpt_easypose .
+
+./extras/easypose/run_container.sh
+```

extras/easypose/detection.py

@@ -0,0 +1,98 @@
+import numpy as np
+from typing import List
+
+from .base_model import BaseModel
+from .utils import letterbox, nms, xywh2xyxy
+
+
+class RTMDet(BaseModel):
+    def __init__(self,
+                 model_path: str,
+                 conf_threshold: float,
+                 iou_threshold: float,
+                 device: str = 'CUDA',
+                 warmup: int = 30):
+        super(RTMDet, self).__init__(model_path, device, warmup)
+        self.conf_threshold = conf_threshold
+        self.iou_threshold = iou_threshold
+        self.dx = 0
+        self.dy = 0
+        self.scale = 0
+
+    def preprocess(self, image: np.ndarray):
+        th, tw = self.input_shape[2:]
+        image, self.dx, self.dy, self.scale = letterbox(image, (tw, th))
+        tensor = (image - np.array((103.53, 116.28, 123.675))) / np.array((57.375, 57.12, 58.395))
+        tensor = np.expand_dims(tensor, axis=0).transpose((0, 3, 1, 2)).astype(np.float32)
+        return tensor
+
+    def postprocess(self, tensor: List[np.ndarray]):
+        boxes = tensor[0]
+        boxes = np.squeeze(boxes, axis=0)
+        boxes[..., [4, 5]] = boxes[..., [5, 4]]
+
+        boxes = nms(boxes, self.iou_threshold, self.conf_threshold)
+        
+        if boxes.shape[0] == 0:
+            return boxes
+
+        human_class = boxes[..., -1] == 0
+        boxes = boxes[human_class][..., :4]
+
+        boxes[:, 0] -= self.dx
+        boxes[:, 2] -= self.dx
+        boxes[:, 1] -= self.dy
+        boxes[:, 3] -= self.dy
+
+        boxes = np.clip(boxes, a_min=0, a_max=None)
+        boxes[:, :4] /= self.scale
+
+        return boxes
+
+
+class Yolov8(BaseModel):
+    def __init__(self,
+                 model_path: str,
+                 conf_threshold: float,
+                 iou_threshold: float,
+                 device: str = 'CUDA',
+                 warmup: int = 30):
+        super(Yolov8, self).__init__(model_path, device, warmup)
+        self.conf_threshold = conf_threshold
+        self.iou_threshold = iou_threshold
+        self.dx = 0
+        self.dy = 0
+        self.scale = 0
+
+    def preprocess(self, image):
+        th, tw = self.input_shape[2:]
+        image, self.dx, self.dy, self.scale = letterbox(image, (tw, th))
+        tensor = image / 255.
+        tensor = np.expand_dims(tensor, axis=0).transpose((0, 3, 1, 2)).astype(np.float32)
+        return tensor
+
+    def postprocess(self, tensor):
+        feature_map = tensor[0]
+        feature_map = np.squeeze(feature_map, axis=0).transpose((1, 0))
+
+        pred_class = feature_map[..., 4:]
+        pred_conf = np.max(pred_class, axis=-1, keepdims=True)
+        pred_class = np.argmax(pred_class, axis=-1, keepdims=True)
+        boxes = np.concatenate([feature_map[..., :4], pred_conf, pred_class], axis=-1)
+
+        boxes = xywh2xyxy(boxes)
+        boxes = nms(boxes, self.iou_threshold, self.conf_threshold)
+
+        if boxes.shape[0] == 0:
+            return boxes
+
+        human_class = boxes[..., -1] == 0
+        boxes = boxes[human_class][..., :4]
+
+        boxes[:, 0] -= self.dx
+        boxes[:, 2] -= self.dx
+        boxes[:, 1] -= self.dy
+        boxes[:, 3] -= self.dy
+        boxes = np.clip(boxes, a_min=0, a_max=None)
+        boxes[:, :4] /= self.scale
+        return boxes

extras/easypose/dockerfile

@@ -0,0 +1,10 @@
+FROM rapidposetriangulation
+
+WORKDIR /
+
+RUN pip3 install --upgrade --no-cache-dir onnxruntime-gpu
+RUN git clone https://github.com/Dominic23331/EasyPose.git --depth=1
+RUN cd /EasyPose/; pip install -v -e .
+
+WORKDIR /RapidPoseTriangulation/
+CMD ["/bin/bash"]

extras/easypose/pipeline.py

@@ -0,0 +1,262 @@
+import os
+
+import numpy as np
+
+from easypose import model
+from easypose.model import detection
+from easypose.model import pose
+from .download import get_url, get_model_path, download
+from .consts import AvailablePoseModels, AvailableDetModels
+from .common import Person, region_of_interest, restore_keypoints
+
+
+def get_pose_model(pose_model_path, pose_model_decoder, device, warmup):
+    if pose_model_decoder == 'Dark':
+        pose_model = pose.Heatmap(pose_model_path, dark=True, device=device, warmup=warmup)
+    else:
+        pose_model = getattr(pose, pose_model_decoder)(pose_model_path, device=device, warmup=warmup)
+    return pose_model
+
+
+def get_det_model(det_model_path, model_type, conf_thre, iou_thre, device, warmup):
+    det_model = getattr(detection, model_type)(det_model_path, conf_thre, iou_thre, device, warmup)
+    return det_model
+
+
+class TopDown:
+    def __init__(self,
+                 pose_model_name,
+                 pose_model_decoder,
+                 det_model_name,
+                 conf_threshold=0.6,
+                 iou_threshold=0.6,
+                 device='CUDA',
+                 warmup=30):
+        if pose_model_name not in AvailablePoseModels.POSE_MODELS:
+            raise ValueError(
+                'The {} human pose estimation model is not in the model repository.'.format(pose_model_name))
+        if pose_model_decoder not in AvailablePoseModels.POSE_MODELS[pose_model_name]:
+            raise ValueError(
+                'No {} decoding head for the {} model was found in the model repository.'.format(pose_model_decoder,
+                                                                                                 pose_model_name))
+        if det_model_name not in AvailableDetModels.DET_MODELS:
+            raise ValueError(
+                'The {} detection model is not in the model repository.'.format(det_model_name))
+
+        pose_model_dir = get_model_path(AvailablePoseModels.POSE_MODELS[pose_model_name][pose_model_decoder],
+                                         detection_model=False)
+        pose_model_path = os.path.join(pose_model_dir,
+                                      AvailablePoseModels.POSE_MODELS[pose_model_name][pose_model_decoder])
+
+        if os.path.exists(pose_model_path):
+            try:
+                self.pose_model = get_pose_model(pose_model_path, pose_model_decoder, device, warmup)
+            except Exception:
+                url = get_url(AvailablePoseModels.POSE_MODELS[pose_model_name][pose_model_decoder],
+                              detection_model=False)
+                download(url, pose_model_dir)
+                self.pose_model = get_pose_model(pose_model_path, pose_model_decoder, device, warmup)
+        else:
+            url = get_url(AvailablePoseModels.POSE_MODELS[pose_model_name][pose_model_decoder],
+                          detection_model=False)
+            download(url, pose_model_dir)
+            self.pose_model = get_pose_model(pose_model_path, pose_model_decoder, device, warmup)
+
+        det_model_dir = get_model_path(AvailableDetModels.DET_MODELS[det_model_name]['file_name'],
+                                        detection_model=True)
+        det_model_path = os.path.join(det_model_dir,
+                                      AvailableDetModels.DET_MODELS[det_model_name]['file_name'])
+        det_model_type = AvailableDetModels.DET_MODELS[det_model_name]['model_type']
+        if os.path.exists(det_model_path):
+            try:
+                self.det_model = get_det_model(det_model_path,
+                                               det_model_type,
+                                               conf_threshold,
+                                               iou_threshold,
+                                               device,
+                                               warmup)
+            except Exception:
+                url = get_url(AvailableDetModels.DET_MODELS[det_model_name]['file_name'],
+                              detection_model=True)
+                download(url, det_model_dir)
+                self.det_model = get_det_model(det_model_path,
+                                               det_model_type,
+                                               conf_threshold,
+                                               iou_threshold,
+                                               device,
+                                               warmup)
+        else:
+            url = get_url(AvailableDetModels.DET_MODELS[det_model_name]['file_name'],
+                          detection_model=True)
+            download(url, det_model_dir)
+            self.det_model = get_det_model(det_model_path,
+                                           det_model_type,
+                                           conf_threshold,
+                                           iou_threshold,
+                                           device,
+                                           warmup)
+
+    def predict(self, image):
+        boxes = self.det_model(image)
+        results = []
+        for i in range(boxes.shape[0]):
+            p = Person()
+            p.box = boxes[i]
+            region = region_of_interest(image, p.box)
+            kp = self.pose_model(region)
+            p.keypoints = restore_keypoints(p.box, kp)
+            results.append(p)
+        return results
+
+
+class Pose:
+    def __init__(self,
+                 pose_model_name,
+                 pose_model_decoder,
+                 device='CUDA',
+                 warmup=30):
+        if pose_model_name not in AvailablePoseModels.POSE_MODELS:
+            raise ValueError(
+                'The {} human pose estimation model is not in the model repository.'.format(pose_model_name))
+        if pose_model_decoder not in AvailablePoseModels.POSE_MODELS[pose_model_name]:
+            raise ValueError(
+                'No {} decoding head for the {} model was found in the model repository.'.format(pose_model_decoder,
+                                                                                                 pose_model_name))
+
+        pose_model_dir = get_model_path(AvailablePoseModels.POSE_MODELS[pose_model_name][pose_model_decoder],
+                                         detection_model=False)
+        pose_model_path = os.path.join(pose_model_dir,
+                                       AvailablePoseModels.POSE_MODELS[pose_model_name][pose_model_decoder])
+
+        if os.path.exists(pose_model_path):
+            try:
+                self.pose_model = get_pose_model(pose_model_path, pose_model_decoder, device, warmup)
+            except Exception:
+                url = get_url(AvailablePoseModels.POSE_MODELS[pose_model_name][pose_model_decoder],
+                              detection_model=False)
+                download(url, pose_model_dir)
+                self.pose_model = get_pose_model(pose_model_path, pose_model_decoder, device, warmup)
+        else:
+            url = get_url(AvailablePoseModels.POSE_MODELS[pose_model_name][pose_model_decoder],
+                          detection_model=False)
+            download(url, pose_model_dir)
+            self.pose_model = get_pose_model(pose_model_path, pose_model_decoder, device, warmup)
+
+    def predict(self, image):
+        p = Person()
+        box = np.array([0, 0, image.shape[3], image.shape[2], 1, 0])
+        p.box = box
+        p.keypoints = self.pose_model(image)
+        return p
+
+
+class CustomTopDown:
+    def __init__(self,
+                 pose_model,
+                 det_model,
+                 pose_decoder=None,
+                 device='CUDA',
+                 iou_threshold=0.6,
+                 conf_threshold=0.6,
+                 warmup=30):
+        if isinstance(pose_model, model.BaseModel):
+            self.pose_model = pose_model
+        elif isinstance(pose_model, str):
+            if pose_model not in AvailablePoseModels.POSE_MODELS:
+                raise ValueError(
+                    'The {} human pose estimation model is not in the model repository.'.format(pose_model))
+            if pose_model not in AvailablePoseModels.POSE_MODELS[pose_model]:
+                raise ValueError(
+                    'No {} decoding head for the {} model was found in the model repository.'.format(pose_decoder,
+                                                                                                     pose_model))
+
+            pose_model_dir = get_model_path(AvailablePoseModels.POSE_MODELS[pose_model][pose_decoder],
+                                            detection_model=False)
+            pose_model_path = os.path.join(pose_model_dir,
+                                           AvailablePoseModels.POSE_MODELS[pose_model][pose_decoder])
+
+            if os.path.exists(pose_model_path):
+                try:
+                    self.pose_model = get_pose_model(pose_model_path, pose_decoder, device, warmup)
+                except Exception:
+                    url = get_url(AvailablePoseModels.POSE_MODELS[pose_model][pose_decoder],
+                                  detection_model=False)
+                    download(url, pose_model_dir)
+                    self.pose_model = get_pose_model(pose_model_path, pose_decoder, device, warmup)
+            else:
+                url = get_url(AvailablePoseModels.POSE_MODELS[pose_model][pose_decoder],
+                              detection_model=False)
+                download(url, pose_model_dir)
+                self.pose_model = get_pose_model(pose_model_path, pose_decoder, device, warmup)
+        else:
+            raise TypeError("Invalid type for pose model, Please write a custom model based on 'BaseModel'.")
+
+        if isinstance(det_model, model.BaseModel):
+            self.det_model = det_model
+        elif isinstance(det_model, str):
+            if det_model not in AvailableDetModels.DET_MODELS:
+                raise ValueError(
+                    'The {} detection model is not in the model repository.'.format(det_model))
+
+            det_model_dir = get_model_path(AvailableDetModels.DET_MODELS[det_model]['file_name'],
+                                           detection_model=True)
+            det_model_path = os.path.join(det_model_dir,
+                                          AvailableDetModels.DET_MODELS[det_model]['file_name'])
+            det_model_type = AvailableDetModels.DET_MODELS[det_model]['model_type']
+            if os.path.exists(det_model_path):
+                try:
+                    self.det_model = get_det_model(det_model_path,
+                                                   det_model_type,
+                                                   conf_threshold,
+                                                   iou_threshold,
+                                                   device,
+                                                   warmup)
+                except Exception:
+                    url = get_url(AvailableDetModels.DET_MODELS[det_model]['file_name'],
+                                  detection_model=True)
+                    download(url, det_model_dir)
+                    self.det_model = get_det_model(det_model_path,
+                                                   det_model_type,
+                                                   conf_threshold,
+                                                   iou_threshold,
+                                                   device,
+                                                   warmup)
+            else:
+                url = get_url(AvailableDetModels.DET_MODELS[det_model]['file_name'],
+                              detection_model=True)
+                download(url, det_model_dir)
+                self.det_model = get_det_model(det_model_path,
+                                               det_model_type,
+                                               conf_threshold,
+                                               iou_threshold,
+                                               device,
+                                               warmup)
+        else:
+            raise TypeError("Invalid type for detection model, Please write a custom model based on 'BaseModel'.")
+
+    def predict(self, image):
+        boxes = self.det_model(image)
+        results = []
+        for i in range(boxes.shape[0]):
+            p = Person()
+            p.box = boxes[i]
+            region = region_of_interest(image, p.box)
+            kp = self.pose_model(region)
+            p.keypoints = restore_keypoints(p.box, kp)
+            results.append(p)
+        return results
+
+
+class CustomSinglePose:
+    def __init__(self, pose_model):
+        if isinstance(pose_model, model.BaseModel):
+            self.pose_model = pose_model
+        else:
+            raise TypeError("Invalid type for pose model, Please write a custom model based on 'BaseModel'.")
+
+    def predict(self, image):
+        p = Person()
+        box = np.array([0, 0, image.shape[3], image.shape[2], 1, 0])
+        p.box = box
+        p.keypoints = self.pose_model(image)
+        return p

extras/easypose/pose.py

@@ -0,0 +1,64 @@
+import numpy as np
+from typing import List
+
+from .base_model import BaseModel
+from .utils import letterbox, get_heatmap_points, \
+    get_real_keypoints, refine_keypoints_dark, refine_keypoints, simcc_decoder
+
+
+class Heatmap(BaseModel):
+    def __init__(self,
+                 model_path: str,
+                 dark: bool = False,
+                 device: str = 'CUDA',
+                 warmup: int = 30):
+        super(Heatmap, self).__init__(model_path, device, warmup)
+        self.use_dark = dark
+        self.img_size = ()
+
+    def preprocess(self, image: np.ndarray):
+        th, tw = self.input_shape[2:]
+        self.img_size = image.shape[:2]
+        image, _, _, _ = letterbox(image, (tw, th))
+        tensor = (image - np.array((103.53, 116.28, 123.675))) / np.array((57.375, 57.12, 58.395))
+        tensor = np.expand_dims(tensor, axis=0).transpose((0, 3, 1, 2)).astype(np.float32)
+        return tensor
+
+    def postprocess(self, tensor: List[np.ndarray]):
+        heatmaps = tensor[0]
+        heatmaps = np.squeeze(heatmaps, axis=0)
+        keypoints = get_heatmap_points(heatmaps)
+        if self.use_dark:
+            keypoints = refine_keypoints_dark(keypoints, heatmaps, 11)
+        else:
+            keypoints = refine_keypoints(keypoints, heatmaps)
+        keypoints = get_real_keypoints(keypoints, heatmaps, self.img_size)
+        return keypoints
+
+
+class SimCC(BaseModel):
+    def __init__(self, model_path: str, device: str = 'CUDA', warmup: int = 30):
+        super(SimCC, self).__init__(model_path, device, warmup)
+        self.dx = 0
+        self.dy = 0
+        self.scale = 0
+
+    def preprocess(self, image: np.ndarray):
+        th, tw = self.input_shape[2:]
+        image, self.dx, self.dy, self.scale = letterbox(image, (tw, th))
+        tensor = (image - np.array((103.53, 116.28, 123.675))) / np.array((57.375, 57.12, 58.395))
+        tensor = np.expand_dims(tensor, axis=0).transpose((0, 3, 1, 2)).astype(np.float32)
+        return tensor
+
+    def postprocess(self, tensor: List[np.ndarray]):
+        simcc_x, simcc_y = tensor
+        simcc_x = np.squeeze(simcc_x, axis=0)
+        simcc_y = np.squeeze(simcc_y, axis=0)
+        keypoints = simcc_decoder(simcc_x,
+                                  simcc_y,
+                                  self.input_shape[2:],
+                                  self.dx,
+                                  self.dy,
+                                  self.scale)
+
+        return keypoints

extras/easypose/run_container.sh

@@ -0,0 +1,11 @@
+#! /bin/bash
+
+xhost +
+docker run --privileged --rm --network host -it \
+  --gpus all --shm-size=16g --ulimit memlock=-1 --ulimit stack=67108864 \
+  --volume "$(pwd)"/:/RapidPoseTriangulation/ \
+  --volume "$(pwd)"/../datasets/:/datasets/ \
+  --volume "$(pwd)"/../skelda/:/skelda/ \
+  --volume /tmp/.X11-unix:/tmp/.X11-unix \
+  --env DISPLAY --env QT_X11_NO_MITSHM=1 \
+  rpt_easypose

extras/easypose/utils.py

@@ -0,0 +1,203 @@
+from itertools import product
+from typing import Sequence
+
+import cv2
+import numpy as np
+
+
+def letterbox(img: np.ndarray, target_size: Sequence[int], fill_value: int = 128):
+    h, w = img.shape[:2]
+    tw, th = target_size
+
+    scale = min(tw / w, th / h)
+    nw, nh = int(w * scale), int(h * scale)
+
+    resized_img = cv2.resize(img, (nw, nh))
+
+    canvas = np.full((th, tw, img.shape[2]), fill_value, dtype=img.dtype)
+
+    dx, dy = (tw - nw) // 2, (th - nh) // 2
+    canvas[dy:dy + nh, dx:dx + nw, :] = resized_img
+
+    return canvas, dx, dy, scale
+
+
+def intersection_over_union(box1: np.ndarray, box2: np.ndarray):
+    area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
+    area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
+
+    x1 = max(box1[0], box2[0])
+    y1 = max(box1[1], box2[1])
+    x2 = min(box1[2], box2[2])
+    y2 = min(box1[3], box2[3])
+
+    intersection = (x2 - x1) * (y2 - y1)
+    union = area1 + area2 - intersection
+    iou = intersection / (union + 1e-6)
+
+    return iou
+
+
+def xywh2xyxy(boxes):
+    boxes[:, 0] -= boxes[:, 2] / 2
+    boxes[:, 1] -= boxes[:, 3] / 2
+    boxes[:, 2] += boxes[:, 0]
+    boxes[:, 3] += boxes[:, 1]
+    return boxes
+
+def nms(boxes: np.ndarray, iou_threshold: float, conf_threshold: float):
+    conf = boxes[..., 4] > conf_threshold
+    boxes = boxes[conf]
+    boxes = list(boxes)
+    boxes.sort(reverse=True, key=lambda x: x[4])
+
+    result = []
+    while boxes:
+        chosen_box = boxes.pop()
+
+        b = []
+        for box in boxes:
+            if box[-1] != chosen_box[-1] or \
+               intersection_over_union(chosen_box, box) \
+               < iou_threshold:
+                b.append(box)
+
+        result.append(chosen_box)
+        boxes = b
+
+    return np.array(result)
+
+
+def get_heatmap_points(heatmap: np.ndarray):
+    keypoints = np.zeros([1, heatmap.shape[0], 3], dtype=np.float32)
+    for i in range(heatmap.shape[0]):
+        h, w = np.nonzero(heatmap[i] == heatmap[i].max())
+        h, w = h[0], w[0]
+        h_fixed = h + 0.5
+        w_fixed = w + 0.5
+        score = heatmap[i][h][w]
+        keypoints[0][i][0] = w_fixed
+        keypoints[0][i][1] = h_fixed
+        keypoints[0][i][2] = score
+    return keypoints
+
+
+def gaussian_blur(heatmaps: np.ndarray, kernel: int = 11):
+    assert kernel % 2 == 1
+
+    border = (kernel - 1) // 2
+    K, H, W = heatmaps.shape
+
+    for k in range(K):
+        origin_max = np.max(heatmaps[k])
+        dr = np.zeros((H + 2 * border, W + 2 * border), dtype=np.float32)
+        dr[border:-border, border:-border] = heatmaps[k].copy()
+        dr = cv2.GaussianBlur(dr, (kernel, kernel), 0)
+        heatmaps[k] = dr[border:-border, border:-border].copy()
+        heatmaps[k] *= origin_max / np.max(heatmaps[k])
+    return heatmaps
+
+
+def refine_keypoints(keypoints: np.ndarray, heatmaps: np.ndarray):
+    N, K = keypoints.shape[:2]
+    H, W = heatmaps.shape[:2]
+
+    for n, k in product(range(N), range(K)):
+        x, y = keypoints[n, k, :2].astype(int)
+
+        if 1 < x < W - 1 and 0 < y < H:
+            dx = heatmaps[k, y, x + 1] - heatmaps[k, y, x - 1]
+        else:
+            dx = 0.
+
+        if 1 < y < H - 1 and 0 < x < W:
+            dy = heatmaps[k, y + 1, x] - heatmaps[k, y - 1, x]
+        else:
+            dy = 0.
+
+        keypoints[n, k] += np.sign([dx, dy, 0], dtype=np.float32) * 0.25
+
+    return keypoints
+
+
+def refine_keypoints_dark(keypoints: np.ndarray, heatmaps: np.ndarray, blur_kernel_size: int = 11):
+    N, K = keypoints.shape[:2]
+    H, W = heatmaps.shape[1:]
+
+    # modulate heatmaps
+    heatmaps = gaussian_blur(heatmaps, blur_kernel_size)
+    np.maximum(heatmaps, 1e-10, heatmaps)
+    np.log(heatmaps, heatmaps)
+
+    for n, k in product(range(N), range(K)):
+        x, y = keypoints[n, k, :2].astype(int)
+        if 1 < x < W - 2 and 1 < y < H - 2:
+            dx = 0.5 * (heatmaps[k, y, x + 1] - heatmaps[k, y, x - 1])
+            dy = 0.5 * (heatmaps[k, y + 1, x] - heatmaps[k, y - 1, x])
+
+            dxx = 0.25 * (
+                    heatmaps[k, y, x + 2] - 2 * heatmaps[k, y, x] +
+                    heatmaps[k, y, x - 2])
+            dxy = 0.25 * (
+                    heatmaps[k, y + 1, x + 1] - heatmaps[k, y - 1, x + 1] -
+                    heatmaps[k, y + 1, x - 1] + heatmaps[k, y - 1, x - 1])
+            dyy = 0.25 * (
+                    heatmaps[k, y + 2, x] - 2 * heatmaps[k, y, x] +
+                    heatmaps[k, y - 2, x])
+            derivative = np.array([[dx], [dy]])
+            hessian = np.array([[dxx, dxy], [dxy, dyy]])
+            if dxx * dyy - dxy ** 2 != 0:
+                hessianinv = np.linalg.inv(hessian)
+                offset = -hessianinv @ derivative
+                offset = np.squeeze(np.array(offset.T), axis=0)
+                keypoints[n, k, :2] += offset
+    return keypoints
+
+
+def get_real_keypoints(keypoints: np.ndarray, heatmaps: np.ndarray, img_size: Sequence[int]):
+    img_h, img_w = img_size
+    heatmap_h, heatmap_w = heatmaps.shape[1:]
+    heatmap_ratio = heatmaps.shape[1] / heatmaps.shape[2]
+    img_ratio = img_h / img_w
+    if heatmap_ratio > img_ratio:
+        resize_w = img_w
+        resize_h = int(img_w * heatmap_ratio)
+    elif heatmap_ratio < img_ratio:
+        resize_h = img_h
+        resize_w = int(img_h / heatmap_ratio)
+    else:
+        resize_w = img_w
+        resize_h = img_h
+
+    keypoints[:, :, 0] = (keypoints[:, :, 0] / heatmap_w) * resize_w - (resize_w - img_w) / 2
+    keypoints[:, :, 1] = (keypoints[:, :, 1] / heatmap_h) * resize_h - (resize_h - img_h) / 2
+
+    keypoints = np.squeeze(keypoints, axis=0)
+
+    return keypoints
+
+
+def simcc_decoder(simcc_x: np.ndarray,
+                  simcc_y: np.ndarray,
+                  input_size: Sequence[int],
+                  dx: int,
+                  dy: int,
+                  scale: float):
+    x = np.argmax(simcc_x, axis=-1, keepdims=True).astype(np.float32)
+    y = np.argmax(simcc_y, axis=-1, keepdims=True).astype(np.float32)
+
+    x_conf = np.max(simcc_x, axis=-1, keepdims=True)
+    y_conf = np.max(simcc_y, axis=-1, keepdims=True)
+    conf = (x_conf + y_conf) / 2
+
+    x /= simcc_x.shape[-1]
+    y /= simcc_y.shape[-1]
+    x *= input_size[1]
+    y *= input_size[0]
+
+    keypoints = np.concatenate([x, y, conf], axis=-1)
+    keypoints[..., 0] -= dx
+    keypoints[..., 1] -= dy
+    keypoints[..., :2] /= scale
+
+    return keypoints

scripts/utils_2d_pose_ep.py

@@ -0,0 +1,61 @@
+import os
+
+import cv2
+import numpy as np
+import easypose as ep
+
+# ==================================================================================================
+
+filepath = os.path.dirname(os.path.realpath(__file__)) + "/"
+
+# ==================================================================================================
+
+
+def load_model():
+    print("Loading mmpose model ...")
+
+    model = ep.TopDown("rtmpose_m", "SimCC", "rtmdet_s")
+
+    print("Loaded mmpose model")
+    return model
+
+
+def load_wb_model():
+    print("Loading mmpose whole body model ...")
+
+    model = None
+
+    print("Loaded mmpose model")
+    return model
+
+
+# ==================================================================================================
+
+
+def get_2d_pose(model, imgs, num_joints=17):
+    """See: https://mmpose.readthedocs.io/en/latest/user_guides/inference.html#basic-usage"""
+
+    new_poses = []
+    for i in range(len(imgs)):
+        img = imgs[i]
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+        poses = []
+        dets = model.predict(img)
+        for pose in dets:
+            pose = pose.keypoints
+            pose = np.asarray(pose)
+
+            scores = pose[:, 2].reshape(-1, 1)
+            scores = np.clip(scores, 0, 1)
+            pose = np.concatenate((pose[:, :2], scores), axis=-1)
+
+            poses.append(pose)
+
+        if len(poses) == 0:
+            poses.append(np.zeros([num_joints, 3]))
+
+        poses = np.array(poses)
+        new_poses.append(poses)
+
+    return new_poses