diff --git a/shenzhen_solitaire/cv/card_finder.py b/shenzhen_solitaire/cv/card_finder.py index f7ba254..3753289 100644 --- a/shenzhen_solitaire/cv/card_finder.py +++ b/shenzhen_solitaire/cv/card_finder.py @@ -19,9 +19,11 @@ def _extract_squares(image: np.ndarray, def get_field_squares(image: np.ndarray, - adjustment: Adjustment) -> List[np.ndarray]: + adjustment: Adjustment, + count_x: int, + count_y: int) -> List[np.ndarray]: squares = [] - for index_x, index_y in itertools.product(range(8), range(5)): + 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) @@ -62,14 +64,6 @@ def simplify(image: np.ndarray) -> Tuple[np.ndarray, Dict[Cardcolor, int]]: return (result_image, result_dict) -def get_simplified_squares(image: np.ndarray, - adjustment: Adjustment) -> List[np.ndarray]: - squares = get_field_squares(image, adjustment) - for index, square in enumerate(squares): - squares[index], _ = simplify(square) - return squares - - 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] @@ -107,12 +101,13 @@ def find_square(search_square: np.ndarray, return (best_square, best_count) -def catalogue_cards(squares: List[np.ndarray]) -> List[Tuple[np.ndarray, Card]]: +def catalogue_cards(squares: List[np.ndarray] + ) -> List[Tuple[np.ndarray, Card]]: 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]arkblack") + "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, @@ -131,7 +126,6 @@ def catalogue_cards(squares: List[np.ndarray]) -> List[Tuple[np.ndarray, Card]]: card_type: Optional[Card] = None if len(card_id) == 1: if card_id not in special_card_map: - print("hi") continue card_type = special_card_map[card_id] elif len(card_id) == 2: @@ -146,6 +140,7 @@ def catalogue_cards(squares: List[np.ndarray]) -> List[Tuple[np.ndarray, Card]]: else: continue assert card_type is not None + print(card_type) result.append((square, card_type)) break diff --git a/shenzhen_solitaire/cv/configuration.py b/shenzhen_solitaire/cv/configuration.py index c5f7bbd..358858c 100644 --- a/shenzhen_solitaire/cv/configuration.py +++ b/shenzhen_solitaire/cv/configuration.py @@ -86,6 +86,6 @@ class Configuration: def generate(image: np.ndarray) -> 'Configuration': """Generate a configuration with user input""" adj = adjustment.adjust_field(image) - squares = card_finder.get_field_squares(image, adj) - catalogue = card_finder.catalogue_cards(squares[:2]) + squares = card_finder.get_field_squares(image, adj, 5, 8) + catalogue = card_finder.catalogue_cards(squares) return Configuration(adj=adj, catalogue=catalogue, meta={}) diff --git a/test/cv_helper.py b/test/cv_helper.py index 4eb90f8..828a522 100644 --- a/test/cv_helper.py +++ b/test/cv_helper.py @@ -11,11 +11,12 @@ import cv2 from .context import shenzhen_solitaire from shenzhen_solitaire.cv import adjustment from shenzhen_solitaire.cv import card_finder +from shenzhen_solitaire.cv import board_parser from shenzhen_solitaire import board from shenzhen_solitaire.cv.configuration import Configuration -def main() -> None: +def generate() -> None: with open("Solitaire.png", 'rb') as fd: img_str = fd.read() nparr = np.frombuffer(img_str, np.uint8) @@ -24,6 +25,17 @@ def main() -> None: a = Configuration.generate(image) a.save('test_config.zip') +def parse() -> board.Board: + with open("Solitaire2.png", 'rb') as fd: + img_str = fd.read() + nparr = np.frombuffer(img_str, np.uint8) + image = cv2.imdecode(nparr, cv2.IMREAD_COLOR) + image = cv2.resize(image, (1000, 629)) + + a = Configuration.load('test_config.zip') + a.field_adjustment = adjustment.adjust_field(image) + return board_parser.parse_board(image, a) + if __name__ == "__main__": - main() + generate()