-
Notifications
You must be signed in to change notification settings - Fork 40
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Make HP detection completely unsupervised (#228)
* Round numbers in the visualizer * Make hp detection completely unsupervised
- Loading branch information
Showing
6 changed files
with
62 additions
and
172 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 |
Binary file not shown.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters