From 5cca608962d71115c3b7852562379cbe58906e13 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lukas=20W=C3=B6lfer?= Date: Mon, 22 Apr 2019 22:17:19 +0200 Subject: [PATCH] Linting --- shenzhen_solitaire/cv/adjustment.py | 44 ++++++++++++++++------------ shenzhen_solitaire/cv/card_finder.py | 24 ++++++++------- 2 files changed, 40 insertions(+), 28 deletions(-) diff --git a/shenzhen_solitaire/cv/adjustment.py b/shenzhen_solitaire/cv/adjustment.py index 3fbf30b..19e22d8 100644 --- a/shenzhen_solitaire/cv/adjustment.py +++ b/shenzhen_solitaire/cv/adjustment.py @@ -1,11 +1,14 @@ +"""Contains functions to find significant pieces of a solitaire screenshot""" + from typing import Optional, Tuple from dataclasses import dataclass -import cv2 -import numpy +import numpy # type: ignore +import cv2 # type: ignore @dataclass class Adjustment: + """Configuration for a grid""" x: int y: int w: int @@ -14,12 +17,13 @@ class Adjustment: dy: int -def get_square(adjustment: Adjustment, ix: int = 0, - iy: int = 0) -> Tuple[int, int, int, int]: - return (adjustment.x + adjustment.dx * ix, - adjustment.y + adjustment.dy * iy, - adjustment.x + adjustment.w + adjustment.dx * ix, - adjustment.y + adjustment.h + adjustment.dy * iy) +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( @@ -30,15 +34,15 @@ def _adjust_squares( if not adjustment: adjustment = Adjustment(0, 0, 0, 0, 0, 0) while True: - B = image.copy() - for ix in range(count_x): - for iy in range(count_y): - square = get_square(adjustment, ix, iy) - cv2.rectangle(B, + working_image = image.copy() + for index_x in range(count_x): + for index_y in 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', B) + cv2.imshow('Window', working_image) k = cv2.waitKey(0) print(k) if k == 27: @@ -80,17 +84,21 @@ def _adjust_squares( return adjustment -def adjust_field(image) -> 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) -> Adjustment: +def adjust_bunker(image: numpy.ndarray) -> Adjustment: + """Open configuration grid for the bunker""" return _adjust_squares(image, 3, 1) -def adjust_hua(image) -> Adjustment: +def adjust_hua(image: numpy.ndarray) -> Adjustment: + """Open configuration grid for the flower card""" return _adjust_squares(image, 1, 1) -def adjust_goal(image) -> Adjustment: +def adjust_goal(image: numpy.ndarray) -> Adjustment: + """Open configuration grid for the goal""" return _adjust_squares(image, 3, 1) diff --git a/shenzhen_solitaire/cv/card_finder.py b/shenzhen_solitaire/cv/card_finder.py index b16abda..9820e94 100644 --- a/shenzhen_solitaire/cv/card_finder.py +++ b/shenzhen_solitaire/cv/card_finder.py @@ -1,8 +1,10 @@ +"""Functions to detect card value""" + from typing import List, Tuple, Optional, Dict import enum import itertools -import cv2 -import numpy as np +import numpy as np # type: ignore +import cv2 # type: ignore from .adjustment import Adjustment, get_square @@ -18,13 +20,13 @@ def _extract_squares(image: np.ndarray, def get_field_squares(image: np.ndarray, adjustment: Adjustment) -> List[np.ndarray]: squares = [] - for ix in range(8): - for iy in range(5): - squares.append(get_square(adjustment, ix, iy)) + for index_x, index_y in itertools.product(range(8), range(5)): + 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) @@ -83,15 +85,16 @@ def _find_single_square(search_square: np.ndarray, count, (x - search_square.shape[0], y - search_square.shape[1])) + assert best_result return best_result def find_square(search_square: np.ndarray, squares: List[np.ndarray]) -> Tuple[np.ndarray, int]: best_set = False - best_square = None + best_square: Optional[np.ndarray] = None best_count = 0 - best_coord = None + best_coord: Optional[Tuple[int, int]] = None for square in squares: count, coord = _find_single_square(square, search_square) if not best_set or count < best_count: @@ -99,11 +102,12 @@ def find_square(search_square: np.ndarray, best_square = square best_count = count best_coord = coord - print(best_square[1]) + assert best_square + assert best_coord cv2.imshow("Window", best_square - search_square[best_coord[0]:best_coord[0] + - best_square.shape[0], best_coord[1]:best_coord[1] + - best_square.shape[1]]) + best_square.shape[0], best_coord[1]:best_coord[1] + + best_square.shape[1]]) while cv2.waitKey(0) != 27: pass cv2.destroyWindow("Window")