cv ¶

def assign_colour(face: np.ndarray,
                  bounding_box: Tuple[int, int, int, int],
                  mass_centre: Tuple[float, float],
                  colour: cube.TileColor,
                  segment_width: float,
                  segment_height: float) -> np.ndarray:
    """
    Assigns the colour value at a mass center in an image to a 3x3
    grid according to the sub-sections in the bounding box.
    Args:
        face: The face to assign the colour to.
        bounding_box: The bounding box (outer grid) surrounding the mass centres.
        mass_centre: The mass centre, which should be alignable in the bounding box in a 3x3 grid.
        colour: The colour of the pixel of the mass centre.
        segment_width: The width of a segment in the bounding box (its width / 3).
        segment_height: The height of a segment in the bounding box (its height / 3).

    Returns:
        A 3x3 numpy array consisting of the TileColors assigned so far.
    """
    for row in range(3):
        for col in range(3):
            if mass_centre[0] > bounding_box[0] + col * segment_width and \
               mass_centre[0] < bounding_box[0] + (col + 1) * segment_width and \
               mass_centre[1] > bounding_box[1] + (row) * segment_height and \
               mass_centre[1] < bounding_box[1] + (row + 1) * segment_height:
                face[row, col] = colour
                return face
    return face

def find_colours(image: np.ndarray) -> np.ndarray:
    """
    Creates a 3x3 numpy array containing TileColors describing the cube which is pictured in the input image in a
    minimal way
    Args:
        image: The input image (BGR) containing a face of a cube.

    Returns:
        A 3x3 numpy array containing the TileColors of the pictured cube.
    """
    image_hsv     = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    image_black   = np.zeros(image.shape)
    image_height, image_width, _ = image.shape

    global COLOURS, COLOUR_RANGES

    coloured_contours = []
    for colour in COLOURS:
        binarized = mask_colour(image_hsv, colour)
        size = 11
        elem = cv2.getStructuringElement(cv2.MORPH_RECT, (size, size))
        binarized = cv2.erode(binarized, elem)

        contours, hierarchy = cv2.findContours(binarized, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        squares_idx = []
        for (idx, contour) in enumerate(contours):
            if is_square_contour(contour, image_width, image_height):
                squares_idx.append(idx)

        for idx in range(len(contours)):
            if idx in squares_idx:
                cv2.drawContours(image_black, contours, idx, get_bgr(colour), -1)
                coloured_contours.append((contours[idx], colour))
            else:
                cv2.drawContours(image_black, contours, idx, (0.5, 0.5, 0.5), 1)

    while len(coloured_contours) > SEGMENTS_PER_FACE:
        # remove outliers until we have the right amount of segments
        sorted_areas   = list(sorted(map(lambda x: cv2.contourArea(x[0]), coloured_contours)))
        mean_area = sorted_areas[len(sorted_areas) // 2] if len(sorted_areas) % 2 != 0 else \
            (sorted_areas[len(sorted_areas) // 2] + sorted_areas[len(sorted_areas) // 2 - 1]) / 2

        max_deviation_idx = -1
        max_deviation     = 0
        for idx, (contour, colour) in enumerate(coloured_contours):
            deviation = abs(cv2.contourArea(contour) - mean_area)
            if deviation > max_deviation:
                max_deviation = deviation
                max_deviation_idx = idx

        coloured_contours.remove(coloured_contours[max_deviation_idx])

    found_contours = [point for sublist in coloured_contours for point in sublist[0]]
    # IMPROVEMENT: make sure the rect is really upright, i.e. by rotation

    if logger.isEnabledFor(logging.DEBUG):
        cv2.imshow("face", image_black)
        cv2.waitKey()

    rect = cv2.boundingRect(np.array(found_contours))
    try:
        result = to_array(image_hsv, coloured_contours, rect)
    except ValueError as e:
        logger.error(e)
        image = cv2.rectangle(image, (rect[0], rect[1]), (rect[0] + rect[2], rect[1] + rect[3]), (255, 255, 255), 1)
        cv2.imwrite("error.png", image)
        cv2.imwrite("cv.png", (image_black * 255).astype(np.uint8))
        raise e

    return result

def get_bgr(colour: str) -> Tuple[float, float, float]:
    """
    Returns the BGR value for a colour supplied as a string.
    Args:
        colour: The colour.

    Returns:
        The BGR value of the colour as a tuple of (B, G, R) in range 0..1

    """
    if colour == 'red':
        return (0.0, 0.0, 1.0)
    elif colour == 'blue':
        return (1.0, 0.0, 0.0)
    elif colour == 'green':
        return (0.0, 1.0 ,0.0)
    elif colour == 'yellow':
        return (0.0, 1.0, 1.0)
    elif colour == 'orange':
        return (0.0, 0.5, 1.0)
    elif colour == 'white':
        return (1.0, 1.0, 1.0)

def get_colour(image_hsv: np.ndarray,
               face: np.ndarray,
               bounding_box: Tuple[int, int, int, int]) -> np.ndarray:
    """
    Creates an array of TileColors described in the passed image.
    More specifically, this function fills in values of the face which are not yet assigned and tries to reconstruct them.
    Args:
        image_hsv: An image in HSV color space containing a face of a Rubik's cube.
        face: The assigned face values so far.
        bounding_box: The bounding box containing all cube segments.

    Returns:
        A 3x3 numpy array containing the TileColors of the pictured cube.
    """
    width, height = bounding_box[2], bounding_box[3]
    segment_width, segment_height = width / 3.0, height / 3.0

    missing_segments = np.argwhere(face == -1)

    for location in missing_segments:

        image_loc_y = int(bounding_box[1] + (location[0] + 0.5) * segment_height)
        image_loc_x = int(bounding_box[0] + (location[1] + 0.5) * segment_width)
        picked_colour = image_hsv[image_loc_y, image_loc_x]

        for colour in COLOURS:
            if mask_colour(np.array([[picked_colour]]), colour)[0][0]:
                face[location[0], location[1]] = cube.TileColor.from_string(colour)

        if face[location[0], location[1]] == -1:
            raise ValueError(
                "Tried to recover colour at " + str(image_loc_x) + " " + str(image_loc_y) + " but got: " + str(
                    picked_colour))

    return face

def is_square_box(box: Tuple[int, int, int, int]) -> bool:
    """
    Returns whether or not a box is square-like.
    Args:
        box: The bounding box in question.
            The rectangle is defined as the opencv rectangle -- a tuple of `(xmin, ymin, width, height)`.

    Returns:
        A bool, indicating whether or not the supplied box is square-like.
    """
    return abs((box[2] / box[3]) - 1.0) < EPSILON

def is_square_contour(contour: List[Tuple[int, int]], width: int, height: int) -> bool:
    """
    Returns whether or not a contour is square-like.
    The checks are based upon:

        - the area of the contour is similar to the area of the bounding box
        - the aspect ratio of the bounding box is close to 1 (a square)
        - the width and height are within an expected range

    Args:
        contour: The contour in question.
            The contour is defined as the opencv contour -- a list of points (tuples) of `(x,y)`.

    Returns:
        A bool, indicating whether or not the supplied contour is square-like.
    """
    contour_area = cv2.contourArea(contour)
    if contour_area == 0.0:
        return False

    rect = cv2.minAreaRect(contour)

    area_measure = (rect[1][0] * rect[1][1]) / contour_area
    aspect_ratio_measure = rect[1][0] / rect[1][1]

    estimated_side_length = np.sqrt(contour_area)

    return (min(width, height) * MIN_SIDE_LENGTH_RATIO) < estimated_side_length < (min(width, height) * MAX_SIDE_LENGTH_RATIO) and \
           (abs(area_measure - 1.0) < EPSILON) and \
           (abs(aspect_ratio_measure - 1.0) < EPSILON)

def mask_colour(image_hsv: np.ndarray, colour: str) -> np.ndarray:
    """
    Masks the supplied image with the supplied colour according to `COLOUR_RANGES`.
    Args:
        image_hsv: The image (HSV) to be masked.
        colour: The colour to mask with.

    Returns:
        The masked image.

    """
    if colour == 'red':
        return cv2.bitwise_or(
            cv2.inRange(image_hsv, COLOUR_RANGES[colour + '_min_0'], COLOUR_RANGES[colour + '_max_0']),
            cv2.inRange(image_hsv, COLOUR_RANGES[colour + '_min_1'], COLOUR_RANGES[colour + '_max_1']))

    return cv2.inRange(image_hsv, COLOUR_RANGES[colour + '_min'], COLOUR_RANGES[colour + '_max'])

def process_images(face_images: List[np.ndarray],
                   use_kmeans: bool = True) -> List[np.ndarray]:
    """
    Processes the images of each face of the cube to a 3x3 numpy array containing the respective TileColors.
    Args:
        face_images: A list of numpy arrays (BGR images) with one side of the cube shown in each of the image.
            Six images are expected in total.
        use_kmeans: Flag, indicating whether or not use the kmeans clustering algorithm and corresponding computer
            to assign the colours. Otherwise, use simpler approach using HSV colour ranges.

    Returns:
        A list of 3x3 numpy arrays containing the TileColors of the respective image.
    """
    assert len(face_images) == len(cube.FacePosition), f"Expected {len(cube.FacePosition)} face images but got: {len(face_images)}"
    for idx, img in enumerate(face_images):
        filename = os.path.join(tempfile.gettempdir(), "img_" + str(idx) + ".png")
        cv2.imwrite(filename, img)

    if not use_kmeans:
        faces = list(map(find_colours, face_images))
    else:
        faces = kmeans_color_detection(face_images)

    logger.debug(faces)

    assert len(faces) == len(cube.FacePosition)
    flattened = np.array(faces).flatten()
    assert (len(flattened) == cube.Cube3D.NR_SEGMENTS_PER_FACE * len(cube.FacePosition))
    for color in cube.TileColor:
        if color is not cube.TileColor.NONE:
            assert (len(np.where(flattened == color)[0]) == cube.Cube3D.NR_SEGMENTS_PER_FACE), \
                f'Have too many / too few tiles for colour {cube.FacePosition.from_tile_color(color).to_string_notation()}: {len(np.where(flattened == color)[0])}'

    return faces

def sort_mass_centres(mass_centre: Tuple[float, float]) -> int:
    """
    Function to sort mass centres by.
    Args:
        mass_centre: The supplied mass centre.

    Returns:
        The value used for sorting.
    """
    x, y = mass_centre[0]
    return x + 3 * y

def to_array(image_hsv: np.ndarray,
             coloured_contours: List[Tuple[List[Tuple[int, int]], str]],
             bounding_box: Tuple[int, int, int, int]) -> np.ndarray:
    """
    Converts a list of coloured contours to a 3x3 numpy array which contains the TileColors depicted in the image.
    Args:
        image_hsv: The input image containing a face of a cube in HSV color space.
        coloured_contours: A list of contours and their respective colour.
        bounding_box: The bounding box containing the contours.

    Returns:
        A 3x3 numpy array containing the TileColors of the pictured cube.
    """
    mass_centres = []
    for (contour, colour) in coloured_contours:
        m = cv2.moments(contour)
        mass_centre = (m['m10'] / m['m00'], m['m01'] / m['m00'])
        mass_centres.append((mass_centre, cube.TileColor.from_string(colour)))

    sorted_mass_centres = sorted(mass_centres, key=sort_mass_centres)

    face = np.full((3, 3), dtype=cube.TileColor, fill_value=-1)
    if len(sorted_mass_centres) <= SEGMENTS_PER_FACE and is_square_box(bounding_box):
        # we can possibly reconstruct the lost faces

        width, height = bounding_box[2], bounding_box[3]
        segment_width, segment_height = width / 3.0, height / 3.0
        for (mass_centre, colour) in sorted_mass_centres:
            face = assign_colour(face, bounding_box, mass_centre, colour, segment_width, segment_height)

        face = get_colour(image_hsv, face, bounding_box)

    elif len(sorted_mass_centres) > SEGMENTS_PER_FACE:
        raise ValueError("Found too many segments")

    elif not is_square_box(bounding_box):
        raise ValueError("Found segments are not bounded by a square box")

    return face