test_mmpose_models.py

# Copyright (c) OpenMMLab. All rights reserved.
import mmengine
import pytest
import torch
from mmengine.config import ConfigDict
from mmengine.structures import InstanceData

from mmdeploy.codebase import import_codebase
from mmdeploy.utils import Backend, Codebase
from mmdeploy.utils.test import WrapModel, check_backend, get_rewrite_outputs

try:
    from torch.testing import assert_close as torch_assert_close
except Exception:
    from torch.testing import assert_allclose as torch_assert_close

try:
    import_codebase(Codebase.MMPOSE)
except ImportError:
    pytest.skip(
        f'{Codebase.MMPOSE} is not installed.', allow_module_level=True)

from .utils import generate_mmpose_deploy_config  # noqa: E402
from .utils import generate_mmpose_task_processor  # noqa: E402


def get_heatmap_head():
    from mmpose.models.heads import HeatmapHead

    model = HeatmapHead(
        2,
        4,
        deconv_out_channels=(16, 16, 16),
        loss=dict(type='KeypointMSELoss', use_target_weight=False))
    model.requires_grad_(False)
    return model


@pytest.mark.parametrize('backend_type', [Backend.ONNXRUNTIME])
def test_heatmaphead_forward(backend_type: Backend):
    check_backend(backend_type, True)
    model = get_heatmap_head()
    model.cpu().eval()
    deploy_cfg = generate_mmpose_deploy_config(backend_type.value)
    feats = [torch.rand(1, 2, 32, 48)]
    wrapped_model = WrapModel(model, 'forward')
    rewrite_inputs = {'feats': feats}
    rewrite_outputs, _ = get_rewrite_outputs(
        wrapped_model=wrapped_model,
        model_inputs=rewrite_inputs,
        deploy_cfg=deploy_cfg,
        run_with_backend=False)
    assert isinstance(rewrite_outputs, torch.Tensor)


def get_msmu_head():
    from mmpose.models.heads import MSPNHead
    model = MSPNHead(
        num_stages=1,
        num_units=1,
        out_shape=(32, 48),
        unit_channels=16,
        level_indices=[1])
    model.requires_grad_(False)
    return model


@pytest.mark.parametrize('backend_type', [Backend.ONNXRUNTIME])
def test_msmuhead_forward(backend_type: Backend):
    check_backend(backend_type, True)
    model = get_msmu_head()
    model.cpu().eval()
    deploy_cfg = generate_mmpose_deploy_config(backend_type.value)
    feats = [[torch.rand(1, 16, 32, 48)]]
    wrapped_model = WrapModel(model, 'forward')
    rewrite_inputs = {'feats': feats}
    rewrite_outputs, _ = get_rewrite_outputs(
        wrapped_model=wrapped_model,
        model_inputs=rewrite_inputs,
        deploy_cfg=deploy_cfg,
        run_with_backend=False)
    assert isinstance(rewrite_outputs, torch.Tensor)


def get_cross_resolution_weighting_model():
    from mmpose.models.backbones.litehrnet import CrossResolutionWeighting

    class DummyModel(torch.nn.Module):

        def __init__(self):
            super().__init__()
            self.model = CrossResolutionWeighting([16, 16], ratio=8)

        def forward(self, x):
            assert isinstance(x, torch.Tensor)
            return self.model([x, x])

    model = DummyModel()
    model.requires_grad_(False)
    return model


@pytest.mark.parametrize('backend_type', [Backend.ONNXRUNTIME])
def test_estimator_forward(backend_type: Backend):
    check_backend(backend_type, True)
    deploy_cfg = generate_mmpose_deploy_config(backend_type.value)
    task_processor = generate_mmpose_task_processor(deploy_cfg=deploy_cfg)
    model = task_processor.build_pytorch_model()
    model.requires_grad_(False)
    model.cpu().eval()
    wrapped_model = WrapModel(model, 'forward', data_samples=None)
    rewrite_inputs = {'inputs': torch.rand(1, 3, 256, 192)}
    rewrite_outputs, _ = get_rewrite_outputs(
        wrapped_model=wrapped_model,
        model_inputs=rewrite_inputs,
        run_with_backend=False,
        deploy_cfg=deploy_cfg)
    assert isinstance(rewrite_outputs, torch.Tensor)


