main.py

import cv2 as cv
import numpy as np
import time
import argparse

PATTERN_SIZES = [(7, 7), (9, 9)]
CELL_PX = 80
ONLY_DARK_SQUARES = True

def try_find_checkerboard(gray, pattern_sizes):
    flags = cv.CALIB_CB_EXHAUSTIVE + cv.CALIB_CB_ACCURACY
    for ps in pattern_sizes:
        ok, corners = cv.findChessboardCornersSB(gray, ps, flags=flags)
        if ok:
            return ps, corners
    return None, None

def homography_from_corners(corners, squares_w, squares_h, cell_px):
    w_inner, h_inner = squares_w - 1, squares_h - 1
    dst_pts = []
    for j in range(h_inner):
        for i in range(w_inner):
            dst_pts.append([(i + 1) * cell_px, (j + 1) * cell_px])
    dst_pts = np.array(dst_pts, dtype=np.float32).reshape(-1, 1, 2)
    H, _ = cv.findHomography(corners, dst_pts, method=cv.RANSAC, ransacReprojThreshold=3.0)
    return H

def warp_board(frame, H, board_size):
    return cv.warpPerspective(frame, H, board_size)

def detect_pieces_on_board(warped_gray, squares_w, squares_h, cell_px, only_dark=True):
    board_matrix = [[None for _ in range(squares_w)] for _ in range(squares_h)]
    centers_board = []
    colors = []
    for j in range(squares_h):
        for i in range(squares_w):
            if only_dark and ((i + j) % 2 == 0):
                continue
            x0, y0 = i * cell_px, j * cell_px
            roi = warped_gray[y0:y0+cell_px, x0:x0+cell_px]
            if roi.size == 0:
                continue
            blur = cv.GaussianBlur(roi, (5, 5), 1.2)
            circles = cv.HoughCircles(
                blur, cv.HOUGH_GRADIENT,
                dp=1.2, minDist=cell_px*0.8,
                param1=120, param2=15,
                minRadius=int(cell_px*0.25), maxRadius=int(cell_px*0.48)
            )
            if circles is None:
                continue
            cx_t, cy_t = cell_px/2, cell_px/2
            circles = np.uint16(np.around(circles))
            best = min(circles, key=lambda c: (c-cx_t)**2 + (c[1]-cy_t)**2)
            cx, cy, r = int(best), int(best[1]), int(best[2])
            if abs(cx - cx_t) > cell_px*0.35 or abs(cy - cy_t) > cell_px*0.35:
                continue
            mask = np.zeros_like(roi, dtype=np.uint8)
            cv.circle(mask, (cx, cy), int(r*0.8), 255, -1)
            mean_val = cv.mean(roi, mask=mask)
            color = 'white' if mean_val > 128 else 'black'
            xb, yb = x0 + cx, y0 + cy
            board_matrix[j][i] = (float(xb), float(yb))
            centers_board.append([xb, yb])
            colors.append(color)
    if len(centers_board) == 0:
        return board_matrix, None, colors
    centers_board = np.array(centers_board, dtype=np.float32).reshape(-1, 1, 2)
    return board_matrix, centers_board, colors

def project_to_original(centers_board, H):
    Hinv = np.linalg.inv(H)
    pts_img = cv.perspectiveTransform(centers_board, Hinv)
    return pts_img.reshape(-1, 2)

def draw_axes_on_frame(frame, H, cell_px=80):
    Hinv = np.linalg.inv(H)
    origin = np.array([[[0, 0]]], dtype=np.float32)
    x_axis = np.array([[[cell_px, 0]]], dtype=np.float32)
    y_axis = np.array([[[0, cell_px]]], dtype=np.float32)
    o = cv.perspectiveTransform(origin, Hinv)[0,0]
    x = cv.perspectiveTransform(x_axis, Hinv)[0,0]
    y = cv.perspectiveTransform(y_axis, Hinv)[0,0]
    o, x, y = tuple(map(int, o)), tuple(map(int, x)), tuple(map(int, y))
    cv.arrowedLine(frame, o, x, (0, 0, 255), 3, tipLength=0.1)  # X red
    cv.arrowedLine(frame, o, y, (0, 255, 0), 3, tipLength=0.1)  # Y green

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", type=str, default="0", help="0 for camera or path to video")
    ap.add_argument("--cells", type=int, default=CELL_PX, help="pixels per cell in rectified view")
    ap.add_argument("--all_squares", action="store_true", help="search pieces on all squares")
    args = ap.parse_args()

    src = 0
    cap = cv.VideoCapture(src)

    H = None
    squares_w = squares_h = None
    board_size = None
    last_print = 0

    while True:
        ok, frame = cap.read()
        if not ok:
            break
        gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)

        if H is None:
            ps, corners = try_find_checkerboard(gray, PATTERN_SIZES)
            if ps is not None:
                squares_w, squares_h = ps + 1, ps[1] + 1
                board_size = (squares_w * args.cells, squares_h * args.cells)
                H = homography_from_corners(corners, squares_w, squares_h, args.cells)

        if H is not None:
            warped = warp_board(frame, H, board_size)
            warped_gray = cv.cvtColor(warped, cv.COLOR_BGR2GRAY)
            board_matrix, centers_board, colors = detect_pieces_on_board(
                warped_gray, squares_w, squares_h, args.cells,
                only_dark=not args.all_squares
            )
            if centers_board is not None:
                for (xb, yb), col in zip(centers_board.reshape(-1, 2), colors):
                    cv.circle(warped, (int(xb), int(yb)), 8, (0, 255, 0), 2)
            if centers_board is not None:
                pts_img = project_to_original(centers_board, H)
                for (x, y), col in zip(pts_img, colors):
                    color = (255, 255, 255) if col == 'white' else (0, 0, 0)
                    cv.circle(frame, (int(x), int(y)), 10, (0, 255, 255), 2)
                    cv.circle(frame, (int(x), int(y)), 4, color, -1)
            draw_axes_on_frame(frame, H, args.cells)

            if time.time() - last_print > 0.5:
                print("board_matrix (rectified px from top-left):")
                for row in board_matrix:
                    print(row)
                print()
                last_print = time.time()

            cv.imshow("Board (top-down)", warped)

        cv.imshow("Original (annotated)", frame)
        if cv.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv.destroyAllWindows()

if __name__ == "__main__":
    main()