def image_shape_to_hw(image_shape):
    """Extract the height and width of the image from HW, HWC or BHWC"""
    img_size = (-1, -1)
    if len(image_shape) == 2:
        img_size = image_shape
    elif len(image_shape) == 3:
        img_size = image_shape[:2]
    elif len(image_shape) == 4:
        img_size = image_shape[1:3]
    else:
        ValueError('image dimensions must be 2, 3, or 4')
    return img_size