def get_scale_norm_model():
    from mmpose.models.utils.rtmcc_block import ScaleNorm

    model = ScaleNorm(48)
    model.requires_grad_(False)
    return model


@pytest.mark.parametrize('backend_type', [Backend.NCNN])
def test_scale_norm_forward(backend_type: Backend):
    check_backend(backend_type, True)
    deploy_cfg = generate_mmpose_deploy_config(backend_type.value)
    model = get_scale_norm_model()
    x = torch.rand(1, 17, 48)
    wrapped_model = WrapModel(model, 'forward')
    model_outputs = model.forward(x)
    rewrite_inputs = {'x': x}
    rewrite_outputs, _ = get_rewrite_outputs(
        wrapped_model=wrapped_model,
        model_inputs=rewrite_inputs,
        deploy_cfg=deploy_cfg,
        run_with_backend=False)
    torch_assert_close(rewrite_outputs, model_outputs)


def get_rtmcc_block_model():
    from mmpose.models.utils.rtmcc_block import RTMCCBlock

    model = RTMCCBlock(48, 48, 48)
    model.requires_grad_(False)
    return model


@pytest.mark.parametrize('backend_type', [Backend.NCNN])
def test_rtmcc_block_forward(backend_type: Backend):
    check_backend(backend_type, True)
    deploy_cfg = generate_mmpose_deploy_config(backend_type.value)
    model = get_rtmcc_block_model()
    inputs = torch.rand(1, 17, 48)
    wrapped_model = WrapModel(model, '_forward')
    model_outputs = model._forward(inputs)
    rewrite_inputs = {'inputs': inputs}
    rewrite_outputs, _ = get_rewrite_outputs(
        wrapped_model=wrapped_model,
        model_inputs=rewrite_inputs,
        deploy_cfg=deploy_cfg,
        run_with_backend=False)
    torch_assert_close(rewrite_outputs, model_outputs)


def get_scale_model():
    from mmpose.models.utils.rtmcc_block import Scale

    model = Scale(48)
    model.requires_grad_(False)
    return model


@pytest.mark.parametrize('backend_type', [Backend.NCNN])
def test_scale_forward(backend_type: Backend):
    check_backend(backend_type, True)
    deploy_cfg = generate_mmpose_deploy_config(backend_type.value)
    model = get_scale_model()
    x = torch.rand(1, 17, 48)
    wrapped_model = WrapModel(model, 'forward')
    model_outputs = model.forward(x)
    rewrite_inputs = {'x': x}
    rewrite_outputs, _ = get_rewrite_outputs(
        wrapped_model=wrapped_model,
        model_inputs=rewrite_inputs,
        deploy_cfg=deploy_cfg,
        run_with_backend=False)
    torch_assert_close(rewrite_outputs, model_outputs)


