def test_get_multi_stage_outputs():
    fake_outputs = [torch.zeros((1, 4, 2, 2))]
    fake_flip_outputs = [torch.ones((1, 4, 2, 2))]
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), outputs_flip=None, num_joints=4, with_heatmaps=[False], with_ae=[True])
    assert (heatmaps == [])
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), outputs_flip=None, num_joints=2, with_heatmaps=[True], with_ae=[True])
    assert (len(heatmaps) == 1)
    flip_index = [1, 0]
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), outputs_flip=fake_flip_outputs, num_joints=2, with_heatmaps=[True], with_ae=[True], flip_index=flip_index)
    assert (len(heatmaps) == 2)
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), tag_per_joint=False, outputs_flip=fake_flip_outputs, num_joints=2, with_heatmaps=[True], with_ae=[True], flip_index=flip_index)
    assert (len(heatmaps) == 2)
    fake_outputs = [torch.zeros((1, 4, 2, 2)), torch.ones((1, 2, 4, 4))]
    fake_flip_outputs = [torch.ones((1, 4, 2, 2)), torch.ones((1, 2, 4, 4))]
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), outputs_flip=None, num_joints=2, with_heatmaps=[True, False], with_ae=[True, True])
    assert torch.allclose(heatmaps[0], torch.tensor(0.0))
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), outputs_flip=fake_flip_outputs, num_joints=2, with_heatmaps=[True, True], with_ae=[True, False])
    assert torch.allclose(heatmaps[0], torch.tensor(0.5))
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), outputs_flip=fake_flip_outputs, num_joints=2, with_heatmaps=[True, False], with_ae=[True, False], flip_index=flip_index)
    assert torch.allclose(heatmaps[0], torch.tensor(0.0))
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), outputs_flip=None, num_joints=2, with_heatmaps=[True, True], with_ae=[True, False], size_projected=(8, 8))
    assert (heatmaps[0].shape == torch.Size([1, 2, 8, 8]))
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), outputs_flip=fake_flip_outputs, num_joints=2, with_heatmaps=[True, True], with_ae=[True, False], align_corners=True)
    assert torch.allclose(heatmaps[0], torch.tensor(0.5))
def test_get_multi_stage_outputs():
    fake_outputs = [torch.zeros((1, 4, 2, 2))]
    fake_flip_outputs = [torch.ones((1, 4, 2, 2))]
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), outputs_flip=None, num_joints=4, with_heatmaps=[False], with_ae=[True])
    assert (heatmaps == [])
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), outputs_flip=None, num_joints=2, with_heatmaps=[True], with_ae=[True])
    assert (len(heatmaps) == 1)
    flip_index = [1, 0]
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), outputs_flip=fake_flip_outputs, num_joints=2, with_heatmaps=[True], with_ae=[True], flip_index=flip_index)
    assert (len(heatmaps) == 2)
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), tag_per_joint=False, outputs_flip=fake_flip_outputs, num_joints=2, with_heatmaps=[True], with_ae=[True], flip_index=flip_index)
    assert (len(heatmaps) == 2)
    fake_outputs = [torch.zeros((1, 4, 2, 2)), torch.ones((1, 2, 4, 4))]
    fake_flip_outputs = [torch.ones((1, 4, 2, 2)), torch.ones((1, 2, 4, 4))]
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), outputs_flip=None, num_joints=2, with_heatmaps=[True, False], with_ae=[True, True])
    assert torch.allclose(heatmaps[0], torch.tensor(0.0))
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), outputs_flip=fake_flip_outputs, num_joints=2, with_heatmaps=[True, True], with_ae=[True, False])
    assert torch.allclose(heatmaps[0], torch.tensor(0.5))
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), outputs_flip=fake_flip_outputs, num_joints=2, with_heatmaps=[True, False], with_ae=[True, False], flip_index=flip_index)
    assert torch.allclose(heatmaps[0], torch.tensor(0.0))
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), outputs_flip=None, num_joints=2, with_heatmaps=[True, True], with_ae=[True, False], size_projected=(8, 8))
    assert (heatmaps[0].shape == torch.Size([1, 2, 8, 8]))
    (outputs, heatmaps, tags) = get_multi_stage_outputs(outputs=copy.deepcopy(fake_outputs), outputs_flip=fake_flip_outputs, num_joints=2, with_heatmaps=[True, True], with_ae=[True, False], align_corners=True)
    assert torch.allclose(heatmaps[0], torch.tensor(0.5))