diff --git a/clashroyalebuildabot/constants.py b/clashroyalebuildabot/constants.py index 7a128fa..de4b544 100644 --- a/clashroyalebuildabot/constants.py +++ b/clashroyalebuildabot/constants.py @@ -82,99 +82,45 @@ ] # Numbers -NUMBER_WIDTH = 32 -NUMBER_HEIGHT = 8 -KING_HP_X = 188 +HP_WIDTH = 40 +HP_HEIGHT = 8 LEFT_PRINCESS_HP_X = 74 RIGHT_PRINCESS_HP_X = 266 ALLY_PRINCESS_HP_Y = 403 ENEMY_PRINCESS_HP_Y = 95 -ALLY_KING_LEVEL_Y = 487 -ENEMY_KING_LEVEL_Y = 19 -KING_LEVEL_X = 134 -KING_LEVEL_2_X = KING_LEVEL_X + NUMBER_WIDTH ELIXIR_BOUNDING_BOX = (100, 628, 350, 643) -NUMBER_CONFIG = [ - [ - "enemy_king_level", - KING_LEVEL_X, - ENEMY_KING_LEVEL_Y, - ], - [ - "enemy_king_level_2", - KING_LEVEL_2_X, - ENEMY_KING_LEVEL_Y, - ], - [ - "ally_king_level", - KING_LEVEL_X, - ALLY_KING_LEVEL_Y, - ], - [ - "ally_king_level_2", - KING_LEVEL_2_X, - ALLY_KING_LEVEL_Y, - ], - ["enemy_king_hp", KING_HP_X, 15], - ["ally_king_hp", KING_HP_X, 495], - [ - "right_ally_princess_hp", +ALLY_HP_LHS_COLOUR = (111, 208, 252) +ALLY_HP_RHS_COLOUR = (63, 79, 112) +ENEMY_HP_LHS_COLOUR = (224, 35, 93) +ENEMY_HP_RHS_COLOUR = (90, 49, 68) +NUMBER_CONFIG = { + "right_ally_princess_hp": [ RIGHT_PRINCESS_HP_X, ALLY_PRINCESS_HP_Y, + ALLY_HP_LHS_COLOUR, + ALLY_HP_RHS_COLOUR, ], - [ - "left_ally_princess_hp", + "left_ally_princess_hp": [ LEFT_PRINCESS_HP_X, ALLY_PRINCESS_HP_Y, + ALLY_HP_LHS_COLOUR, + ALLY_HP_RHS_COLOUR, ], - [ - "right_enemy_princess_hp", + "right_enemy_princess_hp": [ RIGHT_PRINCESS_HP_X, ENEMY_PRINCESS_HP_Y, + ENEMY_HP_LHS_COLOUR, + ENEMY_HP_RHS_COLOUR, ], - [ - "left_enemy_princess_hp", + "left_enemy_princess_hp": [ LEFT_PRINCESS_HP_X, ENEMY_PRINCESS_HP_Y, + ENEMY_HP_LHS_COLOUR, + ENEMY_HP_RHS_COLOUR, ], -] - -# HP -KING_HP = [ - 2400, - 2568, - 2736, - 2904, - 3096, - 3312, - 3528, - 3768, - 4008, - 4392, - 4824, - 5304, - 5832, - 6408, -] -PRINCESS_HP = [ - 1400, - 1512, - 1624, - 1750, - 1890, - 2030, - 2184, - 2352, - 2534, - 2786, - 3052, - 3346, - 3668, - 4032, -] +} # Units - DETECTOR_UNITS = [ Units.ARCHER, Units.ARCHER_QUEEN, diff --git a/clashroyalebuildabot/detectors/detector.py b/clashroyalebuildabot/detectors/detector.py index f544874..3759469 100644 --- a/clashroyalebuildabot/detectors/detector.py +++ b/clashroyalebuildabot/detectors/detector.py @@ -23,9 +23,7 @@ def __init__(self, cards): self.cards = cards self.card_detector = CardDetector(self.cards) - self.number_detector = NumberDetector( - os.path.join(MODELS_DIR, "numbers_S_128x32.onnx") - ) + self.number_detector = NumberDetector() self.unit_detector = UnitDetector( os.path.join(MODELS_DIR, "units_M_480x352.onnx"), self.cards ) diff --git a/clashroyalebuildabot/detectors/number_detector.py b/clashroyalebuildabot/detectors/number_detector.py index eafa8de..ce2f095 100644 --- a/clashroyalebuildabot/detectors/number_detector.py +++ b/clashroyalebuildabot/detectors/number_detector.py @@ -1,113 +1,64 @@ import numpy as np +from PIL import ImageFilter from clashroyalebuildabot.constants import ELIXIR_BOUNDING_BOX -from clashroyalebuildabot.constants import KING_HP -from clashroyalebuildabot.constants import KING_LEVEL_2_X +from clashroyalebuildabot.constants import HP_HEIGHT +from clashroyalebuildabot.constants import HP_WIDTH from clashroyalebuildabot.constants import NUMBER_CONFIG -from clashroyalebuildabot.constants import NUMBER_HEIGHT -from clashroyalebuildabot.constants import NUMBER_WIDTH -from clashroyalebuildabot.detectors.onnx_detector import OnnxDetector from clashroyalebuildabot.namespaces.numbers import NumberDetection from clashroyalebuildabot.namespaces.numbers import Numbers -class NumberDetector(OnnxDetector): - MIN_CONF = 0.5 - +class NumberDetector: @staticmethod - def _calculate_elixir(image): + def _calculate_elixir(image, window_size=10, threshold=50): crop = image.crop(ELIXIR_BOUNDING_BOX) std = np.array(crop).std(axis=(0, 2)) - rolling_std = np.convolve(std, np.ones(10) / 10, mode="valid") - change_points = np.nonzero(rolling_std < 50)[0] + rolling_std = np.convolve( + std, np.ones(window_size) / window_size, mode="valid" + ) + change_points = np.nonzero(rolling_std < threshold)[0] if len(change_points) == 0: elixir = 10 else: - elixir = (change_points[0] + 10) // 25 + elixir = (change_points[0] + window_size) * 10 // crop.width return elixir @staticmethod - def _clean_king_levels(pred): - for side in ["ally", "enemy"]: - vals = [pred[f"{side}_king_level{s}"] for s in ["", "_2"]] - pred[f"{side}_king_level"] = max( - vals, key=lambda x: np.prod(x["confidence"]) - ) - del pred[f"{side}_king_level_2"] - return pred - - @staticmethod - def _clean_king_hp(pred): - for side in ["ally", "enemy"]: - valid_bounding_box = ( - pred[f"{side}_king_level"]["bounding_box"][0] == KING_LEVEL_2_X - ) - valid_king_level = pred[f"{side}_king_level"]["number"] <= 14 - if valid_bounding_box and valid_king_level: - pred[f"{side}_king_hp"]["number"] = KING_HP[ - pred[f"{side}_king_level"]["number"] - 1 - ] - pred[f"{side}_king_hp"]["confidence"] = pred[ - f"{side}_king_level" - ]["confidence"] - return pred - - def _calculate_confidence_and_number(self, pred): - pred = [p for p in pred.tolist() if p[4] > self.MIN_CONF][:4] - pred.sort(key=lambda x: x[0]) - - confidence = [p[4] for p in pred] - if len(confidence) == 0: - confidence = -1 - - number = "".join([str(int(p[5])) for p in pred]) - number = int(number) if len(number) > 0 else 0 + def _calculate_hp(image, bbox, lhs_colour, rhs_colour, threshold=30): + crop = np.array( + image.crop(bbox).filter(ImageFilter.SMOOTH_MORE), dtype=np.float32 + ) - return confidence, number + means = np.array( + [ + np.mean(np.abs(crop - colour), axis=2) + for colour in [lhs_colour, rhs_colour] + ] + ) + best_row = np.argmin(np.sum(np.min(means, axis=0), axis=1)) + means = means[:, best_row, :] + sides = np.argmin(means, axis=0) + avg_min_dist = np.mean(np.where(sides, means[1], means[0])) - def _post_process(self, pred): - clean_pred = {} - for p, (name, x, y) in zip(pred, NUMBER_CONFIG): - confidence, number = self._calculate_confidence_and_number(p) - clean_pred[name] = { - "bounding_box": [x, y, x + NUMBER_WIDTH, y + NUMBER_HEIGHT], - "confidence": confidence, - "number": number, - } - if name == "ally_king_level_2": - clean_pred = self._clean_king_levels(clean_pred) - clean_pred = self._clean_king_hp(clean_pred) - return clean_pred + if avg_min_dist > threshold: + hp = 0.0 + else: + change_point = np.argmin(np.cumsum(2 * sides - 1)) + hp = change_point / (HP_WIDTH - 1) - def _preprocess(self, image): - image, padding = self.resize_pad_transpose_and_scale(image) - return image, padding + return hp def run(self, image): - crops = [] - paddings = [] - for i, (_, x, y) in enumerate(NUMBER_CONFIG): - crop = image.crop([x, y, x + NUMBER_WIDTH, y + NUMBER_HEIGHT]) - crop, padding = self._preprocess(crop) - crops.append(crop) - paddings.append(padding) + pred = {} + for name, (x, y, lhs_colour, rhs_colour) in NUMBER_CONFIG.items(): + bbox = (x, y, x + HP_WIDTH, y + HP_HEIGHT) + hp = self._calculate_hp(image, bbox, lhs_colour, rhs_colour) + pred[name] = NumberDetection(bbox, hp) - preds = self._infer(crops) + elixir = self._calculate_elixir(image) + pred["elixir"] = NumberDetection(ELIXIR_BOUNDING_BOX, elixir) - for i, padding in enumerate(paddings): - preds[i] = self.fix_bboxes( - preds[i], NUMBER_WIDTH, NUMBER_HEIGHT, padding - ) - - pred = self._post_process(preds) - pred = { - k: NumberDetection( - tuple(v["bounding_box"]), v["confidence"], v["number"] - ) - for k, v in pred.items() - } - pred["elixir"] = NumberDetection( - tuple(ELIXIR_BOUNDING_BOX), [1.0], self._calculate_elixir(image) - ) numbers = Numbers(**pred) + return numbers diff --git a/clashroyalebuildabot/models/numbers_S_128x32.onnx b/clashroyalebuildabot/models/numbers_S_128x32.onnx deleted file mode 100644 index 5adbe4b..0000000 Binary files a/clashroyalebuildabot/models/numbers_S_128x32.onnx and /dev/null differ diff --git a/clashroyalebuildabot/namespaces/numbers.py b/clashroyalebuildabot/namespaces/numbers.py index 5e03de6..315d570 100644 --- a/clashroyalebuildabot/namespaces/numbers.py +++ b/clashroyalebuildabot/namespaces/numbers.py @@ -1,22 +1,17 @@ from dataclasses import dataclass -from typing import List, Tuple +from typing import Tuple @dataclass(frozen=True) class NumberDetection: bbox: Tuple[int, int, int, int] - confidence: List[float] - number: int + number: float @dataclass(frozen=True) class Numbers: - enemy_king_level: NumberDetection - enemy_king_hp: NumberDetection left_enemy_princess_hp: NumberDetection right_enemy_princess_hp: NumberDetection - ally_king_level: NumberDetection - ally_king_hp: NumberDetection left_ally_princess_hp: NumberDetection right_ally_princess_hp: NumberDetection elixir: NumberDetection