Worked on detecting card border

shenzhen_solitaire/card_detection/board_parser.py

@@ -2,50 +2,89 @@
 
 import numpy as np
 from .configuration import Configuration
-from ..board import Board
+from ..board import Board, NumberCard, SpecialCard
 from . import card_finder
 import cv2
-from typing import Iterable, Any, List
+from typing import Iterable, Any, List, Tuple, Union
 import itertools
 
 
-def parse_board(image: np.ndarray, conf: Configuration) -> Board:
-    """Parse a screenshot of the game, using a given configuration"""
+def grouper(
+    iterable: Iterable[Any], groupsize: int, fillvalue: Any = None
+) -> Iterable[Iterable[Any]]:
+    "Collect data into fixed-length chunks or blocks"
+    args = [iter(iterable)] * groupsize
+    return itertools.zip_longest(*args, fillvalue=fillvalue)
+
+
+def get_square_iterator(
+    image: np.ndarray, conf: Configuration, row_count: int, column_count: int
+) -> Iterable[Tuple[np.ndarray, np.ndarray]]:
+    """Return iterator for both the square, as well as the matching card border"""
     fake_adjustments = conf.field_adjustment
     fake_adjustments.x -= 5
     fake_adjustments.y -= 5
     fake_adjustments.h += 10
     fake_adjustments.w += 10
-    row_count = 13
-    column_count = 8
-
-    def grouper(iterable: Iterable[Any], groupsize: int, fillvalue: Any = None) -> Iterable[Any]:
-        "Collect data into fixed-length chunks or blocks"
-        args = [iter(iterable)] * groupsize
-        return itertools.zip_longest(*args, fillvalue=fillvalue)
-
     squares = card_finder.get_field_squares(
-        image, conf.field_adjustment, count_x=row_count, count_y=column_count
+        image, fake_adjustments, count_x=row_count, count_y=column_count
+    )
+    border_squares = card_finder.get_field_squares(
+        image, conf.border_adjustment, count_x=row_count, count_y=column_count
     )
     grouped_squares = grouper(squares, row_count)
-    result = Board()
-    for group_index, square_group in enumerate(grouped_squares):
-        group_field = []
-        for index, square in enumerate(square_group):
-            best_val = None
-            best_name = None
-            for template, name in conf.catalogue:
-                res = cv2.matchTemplate(square, template, cv2.TM_CCOEFF_NORMED)
-                min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
-                if best_val is None or max_val > best_val:
-                    best_val = max_val
-                    best_name = name
-            assert best_name is not None
-            group_field.append(best_name)
+    grouped_border_squares = grouper(border_squares, row_count)
+    return zip(grouped_squares, grouped_border_squares)
 
-            # print(f"\t{best_val}: {best_name}")
-            # cv2.imshow("Catalogue", cv2.resize(square, (500, 500)))
-            # cv2.waitKey()
+
+def match_template(template: np.ndarray, search_image: np.ndarray) -> int:
+    """Return matchiness for the template on the search image"""
+    res = cv2.matchTemplate(search_image, template, cv2.TM_CCOEFF_NORMED)
+    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
+    assert isinstance(max_val, int)
+    return max_val
+
+
+def parse_square(
+    square: np.ndarray, border: np.ndarray, conf: Configuration
+) -> Tuple[Union[NumberCard, SpecialCard], bool]:
+    square_fits = [
+        (match_template(template, square), name) for template, name in conf.catalogue
+    ]
+    best_val, best_name = max(square_fits, key=lambda x: x[0])
+
+    best_border = max(
+        match_template(template=template, search_image=border)
+        for template in conf.card_border
+    )
+    best_empty = max(
+        match_template(template=template, search_image=border)
+        for template in conf.empty_card
+    )
+
+    assert best_name is not None
+    assert best_empty is not None
+    assert best_border is not None
+    row_finished = best_empty > best_border
+
+    return (best_name, row_finished)
+
+
+def parse_board(image: np.ndarray, conf: Configuration) -> Board:
+    """Parse a screenshot of the game, using a given configuration"""
+    square_iterator = get_square_iterator(
+        image, conf, row_count=Board.MAX_ROW_SIZE, column_count=8
+    )
+    result = Board()
+    for group_index, (square_group, border_group) in enumerate(square_iterator):
+        group_field = []
+        for index, (square, border_square) in enumerate(
+            zip(square_group, border_group)
+        ):
+            value, row_finished = parse_square(square, border_square, conf)
+            group_field.append(value)
+            if row_finished:
+                break
 
         result.field[group_index] = group_field