Restructured code
This commit is contained in:
0
shenzhen_solitaire/card_detection/__init__.py
Normal file
0
shenzhen_solitaire/card_detection/__init__.py
Normal file
97
shenzhen_solitaire/card_detection/adjustment.py
Normal file
97
shenzhen_solitaire/card_detection/adjustment.py
Normal file
@@ -0,0 +1,97 @@
|
||||
"""Contains functions to find significant pieces of a solitaire screenshot"""
|
||||
|
||||
from typing import Optional, Tuple
|
||||
from dataclasses import dataclass
|
||||
import itertools
|
||||
import numpy
|
||||
import cv2
|
||||
|
||||
|
||||
@dataclass
|
||||
class Adjustment:
|
||||
"""Configuration for a grid"""
|
||||
x: int
|
||||
y: int
|
||||
w: int
|
||||
h: int
|
||||
dx: int
|
||||
dy: int
|
||||
|
||||
|
||||
def get_square(adjustment: Adjustment, index_x: int = 0,
|
||||
index_y: int = 0) -> Tuple[int, int, int, int]:
|
||||
"""Get one square from index and adjustment"""
|
||||
return (adjustment.x + adjustment.dx * index_x,
|
||||
adjustment.y + adjustment.dy * index_y,
|
||||
adjustment.x + adjustment.w + adjustment.dx * index_x,
|
||||
adjustment.y + adjustment.h + adjustment.dy * index_y)
|
||||
|
||||
|
||||
def _adjust_squares(
|
||||
image: numpy.ndarray,
|
||||
count_x: int,
|
||||
count_y: int,
|
||||
adjustment: Optional[Adjustment] = None) -> Adjustment:
|
||||
|
||||
if not adjustment:
|
||||
adjustment = Adjustment(0, 0, 0, 0, 0, 0)
|
||||
|
||||
def _adjustment_step(keycode: int) -> None:
|
||||
assert adjustment is not None
|
||||
x_keys = {81: -1, 83: +1, 104: -10, 115: +10}
|
||||
y_keys = {82: -1, 84: +1, 116: -10, 110: +10}
|
||||
w_keys = {97: -1, 117: +1}
|
||||
h_keys = {111: -1, 101: +1}
|
||||
dx_keys = {59: -1, 112: +1}
|
||||
dy_keys = {44: -1, 46: +1}
|
||||
if keycode in x_keys:
|
||||
adjustment.x += x_keys[keycode]
|
||||
elif keycode in y_keys:
|
||||
adjustment.y += y_keys[keycode]
|
||||
elif keycode in w_keys:
|
||||
adjustment.w += w_keys[keycode]
|
||||
elif keycode in h_keys:
|
||||
adjustment.h += h_keys[keycode]
|
||||
elif keycode in dx_keys:
|
||||
adjustment.dx += dx_keys[keycode]
|
||||
elif keycode in dy_keys:
|
||||
adjustment.dy += dy_keys[keycode]
|
||||
|
||||
while True:
|
||||
working_image = image.copy()
|
||||
for index_x, index_y in itertools.product(
|
||||
range(count_x), range(count_y)):
|
||||
square = get_square(adjustment, index_x, index_y)
|
||||
cv2.rectangle(working_image,
|
||||
(square[0], square[1]),
|
||||
(square[2], square[3]),
|
||||
(0, 0, 0))
|
||||
cv2.imshow('Window', working_image)
|
||||
keycode = cv2.waitKey(0)
|
||||
print(keycode)
|
||||
if keycode == 27:
|
||||
break
|
||||
_adjustment_step(keycode)
|
||||
|
||||
cv2.destroyWindow('Window')
|
||||
return adjustment
|
||||
|
||||
|
||||
def adjust_field(image: numpy.ndarray) -> Adjustment:
|
||||
"""Open configuration grid for the field"""
|
||||
return _adjust_squares(image, 8, 5, Adjustment(42, 226, 15, 15, 119, 24))
|
||||
|
||||
|
||||
def adjust_bunker(image: numpy.ndarray) -> Adjustment:
|
||||
"""Open configuration grid for the bunker"""
|
||||
return _adjust_squares(image, 3, 1)
|
||||
|
||||
|
||||
def adjust_hua(image: numpy.ndarray) -> Adjustment:
|
||||
"""Open configuration grid for the flower card"""
|
||||
return _adjust_squares(image, 1, 1)
|
||||
|
||||
|
||||
def adjust_goal(image: numpy.ndarray) -> Adjustment:
|
||||
"""Open configuration grid for the goal"""
|
||||
return _adjust_squares(image, 3, 1)
|
||||
37
shenzhen_solitaire/card_detection/board_parser.py
Normal file
37
shenzhen_solitaire/card_detection/board_parser.py
Normal file
@@ -0,0 +1,37 @@
|
||||
"""Contains parse_board function"""
|
||||
|
||||
import numpy as np
|
||||
from .configuration import Configuration
|
||||
from ..board import Board
|
||||
from . import card_finder
|
||||
|
||||
|
||||
def parse_board(image: np.ndarray, conf: Configuration) -> Board:
|
||||
"""Parse a screenshot of the game, using a given configuration"""
|
||||
squares = card_finder.get_field_squares(
|
||||
image, conf.field_adjustment, count_x=13, count_y=8)
|
||||
squares = [card_finder.simplify(square)[0] for square in squares]
|
||||
square_rows = [squares[13 * i:13 * (i + 1)] for i in range(8)]
|
||||
empty_square = np.full(
|
||||
shape=(conf.field_adjustment.w,
|
||||
conf.field_adjustment.h),
|
||||
fill_value=card_finder.GREYSCALE_COLOR[card_finder.Cardcolor.Background],
|
||||
dtype=np.uint8)
|
||||
assert empty_square.shape == squares[0].shape
|
||||
result: Board = Board()
|
||||
for row_id, square_row in enumerate(square_rows):
|
||||
for square in square_row:
|
||||
fitting_square, _ = card_finder.find_square(
|
||||
square, [empty_square] + [x[0] for x in conf.catalogue])
|
||||
if np.array_equal(fitting_square, empty_square):
|
||||
print("empty")
|
||||
break
|
||||
for cat_square, cardtype in conf.catalogue:
|
||||
if np.array_equal(fitting_square, cat_square):
|
||||
print(cardtype)
|
||||
result.field[row_id].append(cardtype)
|
||||
break
|
||||
else:
|
||||
print("did not find image")
|
||||
|
||||
return result
|
||||
153
shenzhen_solitaire/card_detection/card_finder.py
Normal file
153
shenzhen_solitaire/card_detection/card_finder.py
Normal file
@@ -0,0 +1,153 @@
|
||||
"""Functions to detect card value"""
|
||||
|
||||
from typing import List, Tuple, Optional, Dict
|
||||
import enum
|
||||
import itertools
|
||||
import numpy as np # type: ignore
|
||||
import cv2 # type: ignore
|
||||
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)
|
||||
|
||||
|
||||
GREYSCALE_COLOR = {
|
||||
Cardcolor.Bai: 50,
|
||||
Cardcolor.Black: 100,
|
||||
Cardcolor.Red: 150,
|
||||
Cardcolor.Green: 200,
|
||||
Cardcolor.Background: 250}
|
||||
|
||||
|
||||
def simplify(image: np.ndarray) -> Tuple[np.ndarray, Dict[Cardcolor, int]]:
|
||||
"""Reduce given image to the colors in Cardcolor"""
|
||||
result_image: np.ndarray = np.zeros(
|
||||
(image.shape[0], image.shape[1]), np.uint8)
|
||||
result_dict: Dict[Cardcolor, int] = {c: 0 for c in Cardcolor}
|
||||
for pixel_x, pixel_y in itertools.product(
|
||||
range(result_image.shape[0]),
|
||||
range(result_image.shape[1])):
|
||||
pixel = image[pixel_x, pixel_y]
|
||||
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: #pylint: disable=E1136
|
||||
best_color = (color, mse)
|
||||
assert best_color
|
||||
result_image[pixel_x, pixel_y] = GREYSCALE_COLOR[best_color[0]]
|
||||
result_dict[best_color[0]] += 1
|
||||
return (result_image, result_dict)
|
||||
|
||||
|
||||
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
|
||||
97
shenzhen_solitaire/card_detection/configuration.py
Normal file
97
shenzhen_solitaire/card_detection/configuration.py
Normal file
@@ -0,0 +1,97 @@
|
||||
"""Contains configuration class"""
|
||||
import zipfile
|
||||
import json
|
||||
from typing import List, Tuple, Dict
|
||||
import io
|
||||
import dataclasses
|
||||
|
||||
import numpy as np
|
||||
from . import adjustment
|
||||
from . import card_finder
|
||||
from .. import board
|
||||
|
||||
|
||||
class Configuration:
|
||||
"""Configuration for solitaire cv"""
|
||||
ADJUSTMENT_FILE_NAME = 'adjustment.json'
|
||||
TEMPLATES_DIRECTORY = 'templates'
|
||||
|
||||
def __init__(self,
|
||||
adj: adjustment.Adjustment,
|
||||
catalogue: List[Tuple[np.ndarray,
|
||||
board.Card]],
|
||||
meta: Dict[str,
|
||||
str]) -> None:
|
||||
self.field_adjustment = adj
|
||||
self.catalogue = catalogue
|
||||
self.meta = meta
|
||||
|
||||
def save(self, filename: str) -> None:
|
||||
"""Save configuration to zip archive"""
|
||||
zip_stream = io.BytesIO()
|
||||
with zipfile.ZipFile(zip_stream, "w") as zip_file:
|
||||
zip_file.writestr(
|
||||
self.ADJUSTMENT_FILE_NAME, json.dumps(
|
||||
dataclasses.asdict(
|
||||
self.field_adjustment)))
|
||||
|
||||
counter = 0
|
||||
for square, card in self.catalogue:
|
||||
counter += 1
|
||||
file_stream = io.BytesIO()
|
||||
np.save(
|
||||
file_stream,
|
||||
card_finder.simplify(square)[0],
|
||||
allow_pickle=False)
|
||||
file_name = ""
|
||||
if isinstance(card, board.SpecialCard):
|
||||
file_name = f's{card.value}-{card.name}-{counter}.npy'
|
||||
elif isinstance(card, board.NumberCard):
|
||||
file_name = f'n{card.suit.value}{card.number}'\
|
||||
f'-{card.suit.name}-{counter}.npy'
|
||||
else:
|
||||
raise AssertionError()
|
||||
zip_file.writestr(
|
||||
self.TEMPLATES_DIRECTORY + f"/{file_name}",
|
||||
file_stream.getvalue())
|
||||
|
||||
with open(filename, 'wb') as zip_archive:
|
||||
zip_archive.write(zip_stream.getvalue())
|
||||
|
||||
@staticmethod
|
||||
def load(filename: str) -> 'Configuration':
|
||||
"""Load configuration from zip archive"""
|
||||
def _parse_file_name(card_filename: str) -> board.Card:
|
||||
assert card_filename.startswith(
|
||||
Configuration.TEMPLATES_DIRECTORY + '/')
|
||||
pure_name = card_filename[
|
||||
len(Configuration.TEMPLATES_DIRECTORY + '/'):]
|
||||
if pure_name[0] == 's':
|
||||
return board.SpecialCard(int(pure_name[1]))
|
||||
if pure_name[0] == 'n':
|
||||
return board.NumberCard(
|
||||
suit=board.NumberCard.Suit(
|
||||
int(pure_name[1])), number=int(pure_name[2]))
|
||||
raise AssertionError()
|
||||
|
||||
catalogue: List[Tuple[np.ndarray, board.Card]] = []
|
||||
with zipfile.ZipFile(filename, 'r') as zip_file:
|
||||
adj = adjustment.Adjustment(
|
||||
**json.loads(
|
||||
zip_file.read(Configuration.ADJUSTMENT_FILE_NAME)))
|
||||
for template_filename in (
|
||||
x for x in zip_file.namelist() if
|
||||
x.startswith(Configuration.TEMPLATES_DIRECTORY + '/')):
|
||||
catalogue.append(
|
||||
(np.load(io.BytesIO(zip_file.read(template_filename))),
|
||||
_parse_file_name(template_filename)))
|
||||
assert catalogue[-1][0] is not None
|
||||
return Configuration(adj=adj, catalogue=catalogue, meta={})
|
||||
|
||||
@staticmethod
|
||||
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, 5, 8)
|
||||
catalogue = card_finder.catalogue_cards(squares)
|
||||
return Configuration(adj=adj, catalogue=catalogue, meta={})
|
||||
Reference in New Issue
Block a user