board_detector.py

import cv2
import numpy as np
import math
# import os
from PIL import Image
import psutil
import time

# global show_cv because I didn't want to have show_cv as an input to every function
show_cv = None
def init_show_cv(val):
    global show_cv
    show_cv = val

# had to write a custom img display function because conflict with picamera2 and opencv made it so I can't use imshow
# (I had to install headless opencv to remove conflict which removes imshow)
def display_img(img_array):
    if len(img_array) < 5: # prevent spamming windows if accidentally input an image instead of array of images
        for cv2_img in img_array:
            rgb_img = cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB)
            pil_img = Image.fromarray(rgb_img) 
            pil_img.show()
    input()
    for proc in psutil.process_iter():
        if proc.name() == "display":
            proc.kill()

def find_longest_lines(img):
    gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray_img, 50, 100, apertureSize=3)

    if (show_cv):
        # cv2.imshow('Canny Filter', edges)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()
        display_img([edges])

    theta_thresh = 60
    horizontal_lines = cv2.HoughLines(edges, rho=1, theta=np.pi/180, threshold=60, min_theta=(theta_thresh/2-1)*np.pi/theta_thresh, max_theta=(theta_thresh/2+1)*np.pi/theta_thresh)
    vertical_lines = cv2.HoughLines(edges, rho=1, theta=np.pi/180, threshold=60, min_theta=-np.pi/theta_thresh, max_theta=np.pi/theta_thresh)

    vertical_line_points = convert_to_cartesian(vertical_lines)
    horizontal_line_points = convert_to_cartesian(horizontal_lines)

    # filter lines too close to each other
    filtered_vertical = filter_lines(vertical_line_points, 50)
    filtered_horizontal = filter_lines(horizontal_line_points, 50)

    # get the 9 largest lines
    sorted_vertical = sorted(filtered_vertical, key=lambda line: min(line[0][1], line[0][3]))[:9]
    sorted_horizontal = sorted(filtered_horizontal, key=lambda line: min(line[0][0], line[0][2]))[:9]

    return sorted_vertical, sorted_horizontal

def convert_to_cartesian(lines):
    line_points = []
    if lines is not None:
        for line in lines:
            rho, theta = line[0]

            cos_theta  = np.cos(theta)
            sin_theta = np.sin(theta)

            x0 = cos_theta * rho
            y0 = sin_theta * rho

            x1 = int(x0 + 1000 * (-sin_theta))
            y1 = int(y0 + 1000 * (cos_theta))
            x2 = int(x0 - 1000 * (-sin_theta))
            y2 = int(y0 - 1000 * (cos_theta))

            line_points.append([[x1,y1,x2,y2]])
    return line_points

def filter_lines(lines, min_distance):
    filtered_lines = []

    # filter out lines too close to each other
    # (this assumes lines are around the same size and parallel)
    # (extremely simplified to improve computational speed because this is all we need)
    if lines is not None:
        for line1 in lines:
            x1, y1, x2, y2 = line1[0]
            line1_x_avg = (x1 + x2) / 2
            line1_y_avg = (y1 + y2) / 2
            keep_line = True
            for line2 in filtered_lines:
                x3, y3, x4, y4 = line2[0]
                line2_x_avg = (x3 + x4) / 2
                line2_y_avg = (y3 + y4) / 2

                # calculate dist between average points of the 2 lines
                dist = np.sqrt((line1_x_avg - line2_x_avg)**2 + (line1_y_avg - line2_y_avg)**2)

                if dist < min_distance:
                    keep_line = False
                    break

            if keep_line:
                filtered_lines.append(line1)
            
    return filtered_lines

