vendor/ViTPose/tests/test_visualization.py

# Copyright (c) OpenMMLab. All rights reserved.
import tempfile

import mmcv
import numpy as np
import pytest

from mmpose.core import (apply_bugeye_effect, apply_sunglasses_effect,
                         imshow_bboxes, imshow_keypoints, imshow_keypoints_3d)


def test_imshow_keypoints():
    # 2D keypoint
    img = np.zeros((100, 100, 3), dtype=np.uint8)
    kpts = np.array([[1, 1, 1], [10, 10, 1]], dtype=np.float32)
    pose_result = [kpts]
    skeleton = [[0, 1]]
    pose_kpt_color = [(127, 127, 127)] * len(kpts)
    pose_link_color = [(127, 127, 127)] * len(skeleton)
    img_vis_2d = imshow_keypoints(
        img,
        pose_result,
        skeleton=skeleton,
        pose_kpt_color=pose_kpt_color,
        pose_link_color=pose_link_color,
        show_keypoint_weight=True)

    # 3D keypoint
    kpts_3d = np.array([[0, 0, 0, 1], [1, 1, 1, 1]], dtype=np.float32)
    pose_result_3d = [{'keypoints_3d': kpts_3d, 'title': 'test'}]
    _ = imshow_keypoints_3d(
        pose_result_3d,
        img=img_vis_2d,
        skeleton=skeleton,
        pose_kpt_color=pose_kpt_color,
        pose_link_color=pose_link_color,
        vis_height=400)


def test_imshow_bbox():
    img = np.zeros((100, 100, 3), dtype=np.uint8)
    bboxes = np.array([[10, 10, 30, 30], [10, 50, 30, 80]], dtype=np.float32)
    labels = ['label 1', 'label 2']
    colors = ['red', 'green']

    with tempfile.TemporaryDirectory() as tmpdir:
        _ = imshow_bboxes(
            img,
            bboxes,
            labels=labels,
            colors=colors,
            show=False,
            out_file=f'{tmpdir}/out.png')

        # test case of empty bboxes
        _ = imshow_bboxes(
            img,
            np.zeros((0, 4), dtype=np.float32),
            labels=None,
            colors='red',
            show=False)

        # test unmatched bboxes and labels
        with pytest.raises(AssertionError):
            _ = imshow_bboxes(
                img,
                np.zeros((0, 4), dtype=np.float32),
                labels=labels[:1],
                colors='red',
                show=False)


def test_effects():
    img = np.zeros((100, 100, 3), dtype=np.uint8)
    kpts = np.array([[10., 10., 0.8], [20., 10., 0.8]], dtype=np.float32)
    bbox = np.array([0, 0, 50, 50], dtype=np.float32)
    pose_results = [dict(bbox=bbox, keypoints=kpts)]
    # sunglasses
    sunglasses_img = mmcv.imread('demo/resources/sunglasses.jpg')
    _ = apply_sunglasses_effect(
        img,
        pose_results,
        sunglasses_img,
        left_eye_index=1,
        right_eye_index=0,
        kpt_thr=0.5)
    _ = apply_sunglasses_effect(
        img,
        pose_results,
        sunglasses_img,
        left_eye_index=1,
        right_eye_index=0,
        kpt_thr=0.9)

    # bug-eye
    _ = apply_bugeye_effect(
        img, pose_results, left_eye_index=1, right_eye_index=0, kpt_thr=0.5)
    _ = apply_bugeye_effect(
        img, pose_results, left_eye_index=1, right_eye_index=0, kpt_thr=0.9)