feat: optimize screenshot handling for batch processing in lesson tasks

ba_auto/detector.py

@@ -128,16 +128,25 @@ def template_visible(template_path, region=None, threshold=0.85):
     return find_template(template_path, region=region, threshold=threshold) is not None
 
 
-def _color_mask(region, rgb_min, rgb_max):
+def capture_screen():
+    """One fresh, decoded screenshot -- for callers that need to read several
+    regions off the same frame instead of paying for a separate scrot
+    capture per region_contains_color()/read_int_on_heart_badge() call (see
+    lesson.py's _scan_open_grid_cells, which reads up to 27 cells per region
+    scan and, unbatched, was firing a fresh capture for nearly every one)."""
+    return driver.read_screenshot(OCR_SHOT_PATH)
+
+
+def _color_mask(region, rgb_min, rgb_max, image=None):
     x1, y1, x2, y2 = region
-    img = driver.read_screenshot(OCR_SHOT_PATH)
+    img = image if image is not None else driver.read_screenshot(OCR_SHOT_PATH)
     crop = img[y1:y2, x1:x2]
     b, g, r = crop[:, :, 0].astype(np.int16), crop[:, :, 1].astype(np.int16), crop[:, :, 2].astype(np.int16)
     (r_lo, g_lo, b_lo), (r_hi, g_hi, b_hi) = rgb_min, rgb_max
     return (r >= r_lo) & (r <= r_hi) & (g >= g_lo) & (g <= g_hi) & (b >= b_lo) & (b <= b_hi)
 
 
-def region_contains_color(region, rgb_min, rgb_max):
+def region_contains_color(region, rgb_min, rgb_max, image=None):
     """Whether any pixel within `region` (x1, y1, x2, y2) falls in the given
     RGB range. Useful for presence checks on small, non-convex glyphs (e.g.
     an arrow chevron) where a single fixed-point probe can land in the
@@ -145,8 +154,12 @@ def region_contains_color(region, rgb_min, rgb_max):
     for story_sweep's region-arrow chevron fell squarely in the notch
     between its two strokes, reading as "absent" even while the arrow was
     clearly rendered a few pixels away. See plan.md Phase 10.
+
+    `image` optionally supplies an already-decoded frame (from
+    capture_screen()) instead of taking a fresh screenshot -- for batched
+    reads of several regions that are known not to change between them.
     """
-    return bool(_color_mask(region, rgb_min, rgb_max).any())
+    return bool(_color_mask(region, rgb_min, rgb_max, image=image).any())
 
 
 def find_color_centroid(region, rgb_min, rgb_max, min_pixels=1):
@@ -269,7 +282,7 @@ def read_int_bordered(region, psm=7, border=20):
     return int(digits) if digits else None
 
 
-def read_int_on_heart_badge(region, psm=7):
+def read_int_on_heart_badge(region, psm=7, image=None):
     """OCR a small dark-navy digit rendered on lesson.py's pink/magenta
     heart-shaped affection badge.
 
@@ -283,9 +296,13 @@ def read_int_on_heart_badge(region, psm=7):
     sampled (fill and outline, light and dark) is R > G, so masking on that
     channel relationship instead of raw brightness cleanly drops the badge
     shape and keeps just the glyph.
+
+    `image` optionally supplies an already-decoded frame (from
+    capture_screen()) instead of taking a fresh screenshot -- see
+    region_contains_color's own `image` param for why.
     """
     x1, y1, x2, y2 = region
-    img = driver.read_screenshot(OCR_SHOT_PATH)
+    img = image if image is not None else driver.read_screenshot(OCR_SHOT_PATH)
     crop = img[y1:y2, x1:x2]
     b, g, r = crop[:, :, 0].astype(np.int16), crop[:, :, 1].astype(np.int16), crop[:, :, 2].astype(np.int16)
     ink = (r < g) & (np.maximum(np.maximum(r, g), b) < 170)

ba_auto/tasks/lesson.py

@@ -177,15 +177,15 @@ def _checkmark_rect(config, row, col, slot):
     return (cx - hx, cy - hy, cx + hx, cy + hy)
 
 
-def _is_slot_already_done(driver, config, row, col, slot):
+def _is_slot_already_done(driver, config, row, col, slot, image=None):
     lo, hi = config.LESSON_GRID_CHECKMARK_RGB
