118 lines
2.2 KiB
Python
118 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
from enum import Enum
|
|
from functools import lru_cache
|
|
from pathlib import Path
|
|
from typing import NamedTuple, Dict, Tuple, List
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from sklearn.externals import joblib
|
|
|
|
here: Path = Path(__file__).parent
|
|
|
|
|
|
class COLOR(Enum):
|
|
WHITE = "white"
|
|
BLACK = "black"
|
|
|
|
def __str__(self) -> str:
|
|
return self.value
|
|
|
|
|
|
class PIECE(Enum):
|
|
KNIGHT = 0
|
|
ROOK = 1
|
|
BISHOP = 2
|
|
PAWN = 3
|
|
QUEEN = 4
|
|
KING = 5
|
|
EMPTY = 6
|
|
|
|
def __str__(self) -> str:
|
|
return self.name.lower()
|
|
|
|
class LESS_PIECE(Enum):
|
|
ROOK = 0
|
|
KNIGHT = 1
|
|
BISHOP = 2
|
|
KING = 3
|
|
QUEEN = 4
|
|
|
|
def __str__(self) -> str:
|
|
return self.name.lower()
|
|
|
|
|
|
PieceAndColor = Tuple[PIECE, COLOR]
|
|
|
|
OUR_PIECES = (
|
|
LESS_PIECE.ROOK,
|
|
LESS_PIECE.KNIGHT,
|
|
LESS_PIECE.BISHOP,
|
|
LESS_PIECE.KING,
|
|
LESS_PIECE.QUEEN
|
|
)
|
|
|
|
|
|
class FILE(int, Enum):
|
|
A = 1
|
|
B = 2
|
|
C = 3
|
|
D = 4
|
|
E = 5
|
|
F = 6
|
|
G = 7
|
|
H = 8
|
|
|
|
|
|
class RANK(int, Enum):
|
|
EIGHT = 8
|
|
SEVEN = 7
|
|
SIX = 6
|
|
FIVE = 5
|
|
FOUR = 4
|
|
THREE = 3
|
|
TWO = 2
|
|
ONE = 1
|
|
|
|
|
|
class _Position(NamedTuple):
|
|
file: FILE
|
|
rank: RANK
|
|
|
|
def __str__(self) -> str:
|
|
return f"{self.file.name}{self.rank}"
|
|
|
|
@property
|
|
def color(self):
|
|
if (self.file + self.rank) % 2:
|
|
return COLOR.WHITE
|
|
return COLOR.BLACK
|
|
|
|
|
|
# POSITION.{A8, A7, ..., H1}
|
|
POSITION = Enum("POSITION", {str(_Position(f, r)): _Position(f, r) for f in FILE for r in RANK}, type=_Position) # NOQA
|
|
|
|
# Squares is a dict mapping positions to square images, i.e. a board container during image processing
|
|
Squares = Dict[POSITION, np.ndarray]
|
|
|
|
|
|
class Board(Dict[POSITION, PIECE]):
|
|
"""Board is a dict mapping positions to a piece, i.e. a board configuration after all image processing"""
|
|
|
|
@property
|
|
def to_array(self) -> List[List[int]]:
|
|
return [[self.get(POSITION((file, rank)), PIECE.EMPTY).value for file in FILE]
|
|
for rank in RANK]
|
|
|
|
|
|
def imwrite(*args, **kwargs):
|
|
Path(args[0]).parent.mkdir(parents=True, exist_ok=True)
|
|
return cv2.imwrite(*args, **kwargs)
|
|
|
|
|
|
@lru_cache()
|
|
def load_classifier(filename):
|
|
# print(f"Loading classifier {filename}")
|
|
return joblib.load(str(here.joinpath(filename)))
|