@pytest.mark.parametrize('backend_type', [Backend.ONNXRUNTIME])
def test_yolox_pose_head(backend_type: Backend):
    try:
        from mmyolo.utils.setup_env import register_all_modules
        from models.yolox_pose_head import YOLOXPoseHead  # noqa: F401,F403
        register_all_modules(True)
    except ImportError:
        pytest.skip(
            'mmpose/projects/yolox-pose is not installed.',
            allow_module_level=True)
    deploy_cfg = mmengine.Config.fromfile(
        'configs/mmpose/pose-detection_yolox-pose_onnxruntime_dynamic.py')
    check_backend(backend_type, True)

    head = YOLOXPoseHead(
        head_module=dict(
            type='YOLOXPoseHeadModule',
            num_classes=1,
            in_channels=256,
            feat_channels=256,
            widen_factor=0.5,
            stacked_convs=2,
            num_keypoints=17,
            featmap_strides=(8, 16, 32),
            use_depthwise=False,
            norm_cfg=dict(type='BN', momentum=0.03, eps=0.001),
            act_cfg=dict(type='SiLU', inplace=True),
        ),
        loss_cls=dict(
            type='mmdet.CrossEntropyLoss',
            use_sigmoid=True,
            reduction='sum',
            loss_weight=1.0),
        loss_bbox=dict(
            type='mmdet.IoULoss',
            mode='square',
            eps=1e-16,
            reduction='sum',
            loss_weight=5.0),
        loss_obj=dict(
            type='mmdet.CrossEntropyLoss',
            use_sigmoid=True,
            reduction='sum',
            loss_weight=1.0),
        loss_pose=dict(
            type='OksLoss',
            metainfo='configs/_base_/datasets/coco.py',
            loss_weight=30.0),
        loss_bbox_aux=dict(
            type='mmdet.L1Loss', reduction='sum', loss_weight=1.0),
        train_cfg=ConfigDict(
            assigner=dict(
                type='PoseSimOTAAssigner',
                center_radius=2.5,
                iou_calculator=dict(type='mmdet.BboxOverlaps2D'),
                oks_calculator=dict(
                    type='OksLoss',
                    metainfo='configs/_base_/datasets/coco.py'))),
        test_cfg=ConfigDict(
            yolox_style=True,
            multi_label=False,
            score_thr=0.001,
            max_per_img=300,
            nms=dict(type='nms', iou_threshold=0.65)))

    class TestYOLOXPoseHeadModel(torch.nn.Module):

        def __init__(self, yolox_pose_head):
            super(TestYOLOXPoseHeadModel, self).__init__()
            self.yolox_pose_head = yolox_pose_head

        def forward(self, x1, x2, x3):
            inputs = [x1, x2, x3]
            data_sample = InstanceData()
            data_sample.set_metainfo(
                dict(ori_shape=(640, 640), scale_factor=(1.0, 1.0)))
            return self.yolox_pose_head.predict(
                inputs, batch_data_samples=[data_sample])

    model = TestYOLOXPoseHeadModel(head)
    model.cpu().eval()

    model_inputs = [
        torch.randn(1, 128, 8, 8),
        torch.randn(1, 128, 4, 4),
        torch.randn(1, 128, 2, 2)
    ]

    with torch.no_grad():
        pytorch_output = model(*model_inputs)[0]
    pred_bboxes = torch.from_numpy(pytorch_output.bboxes).unsqueeze(0)
    pred_bboxes_scores = torch.from_numpy(pytorch_output.scores).reshape(
        1, -1, 1)
    pred_kpts = torch.from_numpy(pytorch_output.keypoints).unsqueeze(0)
    pred_kpts_scores = torch.from_numpy(
        pytorch_output.keypoint_scores).unsqueeze(0).unsqueeze(-1)

    pytorch_output = [
        torch.cat([pred_bboxes, pred_bboxes_scores], dim=-1),
        torch.cat([pred_kpts, pred_kpts_scores], dim=-1)
    ]

    wrapped_model = WrapModel(model, 'forward')
    rewrite_inputs = {
        'x1': model_inputs[0],
        'x2': model_inputs[1],
        'x3': model_inputs[2]
    }
    deploy_cfg.onnx_config.input_names = ['x1', 'x2', 'x3']

    rewrite_outputs, _ = get_rewrite_outputs(
        wrapped_model=wrapped_model,
        model_inputs=rewrite_inputs,
        run_with_backend=True,
        deploy_cfg=deploy_cfg)

    # keep bbox coord >= 0
    rewrite_outputs[0] = rewrite_outputs[0].clamp(min=0)
    torch_assert_close(rewrite_outputs, pytorch_output)