def find_board(img):
    vertical_lines, horizontal_lines = find_longest_lines(img)

    # create bitmasks for vert and horiz so we can get lines and intersections
    height, width, _ = img.shape
    vertical_mask = np.zeros((height, width), dtype=np.uint8)
    horizontal_mask = np.zeros((height, width), dtype=np.uint8)

    for line in vertical_lines:
        x1, y1, x2, y2 = line[0]
        cv2.line(vertical_mask, (x1, y1), (x2, y2), (255), 2)

    for line in horizontal_lines:
        x1, y1, x2, y2 = line[0]
        cv2.line(horizontal_mask, (x1, y1), (x2, y2), (255), 2)

    # get lines and intersections of grid and corresponding contours
    intersection = cv2.bitwise_and(vertical_mask, horizontal_mask)
    board_lines = cv2.bitwise_or(vertical_mask, horizontal_mask)

    contours, hierarchy = cv2.findContours(board_lines, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    if (show_cv):
        intersections, hierarchy = cv2.findContours(intersection, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        # find midpoints of intersection contours (this gives us exact coordinates of the corners of the grid)
        intersection_points = []
        for contour in intersections:
            M = cv2.moments(contour)
            if (M["m00"] != 0):
                midpoint_x = int(M["m10"] / M["m00"])
                midpoint_y = int(M["m01"] / M["m00"])
                intersection_points.append((midpoint_x, midpoint_y))

        # sort the coordinates from left to right then top to bottom
        sorted_intersection_points = sort_square_grid_coords(intersection_points, unpacked=False)

        board_lines_img = img.copy()
        cv2.drawContours(board_lines_img, contours, -1, (255, 255, 0), 2)
        i = 0
        for points in sorted_intersection_points:
            for point in points:
                cv2.circle(board_lines_img, point, 5, (255 - (3 * i), (i % 9) * 28, 3 * i), -1)
                i += 1
        # cv2.imshow('Lines of Board', board_lines_img)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()
        display_img([board_lines_img])

    if (len(vertical_lines) != 9 or len(horizontal_lines) != 9):
        print("Error: Grid does not match expected 9x9")
        print("# of Vertical:",len(vertical_lines))
        print("# of Horizontal:",len(horizontal_lines))
        return None, None

    # find largest contour and get rid of it because it contains weird edges from lines
    max_area = 100000 # we're assuming board is going to be big (hopefully to speed up computation on raspberry pi)
    largest = -1
    # second_largest = -1
    # max_rect = None
    for i, contour in enumerate(contours):
        area = cv2.contourArea(contour)
        if area > max_area:
            max_area = area
            largest = i
    # "largest" is index of largest contour

    # get rid of contour containing the edges of the lines   
    contours = list(contours)
    contours.pop(largest)
    contours = tuple(contours)

    # thicken lines so that connections are made
    contour_mask = np.zeros((height, width), dtype=np.uint8)
    cv2.drawContours(contour_mask, contours, -1, (255), thickness=10)
    thick_contours, _ = cv2.findContours(contour_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # obtain largest contour of the thickened lines (the border) and approximate a 4 sided polygon onto it
    max_area = 100000
    largest = -1
    max_rect = None
    for i, contour in enumerate(thick_contours):
        area = cv2.contourArea(contour)
        if area > max_area:
            epsilon = 0.05 * cv2.arcLength(contour, True)
            rect = cv2.approxPolyDP(contour, epsilon, True) # uses Douglas-Peucker algorithm (probably overkill)
            if (len(rect) == 4):
                max_area = area
                largest = i
                max_rect = rect

    # perspective transform based on rectangle outline of board
    corners = max_rect.reshape(-1, 2) # reshapes it so each row has 2 elements
    corners = [tuple(corner) for corner in corners] # convert to tuples
    corners_sorted = sort_square_grid_coords(corners, unpacked=True)

    tl = corners_sorted[0]
    tr = corners_sorted[1]
    bl = corners_sorted[2]
    br = corners_sorted[3]
    src = np.float32([list(tl), list(tr), list(bl), list(br)])
    dest = np.float32([[0,0], [width, 0], [0, height], [width, height]])
    M = cv2.getPerspectiveTransform(src, dest)
    Minv = cv2.getPerspectiveTransform(dest, src)
    warped_img = img.copy()
    warped_img = cv2.warpPerspective(np.uint8(warped_img), M, (width, height))

    # perspective transform the intersections as well
    M = cv2.getPerspectiveTransform(src, dest)
    Minv = cv2.getPerspectiveTransform(dest, src)
    warped_ip = intersection.copy() # warped intersection points
    warped_ip = cv2.warpPerspective(np.uint8(warped_ip), M, (width, height))

    warped_intersections, hierarchy = cv2.findContours(warped_ip, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    # find midpoints of warped intersection contours (this gives us exact coordinates of the corners of the grid)
    warped_intersection_points = []
    for contour in warped_intersections:
        M = cv2.moments(contour)
        if (M["m00"] != 0):
            midpoint_x = int(M["m10"] / M["m00"])
            midpoint_y = int(M["m01"] / M["m00"])
            warped_intersection_points.append((midpoint_x, midpoint_y))

    # sort the coordinates from left to right then top to bottom
    sorted_warped_points = sort_square_grid_coords(warped_intersection_points, unpacked=False)

    if (show_cv):
        contours_img = img.copy()
        cv2.drawContours(contours_img, thick_contours, -1, (0, 255, 0), 2)
        cv2.drawContours(contours_img, [thick_contours[largest]], -1, (0, 0, 255), 2)
        cv2.drawContours(contours_img, [max_rect], -1, (255, 0, 0), 2)
        for x,y in corners:
            cv2.circle(contours_img, (x, y), 5, (0, 255, 255), -1)
        # cv2.imshow('Contours', contours_img)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()
        display_img([contours_img])

    return warped_img, sorted_warped_points

def find_pieces(warped_img, sorted_warped_points):
    hsv_img = cv2.cvtColor(warped_img, cv2.COLOR_BGR2HSV)

    # threshold to find strongest colors in image
    hsv_mask_sat = cv2.inRange(hsv_img[:,:,1], 60, 255) # saturation mask
    hsv_mask_bright = cv2.inRange(hsv_img[:,:,2], 100, 255) # brightness mask
    hsv_mask = cv2.bitwise_and(hsv_mask_sat, hsv_mask_bright)

    hsv_after = cv2.bitwise_and(hsv_img, hsv_img, mask=hsv_mask)
    bgr_after = cv2.cvtColor(hsv_after, cv2.COLOR_HSV2BGR)
    

    # define color thresholds to use to classify colors later on
    hue_thresh_dict = {'red': (170,190), 'orange':(8,20), 'yellow': (20,40), 'green': (50,70), 'blue': (105,120), 'purple': (120,140), 
                       'teal': (90,105), 'pink': (140,170)} 

    count = 0
    color_grid = []
    # loop through each square of chess board
    for i in range(0,8):
        color_grid.append([])
        for j in range(0,8):
            # establish corners of current square
            tl = sorted_warped_points[i][j]
            tr = sorted_warped_points[i][j+1]
            bl = sorted_warped_points[i+1][j]
            br = sorted_warped_points[i+1][j+1]

            # create a polygon mask for current grid square 
            # (this is because square is not always perfectly square so I can't just loop pixels from tl to tr then bl to br)
            # (this might not be worth the extra computation required to shave off a few pixels from being considered)
            height, width, _ = warped_img.shape
            rect_mask = np.zeros((height, width), dtype=np.uint8)
            poly = np.array([[tl, tr, br, bl]], dtype=np.int32)
            cv2.fillPoly(rect_mask, poly, 255)
            num_pixels = 1
            hue = 0

            # loop through a perfect square that is slightly bigger than actual "square." obtain average hue of color in grid square
            for x in range(min(tl[0],bl[0]), max(tr[0],br[0])):
                for y in range(min(tl[1],tr[1]), max(bl[1],br[1])):
                    if rect_mask[y,x] == 255 and hsv_after[y, x, 0] != 0:
                        num_pixels += 1
                        hue += hsv_after[y, x, 0]
            avg_hue = hue / num_pixels

            # if there is a color in square, then label based on custom hue thresholds and add to color_grid
            piece_found = False
            if avg_hue != 0:
                for color, (lower, upper) in hue_thresh_dict.items():
                    if lower <= avg_hue < upper:
                        color_grid[i].append((color, avg_hue, num_pixels))
                        piece_found = True
                        break # should only do this once
            
            if piece_found == False:
                color_grid[i].append((None, avg_hue, num_pixels))

            count += 1

    # print color_grid. only print when the color is found a lot in the square (> 100 times)
    for row in color_grid:
        print("||", end="")
        for color, avg_hue, num_pixels in row:
            if num_pixels > 100:
                print(color,int(avg_hue), "\t|", end="")
            else:
                print("\t\t|", end="")
        print("|")

    if show_cv:
        # cv2.imshow('Warped', warped_img)

        bgr_after_intersections = bgr_after.copy()
        for points in sorted_warped_points:
                for point in points:
                    cv2.circle(bgr_after_intersections, point, 1, (255, 255, 255), -1)
        display_img([warped_img, bgr_after_intersections])
        # cv2.imshow('hsv_after_intersections', bgr_after_intersections)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()

    return color_grid

def sort_square_grid_coords(coordinates, unpacked):
    # this function assumes there are a perfect square amount of coordinates
    sqrt_len = int(math.sqrt(len(coordinates)))
    sorted_coords = sorted(coordinates, key=lambda coord: coord[1]) # first sort by y values
    # then group rows of the square (for example, 9x9 grid would be 81 coordinates so split into 9 arrays of 9)
    rows = [sorted_coords[i:i+sqrt_len] for i in range(0, len(sorted_coords), sqrt_len)]
    for row in rows:
        row.sort(key=lambda coord: coord[0]) # now sort each row by x

    if (unpacked == False):
        return rows
    
    collapsed = [coord for row in rows for coord in row] # unpack/collapse groups to just be an array of tuples
    return collapsed