132 lines
4.4 KiB
Python
132 lines
4.4 KiB
Python
"""Functions to detect card value"""
|
|
|
|
from typing import List, Tuple, Optional, Dict
|
|
import enum
|
|
import itertools
|
|
import numpy as np
|
|
import cv2
|
|
from .adjustment import Adjustment, get_square
|
|
from ..board import Card, NumberCard, SpecialCard
|
|
|
|
|
|
def _extract_squares(
|
|
image: np.ndarray, squares: List[Tuple[int, int, int, int]]
|
|
) -> List[np.ndarray]:
|
|
return [
|
|
image[square[1] : square[3], square[0] : square[2]].copy() for square in squares
|
|
]
|
|
|
|
|
|
def get_field_squares(
|
|
image: np.ndarray, adjustment: Adjustment, count_x: int, count_y: int
|
|
) -> List[np.ndarray]:
|
|
"""Return all squares in the field, according to the adjustment"""
|
|
squares = []
|
|
for index_x, index_y in itertools.product(range(count_y), range(count_x)):
|
|
squares.append(get_square(adjustment, index_x, index_y))
|
|
return _extract_squares(image, squares)
|
|
|
|
|
|
class Cardcolor(enum.Enum):
|
|
"""Relevant colors for different types of cards"""
|
|
|
|
Bai = (65, 65, 65)
|
|
Black = (0, 0, 0)
|
|
Red = (22, 48, 178)
|
|
Green = (76, 111, 19)
|
|
Background = (178, 194, 193)
|
|
|
|
def _find_single_square(
|
|
search_square: np.ndarray, template_square: np.ndarray
|
|
) -> Tuple[int, Tuple[int, int]]:
|
|
assert search_square.shape[0] >= template_square.shape[0]
|
|
assert search_square.shape[1] >= template_square.shape[1]
|
|
best_result: Optional[Tuple[int, Tuple[int, int]]] = None
|
|
for margin_x, margin_y in itertools.product(
|
|
range(search_square.shape[0], template_square.shape[0] - 1, -1),
|
|
range(search_square.shape[1], template_square.shape[1] - 1, -1),
|
|
):
|
|
search_region = search_square[
|
|
margin_x - template_square.shape[0] : margin_x,
|
|
margin_y - template_square.shape[1] : margin_y,
|
|
]
|
|
count = cv2.countNonZero(search_region - template_square)
|
|
if not best_result or count < best_result[0]: # pylint: disable=E1136
|
|
best_result = (
|
|
count,
|
|
(
|
|
margin_x - template_square.shape[0],
|
|
margin_y - template_square.shape[1],
|
|
),
|
|
)
|
|
assert best_result
|
|
return best_result
|
|
|
|
|
|
def find_square(
|
|
search_square: np.ndarray, squares: List[np.ndarray]
|
|
) -> Tuple[np.ndarray, int]:
|
|
"""Compare all squares in squares with search_square, return best matching one.
|
|
Requires all squares to be simplified."""
|
|
best_set = False
|
|
best_square: Optional[np.ndarray] = None
|
|
best_count = 0
|
|
for square in squares:
|
|
count, _ = _find_single_square(search_square, square)
|
|
if not best_set or count < best_count:
|
|
best_set = True
|
|
best_square = square
|
|
best_count = count
|
|
assert isinstance(best_square, np.ndarray)
|
|
return (best_square, best_count)
|
|
|
|
|
|
def catalogue_cards(squares: List[np.ndarray]) -> List[Tuple[np.ndarray, Card]]:
|
|
"""Run manual cataloging for given squares"""
|
|
cv2.namedWindow("Catalogue", cv2.WINDOW_NORMAL)
|
|
cv2.waitKey(1)
|
|
result: List[Tuple[np.ndarray, Card]] = []
|
|
print("Card ID is [B]ai, [Z]hong, [F]a, [H]ua, [R]ed, [G]reen, [B]lack")
|
|
print("Numbercard e.g. R3")
|
|
special_card_map = {
|
|
"b": SpecialCard.Bai,
|
|
"z": SpecialCard.Zhong,
|
|
"f": SpecialCard.Fa,
|
|
"h": SpecialCard.Hua,
|
|
}
|
|
suit_map = {
|
|
"r": NumberCard.Suit.Red,
|
|
"g": NumberCard.Suit.Green,
|
|
"b": NumberCard.Suit.Black,
|
|
}
|
|
for square in squares:
|
|
while True:
|
|
cv2.imshow("Catalogue", cv2.resize(square, (500, 500)))
|
|
cv2.waitKey(1)
|
|
card_id = input("Card ID:").lower()
|
|
card_type: Optional[Card] = None
|
|
if len(card_id) == 1:
|
|
if card_id not in special_card_map:
|
|
continue
|
|
card_type = special_card_map[card_id]
|
|
elif len(card_id) == 2:
|
|
if not card_id[0] in suit_map:
|
|
continue
|
|
if not card_id[1].isdigit():
|
|
continue
|
|
if card_id[1] == "0":
|
|
continue
|
|
card_type = NumberCard(
|
|
number=int(card_id[1]), suit=suit_map[card_id[0]]
|
|
)
|
|
else:
|
|
continue
|
|
assert card_type is not None
|
|
print(card_type)
|
|
result.append((square, card_type))
|
|
break
|
|
|
|
cv2.destroyWindow("Catalogue")
|
|
assert result is not None
|
|
return result
|