blob: 49c57a6c688c4abe6117fe9e6081d4b6fd8502d1 (
plain)
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
|
import random
import chess.pgn
INPUT = "big_games.pgn"
OUTPUT = "positions.txt"
def get_random_position(game):
moves = list(game.mainline_moves())
# Game is too short
if len(moves) <= 10:
return None
move_number = random.randint(10, min(len(moves) - 1, 40))
board = game.board()
for i, move in enumerate(moves):
board.push(move)
if i == move_number:
break
return board.fen()
def main():
max_games = int(input("Enter number of game to extract: "))
saved = 0
with open(INPUT, "r") as pgn_file, open(OUTPUT, "w") as fen_file:
for i in range(max_games):
game = chess.pgn.read_game(pgn_file)
if game is None:
break
fen = get_random_position(game)
if fen is None:
continue
fen_file.write(fen + "\n")
saved += 1
print(f"{i + 1}/{max_games} > saved {saved}")
print(f"Done. Saved {saved} positions.")
if __name__ == "__main__":
main()
|