-    return detector.region_contains_color(_checkmark_rect(config, row, col, slot), lo, hi)
+    return detector.region_contains_color(_checkmark_rect(config, row, col, slot), lo, hi, image=image)
 
 
-def _read_slot_affection(driver, config, row, col, slot):
+def _read_slot_affection(driver, config, row, col, slot, image=None):
-    if _is_slot_already_done(driver, config, row, col, slot):
+    if _is_slot_already_done(driver, config, row, col, slot, image=image):
         return None
-    value = detector.read_int_on_heart_badge(_badge_rect(config, row, col, slot))
+    value = detector.read_int_on_heart_badge(_badge_rect(config, row, col, slot), image=image)
     if value is not None and value > config.LESSON_GRID_BADGE_MAX_PLAUSIBLE:
         # Contamination from portrait art bleeding into the crop's edge,
         # not a real affection value -- see config.py's comment.
@@ -200,13 +200,21 @@ def _scan_open_grid_cells(driver, config):
     -- `values` is that cell's available slots' affection numbers, in slot
     order. Cells with zero schedulable slots (locked, or every student
     already done/absent) are omitted entirely.
+
+    Reads all 27 (row, col, slot) checkmark/badge probes off ONE captured
+    frame instead of one scrot capture per probe -- this is a pure read
+    with no clicks in between (see _scan_all_regions's own docstring), so
+    nothing on screen changes across the loop; unbatched, this was firing
+    up to ~50 screenshot captures for a single region (see plan.md's
+    "Performance improvement plan").
     """
+    image = detector.capture_screen()
     cells = []
     for row in range(GRID_ROWS):
         for col in range(GRID_COLS):
             values = []
             for slot in range(GRID_SLOTS):
-                value = _read_slot_affection(driver, config, row, col, slot)
+                value = _read_slot_affection(driver, config, row, col, slot, image=image)
                 if value is not None:
                     values.append(value)
             if values:
@@ -230,7 +238,7 @@ def _close_region_grid(driver, config):
         print("[lesson] warning: could not confirm return to the Location Select list after closing the region grid")
 
 
-def _scan_all_regions(driver, config):
+def _scan_all_regions(driver, config, tickets):
     """Open every region's grid once, record its schedulable cells, close it
     again -- a pure read, spends no tickets. Needed because the priority
     below (3-available cells anywhere > 2-available anywhere > lowest
@@ -241,8 +249,21 @@ def _scan_all_regions(driver, config):
     confirmed open is skipped (logged, not fatal), matching this project's
     existing "abort without pressing further keys" convention for a single
     step, not the whole run.
+
+    Stops scanning early once enough triples (3-available cells) have been
+    found to cover every available ticket. _build_priority_queue always
+    runs triples first, in scan order, with no further sort among them, and
+    _run_queue stops the instant tickets hit 0 -- so once triple_count >=
+    tickets, any cell in a not-yet-scanned region can only ever land AFTER
+    enough triples to already exhaust the ticket budget, and _run_queue
+    would never reach it. The queue actually executed is therefore
+    identical to what a full scan would produce; the only difference is
+    fewer regions get looked at when there's no way that data could change
+    the outcome. See plan.md's "Performance improvement plan" for the log
+    analysis this was built from.
     """
     all_cells = []
+    triple_count = 0
     for region_index in range(TOTAL_REGIONS):
         name = config.LESSON_REGION_NAMES[region_index]
         if not _open_region_grid(driver, config, region_index):
@@ -252,7 +273,12 @@ def _scan_all_regions(driver, config):
         print(f"[lesson] scanned {name}: {len(cells)} cell(s) with a schedulable student")
         for row, col, count, values in cells:
             all_cells.append((region_index, row, col, count, values))
+            if count == 3:
+                triple_count += 1
         _close_region_grid(driver, config)
+        if triple_count >= tickets:
+            print(f"[lesson] found {triple_count} triple(s), enough to cover all {tickets} ticket(s) -- stopping scan early")
+            break
     return all_cells
 
 
@@ -367,7 +393,7 @@ def run(driver, config):
         print("[lesson] no lesson tickets available, nothing to do")
     else:
         print("[lesson] scanning all regions for schedulable students")
-        all_cells = _scan_all_regions(driver, config)
+        all_cells = _scan_all_regions(driver, config, tickets)
         queue = _build_priority_queue(all_cells)
         triple_count = sum(1 for c in all_cells if c[3] == 3)
         double_count = sum(1 for c in all_cells if c[3] == 2)