Commit 8953c1f8 authored by narugo1992's avatar narugo1992
Browse files

dev(narugo): fix bug when no detection

parent d5ee3353

imgutils/pose/dwpose.py

+5 −2
Original line number Diff line number Diff line
@@ -410,5 +410,8 @@ def dwpose_estimate(image: ImageTyping, auto_detect: bool = True, 
    resized_img, center, scale = _dwpose_preprocess(np_image, out_bboxes, model_input_size)

    outputs = _dwpose_inference(session, resized_img)

    keypoints, scores = _dwpose_postprocess(outputs, model_input_size, center, scale)

    if keypoints.shape[0] > 0:

        keypoints, scores = _dwpose_reorder_body_points(keypoints, scores)

        return _split_data(keypoints, scores)

    else:

        return []

test/pose/conftest.py

+5 −0
Original line number Diff line number Diff line
@@ -42,3 +42,8 @@ def pose_data_rin(): 
@pytest.fixture()

def pose_data_nad_rin():

    return np.load(get_testfile('pose', 'tohsaka_rin_nad_pose.npy'))





@pytest.fixture()

def img_file_gun():

    return get_testfile('maxim_machine_gun.jpg')

test/pose/test_dwpose.py

+4 −0
Original line number Diff line number Diff line
@@ -31,3 +31,7 @@ class TestPoseDwpose: 
        assert isinstance(poses[1], OP18KeyPointSet)

        assert np.isclose(pose_data_2girls_0, poses[0].all).all()

        assert np.isclose(pose_data_2girls_1, poses[1].all).all()



    def test_dwpose_estimate_nothing(self, img_file_gun):

        poses = dwpose_estimate(img_file_gun)

        assert len(poses) == 0