How to implement a precise detector for the matrix of color blocks in Python?

Question

Let me illustrate the problem with an example first.

Below is an example input image which might be noised and in other kinds of affine transformation.

The output should return a 2-D array/matrix (Use a standard color character to represent each color block) like

[[w, g, w, b],
 [y, o, o, y],
 [g, b, g, y],
 [b, o, y, b]]

So the aim of this problem is to detect the matrix of color blocks for each input image file, and the orientation doesn't matter.

The idea to solve this problem gives me the intuition that the solution is like detecting and parsing a QR code, but I don't know how to work it out specifically.

Could anyone give me some suggestions like

• The idea/procedure to solve this problem.
• Which APIs of which Python packages should I explore and use.
• Any similar classic problem that I should delve into.
• Some proper code.
• ...

Answer 1

We can start by finding the black dots and "rectifying" the image

The corners can be found by thresholding and looking for "circular" blobs. The metric can be found here: https://learnopencv.com/blob-detection-using-opencv-python-c/

Once we have the corners we "rectify" the image (the image may be rotated and/or flipped depending on the ordering of the corners)

Convert to HSV

Mask again on the "v" channel to get the individual squares

Then we calculate the average hue and saturation of each square to smooth over noise (we need the saturation to distinguish white squares)

We can match these numbers up against our pre-defined color values to print out a string:

['b', 'w', 'g', 'w']
['y', 'o', 'o', 'y']
['y', 'g', 'b', 'g']
['b', 'y', 'o', 'b']

Full Code

import cv2
import numpy as np
import math

# get average color of image within mask
# I think this can be done with some mix of numpy commands (it'd be much faster)
def maskedAverageColor(mask, img):
    # convert to hsv
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV);
    h,s,v = cv2.split(hsv);
    height,width = img.shape[:2];
    count = 0;
    hue_ave = 0;
    sat_ave = 0;

    # calculate average
    for y in range(height):
        for x in range(width):
            if mask[y,x] == 255:
                count += 1;
                hue_ave += h[y,x];
                sat_ave += s[y,x];
    hue_ave /= count;
    sat_ave /= count;
    return [hue_ave, sat_ave];

# center of contour
def getCenter(contour):
    M = cv2.moments(contour);
    cx = int(M['m10']/M['m00']);
    cy = int(M['m01']/M['m00']);
    return [cx, cy];

# predefined colors
# white
# green
# blue
# orange
# yellow
# ??? missing one color (need 6 for a cube?)
color_map = {};
color_map['g'] = 53;
color_map['y'] = 27;
color_map['o'] = 10;
color_map['b'] = 120;

# load image
img = cv2.imread("cube.png");

# show
cv2.imshow("Image", img);

# grayscale
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY);

# threshold
mask = cv2.inRange(gray, 0, 50);

# close small black speckles and dilate
kernel = np.ones((3,3),np.uint8);
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel);
mask = cv2.dilate(mask, kernel, iterations = 1);

# find contours
contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE);

# fill in black spaces
for con in contours:
    cv2.drawContours(mask, [con], -1, 255, -1);
cv2.imshow("Mask", mask);

# get the contours again
contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE);

# filter out little contours
filtered = [];
for con in contours:
    area = cv2.contourArea(con);
    if area > 25:
        filtered.append(con);
contours = filtered;

# circularity
corners = [];
for con in contours:
    area = cv2.contourArea(con);
    perm = cv2.arcLength(con, True);
    if area != 0 and perm != 0:
        circ = (4 * math.pi * area) / (perm**2);
        if circ > 0.5:
            # get the corners
            corners.append(getCenter(con));

# sort corners
# find point in the middle of the corners
cx = 0;
cy = 0;
for corner in corners:
    x,y = corner;
    cx += x;
    cy += y;
cx /= 4;
cy /= 4;

# calculate angles
angles = []; # [angle, point]
for corner in corners:
    x, y = corner;
    dx = 1000;
    dy = y - cy;
    angle = math.atan2(dy, dx);
    angles.append([angle, corner]);

# sort by angles (counter-clockwise)
angles = sorted(angles, key = lambda x : x[0]);
corners = [p[1] for p in angles];

cv2.destroyAllWindows();

# unmorph with corners
to_rect = [
[0, 0],
[500, 0],
[500, 500],
[0, 500]
]
corners = np.array(corners, dtype=np.float32);
to_rect = np.array(to_rect, dtype=np.float32);
warp_mat = cv2.getPerspectiveTransform(corners, to_rect);
rectified = cv2.warpPerspective(img, warp_mat, (500,500));
cv2.imshow("Rect", rectified);

# hsv and mask again
hsv = cv2.cvtColor(rectified, cv2.COLOR_BGR2HSV);
h,s,v = cv2.split(hsv);
cv2.imshow("Hue", h);
cv2.imshow("Sat", s);
cv2.imshow("Val", v);
mask = cv2.inRange(v, 0, 150);

# dilate mask
mask = cv2.dilate(mask, kernel, iterations = 5);
mask = cv2.bitwise_not(mask);

# get contours (yes again)
contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE);

# get the squares
squares = [];
for con in contours:
    area = cv2.contourArea(con);
    if area < 100000:
        squares.append(con);

# get the center of squares
drawn = np.copy(rectified);
for square in squares:
    cx, cy = getCenter(square);
    cv2.circle(drawn, (int(cx), int(cy)), 4, (255,0,0), -1);

# split into rows
rows = [];
index = -1;
for y in range(4):
    row = [];
    for x in range(4):
        row.append(squares[index]);
        index -= 1;
    rows.append(row);

# draw rows
drawn = np.copy(rectified);
for row in rows:
    for square in row:
        cx, cy = getCenter(square);
        cv2.circle(drawn, (int(cx), int(cy)), 4, (255,0,0), -1);

# get color of each square
color_rows = [];
for row in rows:
    color_row = [];
    for square in row:
        # get mask
        h,w = rectified.shape[:2];
        mask = np.zeros((h,w), np.uint8);
        mask = cv2.drawContours(mask, [square], -1, 255, -1);

        # calculate average colors
        hue, sat = maskedAverageColor(mask, rectified);

        # convert to string colors
        # white check
        color_str = "NONE";
        if sat < 50:
            color_str = "w";
        else:
            margin = 5;
            for key in color_map:
                if color_map[key] - margin < hue and hue < color_map[key] + margin:
                    color_str = key;
        if color_str == "NONE":
            print("INCREASE MARGIN");
            print("HUE SAT: " + str([hue, sat]));
        color_row.append(color_str);
    color_rows.append(color_row);

# print out results
for row in color_rows:
    print(row);
cv2.waitKey(0);