diff --git a/Solitaire2.png b/Solitaire2.png new file mode 100644 index 0000000..06e62da Binary files /dev/null and b/Solitaire2.png differ diff --git a/shenzhen_solitaire/cv/card_finder.py b/shenzhen_solitaire/cv/card_finder.py index cfc884e..35cb382 100644 --- a/shenzhen_solitaire/cv/card_finder.py +++ b/shenzhen_solitaire/cv/card_finder.py @@ -1,7 +1,9 @@ -from typing import List, Tuple +from typing import List, Tuple, Optional, Dict import numpy from .adjustment import Adjustment, get_square from .. import board +import enum +import itertools def _extract_squares(image: numpy.ndarray, @@ -20,3 +22,26 @@ def get_field_squares(image: numpy.ndarray, for iy in range(5): squares.append(get_square(adjustment, ix, iy)) return _extract_squares(image, squares) + + +class Cardcolor(enum.Enum): + Bai = (65, 65, 65) + Black = (0, 0, 0) + Red = (22, 48, 178) + Green = (76, 111, 19) + Background = (178, 194, 193) + + +def simplify(image: numpy.ndarray) -> Dict[Cardcolor, int]: + result_dict: Dict[Cardcolor, int] = {c: 0 for c in Cardcolor} + for pixel in itertools.chain.from_iterable(image): + best_color: Optional[Tuple[Cardcolor, int]] = None + for color in Cardcolor: + mse = sum((x - y) ** 2 for x, y in zip(color.value, pixel)) + if not best_color or best_color[1] > mse: + best_color = (color, mse) + assert best_color + for i in range(3): + pixel[i] = best_color[0].value[i] + result_dict[best_color[0]] += 1 + return result_dict diff --git a/test/cv_helper.py b/test/cv_helper.py index 3210044..085de03 100644 --- a/test/cv_helper.py +++ b/test/cv_helper.py @@ -1,42 +1,45 @@ from .context import shenzhen_solitaire from shenzhen_solitaire.cv import adjustment from shenzhen_solitaire.cv import card_finder +from typing import Tuple, List, Dict import cv2 +import numpy +import itertools A = cv2.imread("Solitaire.png") -adj = adjustment.adjust_field(A) -X = card_finder.get_field_squares(A, adj) -for h in range(20): - p = {None: 0} - for x in X[h]: - for x2 in ((x1[0], x1[1], x1[2]) for x1 in x): - if x2 in p: - p[x2] += 1 - else: - p[x2] = 1 - B = sorted(p.items(), key=lambda x: x[1]) - print(*B, sep='\n') - T = X[h].copy() - cv2.imshow("Window", T) - while cv2.waitKey(0) != 27: - pass - cv2.destroyWindow("Window") -assert 0 - -for ix, vx in enumerate(T): - for iy, vy in enumerate(vx): - if (vy[0] > 100) and (vy[1] > 100) and (vy[2] > 100): - T[ix, iy] = [255, 255, 255] +def pixelcount(image: numpy.ndarray) -> List[Tuple[Tuple[int, int, int], int]]: + p: Dict[Tuple[int, int, int], int] = {(0, 0, 0): 0} + for pixel in itertools.chain.from_iterable(image): + x = tuple(pixel) + if x in p: + p[x] += 1 else: - T[ix, iy] = [0, 0, 0] + p[x] = 1 + B = sorted(p.items(), key=lambda x: x[1]) + return B -cv2.imshow("Window", T) -cv2.waitKey(0) -cv2.destroyWindow("Window") +def simplify(image: numpy.ndarray) -> None: + cv2.imshow("Window", image) + cv2.waitKey(0) + cv2.destroyWindow("Window") + print(*card_finder.simplify(image).items(), sep='\n') + cv2.imshow("Window", image) + cv2.waitKey(0) + cv2.destroyWindow("Window") -# for j in X: -# cv2.imshow("Window", j) -# cv2.waitKey(0) +def main() -> None: + adj = adjustment.adjust_field(A) + image_squares = card_finder.get_field_squares(A, adj) + for img in image_squares: + print(*pixelcount(img), sep='\n') + cv2.imshow("Window", img) + cv2.waitKey(0) + cv2.destroyWindow("Window") + print() + + +if __name__ == "__main__": + main()