ChessBot_ML/FEN_conversion.py at master · DanielProano/ChessBot_ML · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
import numpy as np
import torch

def fen_to_tensor(fen):
    pieces = {'r': 0, 'n': 1, 'b': 2, 'q': 3, 'k': 4, 'p': 5, 'R': 6, 'N': 7, 'B': 8, 'Q': 9, 'K': 10, 'P': 11}
    board_matrix = np.zeros((12, 8, 8), dtype=np.int8)
    fen_array = fen.split(" ")
    board = fen_array[0].split("/")

    for row_index, row in enumerate(board):
        column_index = 0

        for char in row:
            if char.isdigit():
                column_index += int(char)
            else:
                piece_matrix_index = pieces.get(char)
                board_matrix[piece_matrix_index, row_index, column_index] = 1
                column_index += 1

    return [board_matrix, fen_array[1], fen_array[2], fen_array[3]]

def tensor_to_fen(tensor):
    pieces = {0: 'r', 1: 'n', 2: 'b', 3: 'q', 4: 'k', 5: 'p', 6: 'R', 7: 'N', 8: 'B', 9: 'Q', 10: 'K', 11: 'P'}

    board_matrix = tensor[0]

    board_string = ""

    for rows in range(8):
        if rows != 0:
            board_string += '/'
        row = ""
        number_of_empty_indices = 0
        for col in range(8):
            found_a_piece = False
            for piece_index in range(12):
                if board_matrix[piece_index, rows, col] == 1:
                    if number_of_empty_indices != 0:
                        row += str(number_of_empty_indices)
                    row += pieces.get(piece_index)
                    number_of_empty_indices = 0
                    found_a_piece = True
                    break
            if not found_a_piece:
                number_of_empty_indices += 1
        if number_of_empty_indices != 0:
            row += str(number_of_empty_indices)
        board_string += row

    board_string += " " + tensor[1] + " " + tensor[2] + " " + tensor[3]
    return board_string

tensor = fen_to_tensor("rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR b KQkq -")
fen = tensor_to_fen(tensor)
print(fen)