Model Demo: Pose Estimators Zoo - MMPose

NOTE: The Jupyter Notebook below is included in the Chimera SDK and can be run interactively by running the following CLI command:

$ quadric sdk notebook

From the Jupyter Notebook window in your browser, select the notebook named /quadric/sdk-cli/examples/models/zoo/pose_estimators_zoo/mmpose.ipynb.

MMPose

MMPose is an open source object detection toolbox based on PyTorch. It is a part of the OpenMMLab project.

Please refer the MMPose github for more details.

In this notebook, pose estimation, face estimation, and hand estimation can be experimented.

Environment Setup

!pip3 install -r ../../../requirements.txt -q
!mim install -r ../../../requirements_mim.txt -q

import gc
import os
import pickle
import sys
from pathlib import Path
import shutil
import sys
import zipfile

import matplotlib.pyplot as plt
import numpy as np
import onnx
from onnxruntime import InferenceSession
from PIL import Image

from tvm.contrib.epu.chimera_job.chimera_job import ChimeraJob
from sdk_cli.lib.inference import InferenceEngine, batch_inference
from sdk_cli.lib.quantize import QuantizedONNXModel, quantize_onnx_model
from sdk_cli.node_builtins.outputs.bbox_label_visualizer import draw_bbox
from sdk_cli.utils.models.mmdeploy import MMDeployVariant
from sdk_cli.utils.models.mmpose import MMPoseModelVariant
from sdk_cli.utils.networks import ChimeraNetwork
from sdk_cli.visualizers.layouter import (
    DetectorsLayouter,
    VisualizeHandle,
    VisualizeObject,
    get_visualize_handle,
)
import examples.models.zoo.fix_registry
from examples.models.zoo.zoo_utils import (
    download_file,
    get_mmdeploy_parameters,
    get_transforms_and_dataset,
    install_openmmlab_github,
    onnx_check_and_simplify,
    split_onnx,
)
from pose_estimators_utils import (
    get_mmpose_parameters,
    mmpose_postprocess,
    object_detection_for_estimation,
)

CUR_DIR = os.path.abspath(os.getcwd())

Install MMPose and MMDeploy

github_names = [
    ("mmpose", "1.3.2", None),
    ("mmdeploy", "1.3.1", None),
]
symlink_list = [
    ("configs", "mmpose"),
    ("model-index.yml", "mmpose"),
]
install_openmmlab_github(github_names, symlink_list)

os.chdir(Path(CUR_DIR) / "mmpose")

import mmpose
import mmpose.utils

os.chdir(Path(CUR_DIR) / "mmdeploy")

from mmdeploy.apis import torch2onnx
import torch

os.chdir(CUR_DIR)

Select Model

This notebook can experiment the following models.

• Pose Estimation
  • COCO Dataset
    • td_hm_hourglass52_256x256
    • td_hm_hourglass52_384x384
    • td_hm_mobilenetv2
    • td_hm_mobilenetv2_384x288
    • td_hm_res50_384x288
    • td_hm_res101_dark_384x288
    • td_hm_mspn50
    • td_hm_rsn18
    • td_hm_rsn50
    • td_hm_vgg16
  • MPII Dataset
    • td_hm_hourglass52_mpii
    • td_hm_hourglass52_mpii_384x384
  • CrowdPose Dataset
    • td_hm_res101_crowdpose_320x256
• Face Landmark Estimation
  • td_hm_res50_coco_face
  • td_hm_hourglass52_coco_face
  • td_hm_mobilenetv2_coco_face
• Hand Landmark Estimation
  • td_hm_res50_coco_hand
  • td_hm_hourglass52_coco_hand
  • td_hm_mobilenetv2_coco_hand
• Animal Pose Estimation
  • td_hm_res50_animalpose
  • td_hm_res101_animalpose
  • td_hm_res152_animalpose
• Whole Body Estimation
  • td_hm_res101_coco_wholebody_384x288

Now we are going to experiment td_hm_hourglass52_384x384 as an example.

## Select target model.
MODEL_NAME = MMPoseModelVariant.td_hm_hourglass52_384x384

mmpose_parameters = get_mmpose_parameters(MODEL_NAME)
model_input_size = mmpose_parameters.model_input_size
mmdeploy_variant = MMDeployVariant.mmpose

mmdeploy_parameters = get_mmdeploy_parameters(mmdeploy_variant, mmpose_parameters)

Load Pretrained Model

weights = Path(mmdeploy_parameters.weights_url).name

if not Path(weights).exists():
    download_file(mmdeploy_parameters.weights_url, weights)
    print(f"{weights} is downloaded.")

Export ONNX

onnx_file = f"{MODEL_NAME}.onnx"

## Convert the model to ONNX
torch2onnx(
    "../../../common/calibration/face/daniel_maverick.png",
    work_dir=".",
    save_file=onnx_file,
    deploy_cfg=mmdeploy_parameters.deploy_config,
    model_cfg=mmdeploy_parameters.config,
    model_checkpoint=weights,
    device="cpu",
)

onnx_model = onnx_check_and_simplify(onnx.load(onnx_file), mmpose_parameters.skip_shape_inference)
onnx.save(onnx_model, onnx_file)

print(f"ONNX is exported and simplified as {onnx_file}")
del torch, torch2onnx, mmpose.utils, mmpose, onnx_model
gc.collect();

Split ONNX

session = InferenceSession(onnx.load(onnx_file).SerializeToString())
input_edges = [inp.name for inp in session.get_inputs()]
output_edges = [oup.name for oup in session.get_outputs()]

