advancedskrald/util.py

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)))