feat(single-mode): recognize fan count

Makefile

@@ -4,4 +4,4 @@
 default:
 
 test:
-	py -3.8 -m pytest auto_derby -s
+	py -3.8 -m pytest auto_derby

auto_derby/imagetools.py

@@ -87,12 +87,26 @@ def color_key(img: np.ndarray, color: np.ndarray, threshold: float = 0.8, bit_si
     )).clip(0, 255).astype(np.uint8)
 
     ret = max_value - diff_img
-    mask_img = (ret > (max_value * threshold)).astype(np.uint8)
-    ret *= mask_img
+    if threshold > 0:
+        mask_img = (ret > (max_value * threshold)).astype(np.uint8)
+        ret *= mask_img
     ret = ret.clip(0, 255)
     ret = ret.astype(np.uint8)
     return ret
 
+def constant_color_key(img: np.ndarray, *colors: Tuple[int, ...], threshold: float = 0.8, bit_size: int = 8) -> np.ndarray:
+    ret = np.zeros(img.shape[:2])
+
+    for color in colors:
+        match_img = color_key(
+            img,
+            np.full_like(img, color),
+            threshold=threshold,
+            bit_size=bit_size,
+        )
+        ret = np.array(np.maximum(ret, match_img))
+
+    return ret
 
 def sharpen(img: np.ndarray, size: float = 1, *, bit_size: int = 8) -> np.ndarray:
     return cv2.filter2D(
@@ -120,7 +134,7 @@ def bg_mask_by_outline(outline_img: np.ndarray) -> np.ndarray:
 
     fill_mask_img = cv2.copyMakeBorder(
         outline_img, 1, 1, 1, 1, cv2.BORDER_CONSTANT)
-    bg_mask_img = outline_img.copy()
+    bg_mask_img = np.zeros_like(outline_img)
     for i in border_points:
         x, y = i
         if outline_img[y, x] != 0:
@@ -144,6 +158,18 @@ def resize_by_heihgt(img: Image, height: int) -> Image:
     return img.resize((w, h))
 
 
+def fill_area(img: np.ndarray, color: Tuple[int, ...], *, size_lt: int):
+    contours, _ = cv2.findContours(
+        (img * 255).astype(np.uint8),
+        cv2.RETR_EXTERNAL,
+        cv2.CHAIN_APPROX_NONE,
+    )
+    for i in contours:
+        size = cv2.contourArea(i)
+        if size < size_lt:
+            cv2.drawContours(img, [i], -1, color, cv2.FILLED)
+
+
 _WINDOW_ID: Dict[Literal["value"], int] = {"value": 0}
 
 

auto_derby/jobs/nurturing.py

@@ -210,11 +210,17 @@ def _handle_race_result():
         action.click(pos)
 
 
-def _handle_race():
+def _handle_race(ctx: Context):
     action.wait_click_image(templates.NURTURING_RACE_START_BUTTON)
     action.wait_click_image(templates.NURTURING_RACE_START_BUTTON)
     _handle_race_result()
-    action.wait_click_image(templates.NURTURING_RACE_NEXT_BUTTON)
+    _, pos = action.wait_image(templates.NURTURING_RACE_NEXT_BUTTON)
+
+    with action.recover_cursor():
+        action.move((pos[0], pos[1]-100))
+        action.wheel(10)
+        ctx.update_by_race_result_scene(template.screenshot())
+        action.click(pos)
 
 
 ALL_OPTIONS = [
@@ -259,11 +265,11 @@ def nurturing():
             x, y = pos
             y += 60
             action.click((x, y))
-            _handle_race()
+            _handle_race(ctx)
         elif name == templates.NURTURING_URA_FINALS:
             ctx.next_turn()
             action.click(pos)
-            _handle_race()
+            _handle_race(ctx)
         elif name == templates.NURTURING_COMMAND_TRAINING:
             time.sleep(0.2)  # wait animation
             ctx.update_by_command_scene(template.screenshot(max_age=0))
@@ -272,7 +278,7 @@ def nurturing():
             if action.click_image(templates.NURTURING_SCHEDULED_RACE_OPENING_BANNER):
                 action.wait_click_image(
                     templates.NURTURING_GO_TO_SCHEDULED_RACE_BUTTON)
-                _handle_race()
+                _handle_race(ctx)
                 continue
 
             if ctx.turn_count() >= ctx.total_turn_count() - 2:

auto_derby/launcher/launcher.ps1

@@ -118,6 +118,7 @@ if ($data.Debug) {
     )
     ""
     $env:DEBUG = "auto_derby"
+    $env:AUTO_DERBY_LAST_SCREENSHOT_SAVE_PATH = "last_screenshot.local.png"
 }
 
 if ($data.SingleModeChoicesDataPath) {

auto_derby/ocr.py

@@ -46,15 +46,6 @@ def _pad_img(img: np.ndarray, padding: int = _PREVIEW_PADDING) -> np.ndarray:
     return cv2.copyMakeBorder(img, p, p, p, p, cv2.BORDER_CONSTANT)
 
 
-def _auto_level(img: np.ndarray) -> np.ndarray:
-    black = np.percentile(img, 5)
-    white = np.percentile(img, 95)
-    if black == white:
-        return img
-
-    return np.clip((img - black) / (white - black) * 255, 0, 255).astype(np.uint8)
-
-
 def _query(h: Text) -> Tuple[Text, Text, float]:
     # TODO: use a more efficient data structure, maybe vp-tree
     if not _LABELS:
@@ -168,10 +159,11 @@ def text(img: Image) -> Text:
                                  for i in contours), key=lambda x: x[1][0])
 
     max_char_width = max(bbox[2] - bbox[0] for _, bbox in contours_with_bbox)