target_onnx_file, second_onnx_file = split_onnx(
    onnx_file,
    mmpose_parameters,
    input_edges=input_edges,
    output_edges=output_edges,
    backbone_head=True,
)
del session
gc.collect();

Quantization

transforms, coco_like_subset_of_dataset = get_transforms_and_dataset(mmdeploy_parameters)

quantized_onnx_model: QuantizedONNXModel = quantize_onnx_model(
    target_onnx_file,
    coco_like_subset_of_dataset,
    asymmetric_activation=mmpose_parameters.asymmetric_activation,
)
print(
    f"quantized onnx: {quantized_onnx_model.model_path}, tranges file: {quantized_onnx_model.tensor_ranges_path}"
)
del coco_like_subset_of_dataset
gc.collect();

Compilation

cgc_job = ChimeraJob(
    model_p=str(quantized_onnx_model.model_path),
    trange_file=str(quantized_onnx_model.tensor_ranges_path),
)
cgc_job.compile(quiet=True)
print(cgc_job)

Demo

Object Detection

if mmpose_parameters.visualize_object.startswith("PoseEstimation"):
    all_image_paths = [
        "../../../common/calibration/face/daniel_maverick.png",
    ]
elif mmpose_parameters.visualize_object == VisualizeObject.FaceEstimation:
    all_image_paths = [
        "../../../common/calibration/face/veer_daniel.jpg",
    ]
elif mmpose_parameters.visualize_object == VisualizeObject.HandEstimation:
    all_image_paths = [
        "../../mediapipe/images/fingers/daniel_hand.jpg",
    ]
else:
    raise ValueError(f"{mmpose_parameters.visual_object} is not implemented.")

bboxes_per_image = object_detection_for_estimation(
    mmpose_parameters.visualize_object, all_image_paths
)

%matplotlib inline
pic_len = len(all_image_paths)

ax, idx = {}, 0
fig = plt.figure(figsize=(5 * pic_len, 5), tight_layout=True)
for image_file, detections in bboxes_per_image.items():
    idx += 1
    ax[idx] = fig.add_subplot(int("1%s%s" % (pic_len, idx)))
    ax[idx].imshow(
        draw_bbox(
            np.array(Image.open(image_file)),
            detections,
            show_class=False,
            show_score=False,
        )
    )
    ax[idx].set_title(f"Bounding Boxes: {Path(image_file).name}")
    ax[idx].axis("off")

fig.show()

Pose Estimation

Preprocess for Pose Estimation

There could be multiple humans detected in an image. This preprocessing is to make a list of the sub images corresponding to bounding boxes per image.

persons_per_image = {}

for image_path, detections in bboxes_per_image.items():
    frame = np.array(Image.open(image_path))
    persons_per_image[image_path] = []
    for bbox in detections:
        x1, y1, x2, y2 = bbox[:4].astype(np.int32)
        extracted_image = frame[y1:y2, x1:x2, :]
        persons_per_image[image_path].append(Image.fromarray(extracted_image))

Inference

Pose Estimation inference is done on image portions which the detection model predicts as humans. The batch inference is done on those image portions per image.

engines = {
    InferenceEngine.CHIMERA_ORT_INT8: cgc_job,
    InferenceEngine.CHIMERA_ISS_INT8: cgc_job,
}

head_session = (
    None
    if second_onnx_file is None
    else InferenceSession(onnx.load(second_onnx_file).SerializeToString())
)

outputs_per_image = {}
for image_path, all_images in persons_per_image.items():
    all_transformed_images = []
    for image in all_images:
        transformed_image = transforms(image)

        all_transformed_images.append(
            np.expand_dims(transformed_image.numpy(), axis=0).astype(np.float32)
        )
    outputs_per_image[image_path] = {}
    THREADS = min(len(all_transformed_images), 6)
    for inference_engine, engine in engines.items():
        outputs_per_image[image_path][inference_engine] = batch_inference(
            inference_engine,
            engine,
            all_transformed_images,
            head_session=head_session,
            threads=THREADS,
        )

Run Statistics

print(cgc_job)
cgc_job.plot_run_statistics()

Display Inference Results

%matplotlib inline

visualize_object = mmpose_parameters.visualize_object
visualize_handle = get_visualize_handle(
    visualize_object, mmpose_parameters.visual_threshold, random_seed=17
)

for (image, outputs_per_inference_engine), detections in zip(
    outputs_per_image.items(), bboxes_per_image.values()
):
    layouter = DetectorsLayouter(
        visualize_handle,
        image=image,
    )
    for inference_engine, all_outputs in outputs_per_inference_engine.items():
        for outputs, bboxes in zip(all_outputs, detections):
            _, _, keypoints, keypoints_scores = mmpose_postprocess(
                outputs[0], bboxes, mmpose_parameters
            )
            layouter.add_data(
                keypoints=keypoints,
                keypoints_scores=keypoints_scores,
                title=f"{MODEL_NAME}: {str(inference_engine).upper()}",
            )

    layouter.display()

Citation

@misc{mmpose2020,
    title={OpenMMLab Pose Estimation Toolbox and Benchmark},
    author={MMPose Contributors},
    howpublished = {\url{https://github.com/open-mmlab/mmpose}},
    year={2020}
}

@misc{=mmdeploy,
    title={OpenMMLab's Model Deployment Toolbox.},
    author={MMDeploy Contributors},
    howpublished = {\url{https://github.com/open-mmlab/mmdeploy}},
    year={2021}
}