def center_crop(image):
    """
    Center crop images along the maximum dimension
    Args:
        image:
            A numpy array of an image HW3
    Returns:
        cropped_image: A cropped image with H == W
    """
    h, w = image.shape[:2]
    if h > w:
        h_st = (h - w) // 2
        h_en = h_st + w
        cropped_image = image[h_st: h_en, :, :]
    else:
        w_st = (w - h) // 2
        w_en = w_st + h
        cropped_image = image[:, w_st: w_en, :]
    assert cropped_image.shape[0] == cropped_image.shape[1]
    return cropped_image