+    max_char_height = max(bbox[3] - bbox[1] for _, bbox in contours_with_bbox)
 
     char_img_list: List[Tuple[Tuple[int, int, int, int], np.ndarray]] = []
     char_parts: List[np.ndarray] = []
-    char_parts_bbox = None
+    char_parts_bbox = contours_with_bbox[0][1]
 
     def _push_char():
         if not char_parts:
@@ -194,15 +186,29 @@ def _push_char():
         char_img_list.append((char_parts_bbox, char_img))
 
     for i, bbox in contours_with_bbox:
+        l, t, r, b = bbox
         is_new_char = (
-            char_parts_bbox and
-            bbox[0] > char_parts_bbox[0] + max_char_width * 0.6 and
+            char_parts and
+            l > char_parts_bbox[2] and
+            l > char_parts_bbox[0] + max_char_width * 0.5 and
             not _bbox_contains(_pad_bbox(char_parts_bbox, 2), bbox)
         )
         if is_new_char:
+            space_w = l - char_parts_bbox[2]
+            divide_x = int(l-space_w * 0.5 - 1)
+            last_r = min(divide_x, char_parts_bbox[0] + max_char_width)
+            char_parts_bbox = _union_bbox(
+                char_parts_bbox,
+                (last_r, t, last_r, t),
+            )
             _push_char()
             char_parts = []
-            char_parts_bbox = None
+            char_parts_bbox = (
+                max(divide_x + 1, r-max_char_width),
+                char_parts_bbox[1],
+                r,
+                int(char_parts_bbox[1] + max_char_height),
+            )
         char_parts.append(i)
         char_parts_bbox = _union_bbox(char_parts_bbox, bbox)
     _push_char()

auto_derby/single_mode/context.py

@@ -12,28 +12,32 @@
 
 from .. import ocr, imagetools
 
+import os
+
 
 def _ocr_date(img: Image) -> Tuple[int, int, int]:
     img = imagetools.resize_by_heihgt(img, 32)
     cv_img = np.asarray(img.convert("L"))
     sharpened_img = imagetools.sharpen(cv_img)
-    _, outline_img = cv2.threshold(imagetools.sharpen(
-        cv_img, 1.2), 200, 255, cv2.THRESH_BINARY)
-    bg_mask_img = imagetools.bg_mask_by_outline(
-        outline_img,
-    )
 
     _, binary_img = cv2.threshold(
         sharpened_img,
         100,
         255,
         cv2.THRESH_BINARY_INV,
     )
-    masked_img = cv2.copyTo(binary_img, 255 - bg_mask_img)
+    imagetools.fill_area(binary_img, (0,), size_lt=3)
+
+    if os.getenv("DEBUG") == __name__:
+        cv2.imshow("cv_img", cv_img)
+        cv2.imshow("sharpened_img", sharpened_img)
+        cv2.imshow("binary_img", binary_img)
+        cv2.waitKey()
+        cv2.destroyAllWindows()
 
     text = ocr.text(
         image_from_array(
-            masked_img,
+            binary_img,
         ),
     )
 
@@ -87,6 +91,16 @@ def _recognize_mood(rgb_color: Tuple[int, int, int]) -> float:
     raise ValueError("_recognize_mood: unknown mood color: %s" % (rgb_color,))
 
 
+def _recognize_fan_count(img: Image) -> int:
+    cv_img = cv2.cvtColor(np.asarray(img.convert("RGB")), cv2.COLOR_RGB2BGR)
+    mask_img = imagetools.color_key(
+        cv_img,
+        np.full_like(cv_img, (29, 69, 125))
+    )
+    text = ocr.text(image_from_array(mask_img))
+    return int(text.replace(",", ""))
+
+
 class Context:
     MOOD_VERY_BAD: float = 0.8
     MOOD_BAD: float = 0.9
