feat(single-mode): recognize aoharu normal training effect

auto_derby/scenes/single_mode/training.py

@@ -2,6 +2,7 @@
 # pyright: strict
 
 from __future__ import annotations
+from auto_derby.single_mode.context import Context
 
 from concurrent import futures
 from typing import Tuple
@@ -51,11 +52,23 @@ def __init__(self):
         self.trainings: Tuple[Training, ...] = ()
 
     def recognize(self) -> None:
+        # TODO: remove old api at next major version
+        import warnings
+
+        warnings.warn(
+            "use recognize_v2 instead",
+            DeprecationWarning,
+        )
+        ctx = Context()
+        ctx.scenario = ctx.SCENARIO_URA
+        return self.recognize_v2(ctx)
+
+    def recognize_v2(self, ctx: Context) -> None:
         with futures.ThreadPoolExecutor() as pool:
             self.trainings = tuple(
                 i.result()
                 for i in [
-                    pool.submit(Training.from_training_scene, j)
+                    pool.submit(Training.from_training_scene_v2, ctx, j)
                     for j in _iter_training_images()
                 ]
             )

auto_derby/single_mode/commands/_from_context.py

@@ -42,6 +42,6 @@ def from_context(ctx: Context) -> Iterator[Command]:
         yield RaceCommand(i)
     if not g.ignore_training_commands(ctx):
         scene = TrainingScene.enter(ctx)
-        scene.recognize()
+        scene.recognize_v2(ctx)
         for i in scene.trainings:
             yield TrainingCommand(i)

auto_derby/single_mode/training/__snapshots__/training_test/test_from_training_scene.training_scene_10+aoharu+.txt

@@ -0,0 +1 @@
+Training<lv=1 sta=13 gut=4 ski=4 ptn=spd@1,oth@2^>

auto_derby/single_mode/training/__snapshots__/training_test/test_from_training_scene.training_scene_11+aoharu+.txt

@@ -0,0 +1 @@
+Training<lv=1 gut=10 pow=8 spd=6 ski=5 ptn=spd@1^,oth@1^>

auto_derby/single_mode/training/__snapshots__/training_test/test_from_training_scene.training_scene_12+aoharu+.txt

@@ -0,0 +1 @@
+Training<lv=1 spd=12 pow=5 ski=2>

auto_derby/single_mode/training/training.py

@@ -214,7 +214,28 @@ def __init__(self):
         self.partners: Tuple[Partner, ...] = tuple()
 
     @classmethod
-    def from_training_scene(cls, img: Image) -> Training:
+    def from_training_scene(
+        cls,
+        img: Image,
+    ) -> Training:
+        # TODO: remove old api at next major version
+        import warnings
+
+        warnings.warn(
+            "use from_training_scene_v2 instead",
+            DeprecationWarning,
+        )
+        ctx = Context()
+        ctx.scenario = ctx.SCENARIO_URA
+        return cls.from_training_scene_v2(ctx, img)
+
+    @classmethod
+    def from_training_scene_v2(
+        cls,
+        ctx: Context,
+        img: Image,
+    ) -> Training:
+
         if g.image_path:
             image_id = imagetools.md5(
                 imagetools.cv_image(img.convert("RGB")),
@@ -256,13 +277,31 @@ def from_training_scene(cls, img: Image) -> Training:
             tuple(cast.list_(img.getpixel(rp.vector2((10, 200), 540)), int))
         )
 
-        t, b = 503, 532
-        self.speed = _ocr_training_effect(img.crop(rp.vector4((18, t, 91, b), 466)))
-        self.stamina = _ocr_training_effect(img.crop(rp.vector4((91, t, 163, b), 466)))
-        self.power = _ocr_training_effect(img.crop(rp.vector4((163, t, 237, b), 466)))
-        self.guts = _ocr_training_effect(img.crop(rp.vector4((237, t, 309, b), 466)))
-        self.wisdom = _ocr_training_effect(img.crop(rp.vector4((309, t, 382, b), 466)))
-        self.skill = _ocr_training_effect(img.crop(rp.vector4((387, t, 450, b), 466)))
+        bbox_data = {
+            ctx.SCENARIO_URA: (
+                rp.vector4((18, 503, 91, 532), 466),
+                rp.vector4((91, 503, 163, 532), 466),
+                rp.vector4((163, 503, 237, 532), 466),
+                rp.vector4((237, 503, 309, 532), 466),
+                rp.vector4((309, 503, 382, 532), 466),
+                rp.vector4((387, 503, 450, 532), 466),
+            ),
+            ctx.SCENARIO_AOHARU: (
+                rp.vector4((18, 597, 104, 625), 540),
+                rp.vector4((104, 597, 190, 625), 540),
+                rp.vector4((190, 597, 273, 625), 540),
+                rp.vector4((273, 597, 358, 625), 540),
+                rp.vector4((358, 597, 441, 625), 540),
+                rp.vector4((448, 597, 521, 625), 540),
+            ),
+        }[ctx.scenario]
+
+        self.speed = _ocr_training_effect(img.crop(bbox_data[0]))
+        self.stamina = _ocr_training_effect(img.crop(bbox_data[1]))
+        self.power = _ocr_training_effect(img.crop(bbox_data[2]))
+        self.guts = _ocr_training_effect(img.crop(bbox_data[3]))
+        self.wisdom = _ocr_training_effect(img.crop(bbox_data[4]))
+        self.skill = _ocr_training_effect(img.crop(bbox_data[5]))
         # TODO: recognize vitality
         self._use_estimate_vitality = True
         # TODO: recognize failure rate

auto_derby/single_mode/training/training_test.py

@@ -1,3 +1,4 @@
+from auto_derby.single_mode.context import Context
 import time
 import timeit
 from concurrent import futures
@@ -16,7 +17,11 @@
 )
 def test_from_training_scene(name: Text):
     img, _ = _test.use_screenshot(f"single_mode/{name}.png")
-    training = Training.from_training_scene(img)
+    ctx = Context.new()
+    ctx.scenario = ctx.SCENARIO_URA
+    if "+aoharu+" in name:
+        ctx.scenario = ctx.SCENARIO_AOHARU
+    training = Training.from_training_scene_v2(ctx, img)
     _test.snapshot_match(training, name=name)
 
 

scripts/recognize_training_scene.py

@@ -15,18 +15,15 @@
 
 import PIL.Image
 from auto_derby import imagetools, single_mode, template, config
-
-
-def recognize_training(img: PIL.Image.Image):
-    training = single_mode.Training.from_training_scene(img)
-    print(training)
+from auto_derby.single_mode import Context
 
 
 def main():
 
     parser = argparse.ArgumentParser()
 
     parser.add_argument("image", default="debug/last_screenshot.png")
+    parser.add_argument("-s", "--scenario", default="ura")
     parser.add_argument("--debug", action="store_true")
     parser.add_argument("--debug-partner", action="store_true")
     args = parser.parse_args()
@@ -42,8 +39,13 @@ def main():
             continue
         logging.getLogger(i).setLevel(logging.DEBUG)
     image_path = args.image
+    scenario = {
+        "ura": Context.SCENARIO_URA,
+        "aoharu": Context.SCENARIO_AOHARU,
+    }.get(args.scenario, args.scenario)
     image = imagetools.resize(PIL.Image.open(image_path), width=template.TARGET_WIDTH)
-    recognize_training(image)
+    training = single_mode.Training.from_training_scene(image, scenario=scenario)
+    print(training)
 
 
 if __name__ == "__main__":