98 lines
1.8 KiB
Python
98 lines
1.8 KiB
Python
from __future__ import annotations
|
|
|
|
import cv2
|
|
from enum import Enum
|
|
from functools import lru_cache
|
|
from pathlib import Path
|
|
from typing import NewType, NamedTuple, Dict, Tuple
|
|
|
|
import numpy as np
|
|
from sklearn.externals import joblib
|
|
|
|
|
|
class COLOR(Enum):
|
|
WHITE = "white"
|
|
BLACK = "black"
|
|
|
|
def __str__(self) -> str:
|
|
return self.value
|
|
|
|
|
|
class PIECE(Enum):
|
|
PAWN = "pawn"
|
|
ROOK = "rook"
|
|
KNIGHT = "knight"
|
|
BISHOP = "bishop"
|
|
QUEEN = "queen"
|
|
KING = "king"
|
|
EMPTY = "empty"
|
|
|
|
def __str__(self) -> str:
|
|
return self.value
|
|
|
|
|
|
PieceAndColor = Tuple[PIECE, COLOR]
|
|
|
|
OUR_PIECES = (
|
|
PIECE.ROOK,
|
|
PIECE.KNIGHT,
|
|
)
|
|
|
|
|
|
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)
|
|
|
|
# Squares is a dict mapping positions to square images, i.e. a board container during image processing
|
|
Squares = NewType("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"""
|
|
|
|
|
|
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(filename)
|