@@ -104,6 +118,7 @@ def __init__(self) -> None:
         self.date = (0, 0, 0)
         self.vitality = 0.0
         self.mood = Context.MOOD_NORMAL
+        self.fan_count = 0
 
         self._extra_turn_count = 0
 
@@ -143,6 +158,10 @@ def update_by_command_scene(self, screenshot: Image) -> None:
         self.intelligence = int(
             ocr.text(PIL.ImageOps.invert(screenshot.crop(intelligence_bbox))))
 
+    def update_by_race_result_scene(self, screenshot: Image) -> None:
+        fan_count_bbox = (128, 698, 330, 716)
+        self.fan_count = _recognize_fan_count(screenshot.crop(fan_count_bbox))
+
     def __str__(self):
         return (
             "Context<"
@@ -153,7 +172,8 @@ def __str__(self):
             f"sta={self.stamina},"
             f"pow={self.power},"
             f"per={self.perservance},"
-            f"int={self.intelligence}"
+            f"int={self.intelligence},"
+            f"fan={self.fan_count}"
             ">"
         )
 

auto_derby/single_mode/context_test.py

@@ -7,6 +7,51 @@
 _TEST_DATA_PATH = Path(__file__).parent / "test_data"
 
 
+def test_update_by_command_scene():
+    img = PIL.Image.open(
+        _TEST_DATA_PATH / "command_scene.png").convert("RGB")
+    ctx = Context()
+    ctx.update_by_command_scene(img)
+    assert ctx.date == (1, 12, 2), ctx.date
+    assert ctx.vitality == 0.9162011173184358, ctx.vitality
+    assert ctx.speed == 281, ctx.speed
+    assert ctx.stamina == 217, ctx.stamina
+    assert ctx.power == 210, ctx.power
+    assert ctx.perservance == 187, ctx.perservance
+    assert ctx.intelligence == 266, ctx.intelligence
+    assert ctx.mood == ctx.MOOD_VERY_GOOD, ctx.mood
+
+
+def test_update_by_command_scene_2():
+    img = PIL.Image.open(
+        _TEST_DATA_PATH / "command_scene_2.png").convert("RGB")
+    ctx = Context()
+    ctx.update_by_command_scene(img)
+    assert ctx.date == (2, 1, 1), ctx.date
+    assert ctx.vitality == 0.9162011173184358, ctx.vitality
+    assert ctx.speed == 281, ctx.speed
+    assert ctx.stamina == 217, ctx.stamina
+    assert ctx.power == 210, ctx.power
+    assert ctx.perservance == 198, ctx.perservance
+    assert ctx.intelligence == 266, ctx.intelligence
+    assert ctx.mood == ctx.MOOD_VERY_GOOD, ctx.mood
+
+
+def test_update_by_command_scene_3():
+    img = PIL.Image.open(
+        _TEST_DATA_PATH / "command_scene_3.png").convert("RGB")
+    ctx = Context()
+    ctx.update_by_command_scene(img)
+    assert ctx.date == (3, 1, 1), ctx.date
+    assert ctx.vitality == 1, ctx.vitality
+    assert ctx.speed == 589, ctx.speed
+    assert ctx.stamina == 375, ctx.stamina
+    assert ctx.power == 461, ctx.power
+    assert ctx.perservance == 263, ctx.perservance
+    assert ctx.intelligence == 386, ctx.intelligence
+    assert ctx.mood == ctx.MOOD_VERY_GOOD, ctx.mood
+
+
 def test_update_by_command_scene_issue7():
     img = PIL.Image.open(
         _TEST_DATA_PATH / "command_scene_issue7.png").convert("RGB")
@@ -80,3 +125,27 @@ def test_update_by_command_scene_issue17_2():
     assert ctx.perservance == 156, ctx.perservance
     assert ctx.intelligence == 233, ctx.intelligence
     assert ctx.mood == ctx.MOOD_VERY_GOOD, ctx.mood
+
+
+def test_update_by_race_result_scene():
+    img = PIL.Image.open(
+        _TEST_DATA_PATH / "race_result_scene.png").convert("RGB")
+    ctx = Context()
+    ctx.update_by_race_result_scene(img)
+    assert ctx.fan_count == 1179
+
+
+def test_update_by_race_result_scene_2():
+    img = PIL.Image.open(
+        _TEST_DATA_PATH / "race_result_scene_2.png").convert("RGB")
+    ctx = Context()
+    ctx.update_by_race_result_scene(img)
+    assert ctx.fan_count == 4073
+
+
+def test_update_by_race_result_scene_3():
+    img = PIL.Image.open(
+        _TEST_DATA_PATH / "race_result_scene_3.png").convert("RGB")
+    ctx = Context()
+    ctx.update_by_race_result_scene(img)
+    assert ctx.fan_count == 134344