opencv和numpy实现五子棋

import numpy as np
import cv2
from PIL import Image, ImageFont, ImageDraw


class Gobang:
    """
    五子棋
    """

    def __init__(self):
        # 棋盘大小  18*18
        self.BOARD_SIZE = 18
        # 放大倍数：50。即50个像素表示一个棋盘格子
        self.MAGNIFICATION = 50

        # 0表示空，1表示玩家a的黑子，2表示玩家b的白子
        self.CHESS_NONE_CODE = 0
        self.CHESS_A_CODE = 1
        self.CHESS_B_CODE = 2

        # 窗口名称
        self.WINDOW_NAME = 'Gobang'

        # 黑棋白棋颜色
        self.CHESS_A_COLOR = (0, 0, 0)
        self.CHESS_B_COLOR = (255, 255, 255)

        # 初始化棋子数据
        self.chess_arr = np.full(shape=[self.BOARD_SIZE, self.BOARD_SIZE],
                                 fill_value=self.CHESS_NONE_CODE, dtype=np.uint8)

        # 玩家A和玩家B，默认玩家A先手
        self.is_a_turn = True
        # 棋子计数器
        self.a_counter = 0
        self.b_counter = 0

        # 画棋盘时，向下和右偏移的量
        self.OFFSET = self.MAGNIFICATION // 2

        # 绘制背景
        bg_color = (155, 100, 100)
        r_padding = self.BOARD_SIZE * 20
        self.img = np.full(
            shape=[self.BOARD_SIZE * self.MAGNIFICATION,
                   self.BOARD_SIZE * self.MAGNIFICATION + r_padding, 3],
            fill_value=bg_color,
            dtype=np.uint8)

        self.game_over = False
        # 是否和棋
        self.is_deuce = False
        # 绘制棋盘
        self.paint_board()

        # 显示棋盘 设置鼠标回调事件
        cv2.namedWindow(self.WINDOW_NAME)
        cv2.setMouseCallback(self.WINDOW_NAME, self.on_mouse)
        cv2.imshow(self.WINDOW_NAME, self.img)
        # 按键功能：ESC退出，r键重开
        k = cv2.waitKey(0) & 0xFF
        if k == 27:
            cv2.destroyAllWindows()
        elif k == ord('r'):
            self.__init__()

    # 绘制棋盘
    def paint_board(self):
        height = self.BOARD_SIZE
        width = self.BOARD_SIZE

        # 画棋盘网格线
        line_color = (0, 0, 0)  # 黑色
        for i in range(height):
            for j in range(width):
                # 纵向画线
                cv2.line(self.img,
                         (i * self.MAGNIFICATION + self.OFFSET, self.OFFSET),
                         (i * self.MAGNIFICATION + self.OFFSET,
                          (height - 1) * self.MAGNIFICATION + self.OFFSET),
                         line_color, 2)
                # 横向画线
                cv2.line(self.img, (self.OFFSET, j * self.MAGNIFICATION + self.OFFSET),
                         ((width - 1) * self.MAGNIFICATION + self.OFFSET,
                          j * self.MAGNIFICATION + self.OFFSET),
                         line_color, 2)

    # 鼠标回调函数
    def on_mouse(self, event, x, y, flags, param):
        # 鼠标左键点击
        if event == cv2.EVENT_LBUTTONDOWN:
            if not (self.game_over or self.is_deuce):
                # x，y为放大后的鼠标坐标，index_x, index_y为映射到数组的真正坐标
                index_x = round((x - self.OFFSET) / self.MAGNIFICATION)
                index_y = round((y - self.OFFSET) / self.MAGNIFICATION)
                # 检查输入
                if index_x >= self.BOARD_SIZE or index_y >= self.BOARD_SIZE:
                    print('warning：超出棋盘范围')
                    return
                elif self.chess_arr[index_x, index_y] in [self.CHESS_A_CODE, self.CHESS_B_CODE]:
                    print('warning:此处有子，重新落子')
                    return
                # 玩家落子
                self.chess_step_player(index_x, index_y)
                # 判断胜负
                if self.check_win():
                    self.game_over = True
                    print(f'玩家{"a" if self.is_a_turn else "b"}获胜!')
                # 判断和棋 所有位置下满
                if len(np.argwhere(self.chess_arr == 0)) == 0:
                    self.is_deuce = True
                self.refresh()
                self.is_a_turn = not self.is_a_turn

    # 刷新绘制
    def refresh(self):
        # cv2.imshow(self.WINDOW_NAME, self.img)

        # 绘制黑白双方的步数
        img_temp = self.img
        img_temp = self.draw_text(img_temp, f'黑:{self.a_counter}', (self.BOARD_SIZE * self.MAGNIFICATION,
                                                                    self.BOARD_SIZE * self.MAGNIFICATION // 3))
        img_temp = self.draw_text(img_temp, f'白:{self.b_counter}', (self.BOARD_SIZE * self.MAGNIFICATION,
                                                                    self.BOARD_SIZE * self.MAGNIFICATION // 2))
        if self.game_over:
            # 游戏结束时，绘制结算面板
            img_temp = self.draw_text(img_temp, f'{"黑方" if self.is_a_turn else "白方"}胜!',
                                      (self.BOARD_SIZE * self.MAGNIFICATION,
                                       self.BOARD_SIZE * self.MAGNIFICATION // 3 * 2))
        if self.is_deuce:
            img_temp = self.draw_text(img_temp, '和棋',
                                      (self.BOARD_SIZE * self.MAGNIFICATION,
                                       self.BOARD_SIZE * self.MAGNIFICATION // 3 * 2))
        cv2.imshow(self.WINDOW_NAME, img_temp)

    def draw_text(self, img, text, org):
        # cv2.putText(img, text, org,
        #             cv2.FONT_HERSHEY_TRIPLEX,
        #             self.BOARD_SIZE // 5, (0, 255, 255),
        #             thickness=self.BOARD_SIZE // 5)
        # cv的putText无法自定义字体，因此使用PIL进行字体绘制
        image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        font = ImageFont.truetype(font=r'C:\Windows\Fonts\simsun.ttc', size=self.BOARD_SIZE * 6)
        draw = ImageDraw.Draw(image)
        draw.text(org, text=text, fill=(0, 255, 255), font=font, stroke_width=self.BOARD_SIZE // 10)
        # return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
        return np.array(image)

    # 落子
    def chess_step_player(self, x, y):
        chess_code = self.CHESS_A_CODE if self.is_a_turn else self.CHESS_B_CODE

        self.chess_arr[x][y] = chess_code
        # 棋子计数
        if chess_code == self.CHESS_A_CODE:
            self.a_counter += 1
        else:
            self.b_counter += 1
        print(f'玩家{"a" if self.is_a_turn else "b"}落子：{x}, {y}')
        # 用cv绘制棋子
        self.paint_chess(x, y)

    # 绘制棋子
    def paint_chess(self, x, y):
        # a的回合画黑子，b的回合画白子
        chess_color = self.CHESS_A_COLOR if self.is_a_turn else self.CHESS_B_COLOR
        cv2.circle(
            self.img,
            (x * self.MAGNIFICATION + self.OFFSET, y * self.MAGNIFICATION + self.OFFSET),
            self.MAGNIFICATION // 2,
            chess_color, thickness=-1,
            lineType=8, shift=0)

    # 检查是否有玩家胜利，如果有返回True
    def check_win(self):
        # 判断五子连珠
        if (  # 横向连续五个棋子判断胜利
                self.check_horizontal(self.chess_arr)
                # 纵向检查，这里直接转置数据，调用同样的函数
                or self.check_horizontal(self.chess_arr.T)
                # 对角线方向
                or self.check_oblique(self.chess_arr)
                # 反对角线方向，这里水平反转数据，调用同样的函数
                or self.check_oblique(np.fliplr(self.chess_arr))
        ):
            return True
        else:
            return False

    # 水平方向判断
    def check_horizontal(self, arr):
        chess_code = self.CHESS_A_CODE if self.is_a_turn else self.CHESS_B_CODE

        for i in range(self.BOARD_SIZE):
            for j in range(self.BOARD_SIZE - 4):
                if np.all(arr[i, j:j + 5] == chess_code):
                    return True
        return False

    # 对角方向判断（左上到右下）
    def check_oblique(self, arr):

        # 所有斜线方向判断 这里使用np.diagonal(a, offset=)取出对角线元素
        # offset默认0 当offset=1时，对角线向上移动1个位置。当offset取负值时，向下移动相应的位置数量。
        # 先将对角线转为一维再进行判断  对角线元素小于5个时，无法达成五子连珠，将不再判断
        chess_code = self.CHESS_A_CODE if self.is_a_turn else self.CHESS_B_CODE

        offset_range = self.BOARD_SIZE - 4
        for i in range(-offset_range, offset_range):
            # 转换为一维，偏移量为i
            diag = np.diagonal(arr, offset=i)
            for j in range(len(diag) - 4):
                if np.all(diag[j:j + 5] == chess_code):
                    return True
        return False


if __name__ == '__main__':
    gobang = Gobang()```