captchakraken 3.2.1 → 4.0.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (60) hide show
  1. package/AGENTS.md +5 -9
  2. package/LICENSE +165 -145
  3. package/NOTICE +4 -4
  4. package/README.md +10 -11
  5. package/dist/cli-invocation.d.ts +6 -3
  6. package/dist/cli-invocation.js +26 -10
  7. package/dist/humanize.d.ts +37 -1
  8. package/dist/humanize.js +143 -9
  9. package/dist/index.d.ts +13 -12
  10. package/dist/index.js +25 -33
  11. package/dist/kinds.d.ts +15 -2
  12. package/dist/kinds.js +12 -4
  13. package/dist/playwright-types.d.ts +3 -0
  14. package/dist/selectors.d.ts +93 -11
  15. package/dist/selectors.js +226 -116
  16. package/dist/solver.d.ts +92 -12
  17. package/dist/solver.js +562 -203
  18. package/dist/types.d.ts +53 -8
  19. package/dist/verdicts.d.ts +3 -0
  20. package/dist/verdicts.js +6 -0
  21. package/package.json +4 -4
  22. package/python/AGENTS.md +7 -11
  23. package/python/LICENSE +165 -145
  24. package/python/NOTICE +4 -4
  25. package/python/README.md +10 -9
  26. package/python/pyproject.toml +27 -13
  27. package/python/src/captchakraken/__init__.py +43 -46
  28. package/python/src/captchakraken/action_types.py +32 -1
  29. package/python/src/captchakraken/captchakraken.config.json +416 -0
  30. package/python/src/captchakraken/cli.py +66 -30
  31. package/python/src/captchakraken/clipvideo.py +258 -0
  32. package/python/src/captchakraken/humanize.py +141 -10
  33. package/python/src/captchakraken/image_processor.py +3 -7
  34. package/python/src/captchakraken/keyframes.py +96 -14
  35. package/python/src/captchakraken/kinds.py +14 -2
  36. package/python/src/captchakraken/models.json +28 -3
  37. package/python/src/captchakraken/page_solver.py +595 -251
  38. package/python/src/captchakraken/planner.py +298 -47
  39. package/python/src/captchakraken/pointer/__init__.py +20 -0
  40. package/python/src/captchakraken/pointer/detector.py +333 -0
  41. package/python/src/captchakraken/pointer/elements.py +273 -0
  42. package/python/src/captchakraken/pointer/heads.py +163 -0
  43. package/python/src/captchakraken/pointer/layout.py +319 -0
  44. package/python/src/captchakraken/pointer/marks.py +184 -0
  45. package/python/src/captchakraken/pointer/server.py +146 -0
  46. package/python/src/captchakraken/pointer/service.py +553 -0
  47. package/python/src/captchakraken/pointer/vllm_plugin.py +109 -0
  48. package/python/src/captchakraken/prompts.py +111 -34
  49. package/python/src/captchakraken/selectors.py +221 -135
  50. package/python/src/captchakraken/server_manager.py +18 -0
  51. package/python/src/captchakraken/solver.py +169 -88
  52. package/python/src/captchakraken/tool_calls/board_painted.py +11 -6
  53. package/python/src/captchakraken/tool_calls/fill_tracker.py +152 -0
  54. package/python/src/captchakraken/tool_calls/find_checkbox.py +14 -59
  55. package/python/src/captchakraken/tool_calls/find_grid.py +273 -1191
  56. package/python/src/captchakraken/tool_calls/slide_tracker.py +1102 -0
  57. package/python/src/captchakraken/verdicts.py +16 -0
  58. package/dist/slide-geometry.d.ts +0 -6
  59. package/dist/slide-geometry.js +0 -38
  60. package/python/src/captchakraken/tool_calls/track_piece.py +0 -120
@@ -0,0 +1,333 @@
1
+ """The element detector: a board in, every element on it out, each with a class,
2
+ a box, a score and a mask.
3
+
4
+ torchvision's Mask R-CNN, in two sizes (`WEIGHTS`): `standard`, a ResNet-18 FPN,
5
+ and `lite`, a MobileNetV3-Large FPN ~6x faster on one CPU. Two runtimes, one network:
6
+
7
+ Detector the torch checkpoint, for serving next to a GPU model (needs the
8
+ `serve` extra: torch, torchvision)
9
+ OnnxDetector the same weights exported to ONNX, on onnxruntime's CPU kernels.
10
+ This is what the client runs to find a grid or a checkbox before
11
+ any model is asked; `local()` loads the released weights.
12
+
13
+ Both return the same detections, so code above them never knows which ran.
14
+ """
15
+ from __future__ import annotations
16
+
17
+ import hashlib
18
+ import json
19
+ import os
20
+ import tempfile
21
+ from functools import lru_cache, partial
22
+ from pathlib import Path
23
+ from typing import Dict, List, Optional, Sequence, Union
24
+
25
+ import numpy as np
26
+
27
+ #: h/w: text inputs are ~9:1 wide, which the stock (0.5, 1, 2) cannot propose.
28
+ ASPECTS = (0.1, 0.25, 0.5, 1.0, 2.0)
29
+ #: Detections this weak are dropped here; `elements.number` applies its own bar.
30
+ MIN_SCORE = 0.3
31
+
32
+ #: The released ONNX weights, pinned by commit and content: `standard` (ResNet-18
33
+ #: FPN) and `lite` (MobileNetV3-Large FPN, ~6x faster on one CPU).
34
+ WEIGHTS = {
35
+ "standard": {
36
+ "repo": "CaptchaKraken/CaptchaKraken-v2-Detector",
37
+ "revision": "2c47ab53e0158e35299832484889672703058dc7",
38
+ "file": "detector.onnx",
39
+ "sha256": "725f18914ad0f8621ed332752b6fde4949bbbe2b589bbac406f47e07fe2a6a71",
40
+ },
41
+ "lite": {
42
+ "repo": "CaptchaKraken/CaptchaKraken-v2-Detector",
43
+ "revision": "184987101f0b2ad3f2d5e1dfc4697108c4f2d967",
44
+ "file": "detector-lite.onnx",
45
+ "sha256": "b654e6b02028e29ccdc4c4ca0f52ba7110c9f064f97d507f03ac308745eb41ad",
46
+ },
47
+ }
48
+ #: What `auto` may choose from, and the names a checkpoint's `arch` has had.
49
+ MODELS = ("auto", *WEIGHTS)
50
+ _ARCH_NAMES = {"standard": "standard", "lite": "lite", "normal": "standard", "light": "lite"}
51
+ #: `auto` runs `lite` at or below this many usable CPUs without measuring anything.
52
+ LITE_AT_CPUS = 2
53
+ #: Above that, `auto` times `lite` once and runs `standard` only when it is predicted
54
+ #: to finish a board inside this budget. `standard` takes this many times as long as
55
+ #: `lite` on the same CPUs. Measured on 520x570 grid boards, one x86 core per thread:
56
+ #: 5.3x at 1 thread, 5.1x at 2, 4.6x at 4, 3.9x at 8. The ratio falls as threads are
57
+ #: added, so the worst of them, rounded up, never predicts `standard` too fast.
58
+ AUTO_BUDGET_S = 0.2
59
+ STANDARD_OVER_LITE = 5.5
60
+
61
+
62
+ def build(num_classes: int, pretrained: bool = False, min_size: Sequence[int] = (512,),
63
+ max_size: int = 1024, arch: str = "standard"):
64
+ """`num_classes` excludes background. `pretrained` loads ImageNet weights into
65
+ the backbone only; the heads always start fresh. `arch` is the checkpoint's
66
+ (one without it is `standard`); the Finetune repo's `src/detector.py` builds the same."""
67
+ import torchvision.models as tm
68
+ from torchvision.models.detection import MaskRCNN
69
+ from torchvision.models.detection.anchor_utils import AnchorGenerator
70
+ anchors = AnchorGenerator(((32,), (64,), (128,), (256,), (512,)), (ASPECTS,) * 5)
71
+ if arch not in _ARCH_NAMES:
72
+ raise ValueError(f"unknown detector arch {arch!r}: one of {tuple(WEIGHTS)}")
73
+ arch = _ARCH_NAMES[arch]
74
+ if arch == "standard":
75
+ from torchvision.models.detection.backbone_utils import resnet_fpn_backbone
76
+ backbone = resnet_fpn_backbone(backbone_name="resnet18",
77
+ weights=tm.ResNet18_Weights.DEFAULT if pretrained else None,
78
+ trainable_layers=5)
79
+ return MaskRCNN(backbone, num_classes=num_classes + 1, rpn_anchor_generator=anchors,
80
+ min_size=tuple(min_size), max_size=max_size)
81
+ import torch.nn as nn
82
+ from torchvision.models.detection.backbone_utils import BackboneWithFPN
83
+ from torchvision.models.detection.faster_rcnn import FastRCNNPredictor, TwoMLPHead
84
+ from torchvision.models.detection.mask_rcnn import MaskRCNNHeads, MaskRCNNPredictor
85
+ from torchvision.ops.feature_pyramid_network import LastLevelMaxPool
86
+ from torchvision.ops.misc import FrozenBatchNorm2d
87
+ features = tm.mobilenet_v3_large(
88
+ weights=tm.MobileNet_V3_Large_Weights.DEFAULT if pretrained else None,
89
+ norm_layer=FrozenBatchNorm2d if pretrained else nn.BatchNorm2d).features
90
+ # The last block at each of strides 4, 8, 16 and 32: small items need stride 4.
91
+ taps, channels = {"3": "0", "6": "1", "12": "2", "16": "3"}, [24, 40, 112, 960]
92
+ # A 1024-wide box head like `standard`'s: at 256 the GeeTest bar (a box filling the whole shot)
93
+ # scored 0.25-0.70 as a checkbox where `standard` scores 0.999, for 12 ms less on one CPU.
94
+ ch, rep = 128, 1024
95
+ # GroupNorm: MobileNetV3's stride-4 and stride-8 outputs are unnormalised, 11x and 5x ResNet's scale.
96
+ backbone = BackboneWithFPN(features, taps, channels, ch, extra_blocks=LastLevelMaxPool(),
97
+ norm_layer=partial(nn.GroupNorm, 32))
98
+ return MaskRCNN(backbone, num_classes=None, rpn_anchor_generator=anchors,
99
+ min_size=tuple(min_size), max_size=max_size,
100
+ box_head=TwoMLPHead(ch * 7 * 7, rep), box_predictor=FastRCNNPredictor(rep, num_classes + 1),
101
+ mask_head=MaskRCNNHeads(ch, [ch, ch], 1),
102
+ mask_predictor=MaskRCNNPredictor(ch, ch, num_classes + 1),
103
+ # 1000 before NMS: GeeTest's five-in-a-row balls sit on the stride-4 level, the most
104
+ # crowded, and at 300 two in three were never proposed. 300 after keeps it at ~100 ms.
105
+ rpn_pre_nms_top_n_test=1000, rpn_post_nms_top_n_test=300, box_detections_per_img=64)
106
+
107
+
108
+ def _rgb(image):
109
+ from PIL import Image
110
+ if isinstance(image, (str, Path)):
111
+ with Image.open(image) as im:
112
+ return im.convert("RGB")
113
+ return image.convert("RGB")
114
+
115
+
116
+ def _detections(boxes, labels, scores, masks, classes: Sequence[str], min_score: float) -> List[Dict]:
117
+ """Raw network output -> [{"class", "box": [x0, y0, x1, y1] px, "score",
118
+ "mask": ((x, y), bool array)}], best first, the mask cropped to its pixels."""
119
+ out = []
120
+ for b, c, s, m in zip(np.asarray(boxes).tolist(), np.asarray(labels).tolist(),
121
+ np.asarray(scores).tolist(), np.asarray(masks)[:, 0] > 0.5):
122
+ if s < min_score:
123
+ continue
124
+ ys, xs = np.nonzero(m)
125
+ mask = (((int(xs.min()), int(ys.min())), m[ys.min():ys.max() + 1, xs.min():xs.max() + 1])
126
+ if len(xs) else ((int(b[0]), int(b[1])), np.zeros((0, 0), bool)))
127
+ out.append({"class": classes[c - 1], "box": [round(v, 1) for v in b],
128
+ "score": round(s, 4), "mask": mask})
129
+ return out
130
+
131
+
132
+ class Detector:
133
+ """A checkpoint the Finetune repo's `scripts/train_detector.py` wrote
134
+ ({"classes", "min_size", "max_size", "state_dict"}); `detect` one image."""
135
+
136
+ def __init__(self, checkpoint: Union[str, Path], device: Optional[str] = None):
137
+ import torch
138
+ self.torch = torch
139
+ self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
140
+ ck = torch.load(str(checkpoint), map_location=self.device, weights_only=False)
141
+ self.classes = list(ck["classes"])
142
+ self.model = build(len(self.classes), min_size=ck["min_size"], max_size=ck["max_size"],
143
+ arch=ck.get("arch", "standard"))
144
+ self.model.load_state_dict(ck["state_dict"])
145
+ self.model.to(self.device).eval()
146
+
147
+ def detect(self, image, min_score: float = MIN_SCORE) -> List[Dict]:
148
+ """PIL image or path -> detections, best first (see `_detections`)."""
149
+ from torchvision.transforms.functional import pil_to_tensor
150
+ x = pil_to_tensor(_rgb(image)).float().div(255).to(self.device)
151
+ with self.torch.inference_mode():
152
+ o = self.model([x])[0]
153
+ return _detections(o["boxes"].cpu(), o["labels"].cpu(), o["scores"].cpu(), o["masks"].cpu(),
154
+ self.classes, min_score)
155
+
156
+
157
+ class OnnxDetector:
158
+ """An ONNX export of the same network (image (3, H, W) float 0-1 in; boxes,
159
+ labels, scores, masks out; the class names in the model's `classes`
160
+ metadata). CPU only: a board is one image, and the client has no GPU to assume."""
161
+
162
+ def __init__(self, model: Union[str, Path]):
163
+ import onnxruntime as ort
164
+ so = ort.SessionOptions()
165
+ # One thread per CPU the process may use, not per core it can see: a container held to
166
+ # one vCPU would otherwise run a thread for every host core on it.
167
+ cpus = usable_cpus()
168
+ if cpus < (os.cpu_count() or 1):
169
+ so.intra_op_num_threads = max(1, int(cpus))
170
+ self.session = ort.InferenceSession(str(model), so, providers=["CPUExecutionProvider"])
171
+ meta = self.session.get_modelmeta().custom_metadata_map
172
+ if "classes" not in meta:
173
+ raise ValueError(f"{model} carries no class names: it was not exported for this client")
174
+ self.classes = json.loads(meta["classes"])
175
+
176
+ def detect(self, image, min_score: float = MIN_SCORE) -> List[Dict]:
177
+ """PIL image or path -> detections, best first (see `_detections`)."""
178
+ x = np.ascontiguousarray(np.asarray(_rgb(image), np.float32).transpose(2, 0, 1) / 255.0)
179
+ boxes, labels, scores, masks = self.session.run(None, {"image": x})
180
+ return _detections(boxes, labels, scores, masks, self.classes, min_score)
181
+
182
+
183
+ def usable_cpus() -> float:
184
+ """The CPUs this process may run on: its affinity, capped by a cgroup quota. A
185
+ one-vCPU container usually sees every core of its host and is held to one by
186
+ `cpu.max`, so the core count alone would pick the heavy model there."""
187
+ try:
188
+ n = float(len(os.sched_getaffinity(0)))
189
+ except (AttributeError, OSError):
190
+ n = float(os.cpu_count() or 1)
191
+ try:
192
+ quota, period = Path("/sys/fs/cgroup/cpu.max").read_text().split()[:2]
193
+ if quota != "max":
194
+ n = min(n, int(quota) / int(period))
195
+ except (OSError, ValueError):
196
+ pass
197
+ return n
198
+
199
+
200
+ def model_name(choice: Optional[str] = None) -> str:
201
+ """The released model to run. `choice` is the config's `detector` (`auto`,
202
+ `standard` or `lite`); `CAPTCHA_KRAKEN_DETECTOR_MODEL` overrides it."""
203
+ name = (os.getenv("CAPTCHA_KRAKEN_DETECTOR_MODEL") or choice or "auto").strip().lower()
204
+ if name not in MODELS:
205
+ raise ValueError(f"detector model {name!r}: one of {', '.join(MODELS)}")
206
+ return _auto() if name == "auto" else name
207
+
208
+
209
+ def _auto() -> str:
210
+ """`lite` on LITE_AT_CPUS usable CPUs or fewer. Otherwise `lite` is timed on a board
211
+ once, and `standard` runs when STANDARD_OVER_LITE times that fits AUTO_BUDGET_S. The
212
+ verdict is kept beside the weights, keyed on what it measured, so it runs once a machine."""
213
+ cpus = usable_cpus()
214
+ if cpus <= LITE_AT_CPUS:
215
+ return "lite"
216
+ import platform
217
+ from ..config import state_dir
218
+ key = f"{platform.machine()}|{platform.processor()}|{os.cpu_count()}|{cpus:g}|{WEIGHTS['lite']['sha256'][:16]}"
219
+ cache = state_dir() / "detector" / "auto.json"
220
+ try:
221
+ seen = json.loads(cache.read_text())
222
+ if seen.get("key") == key and seen.get("model") in WEIGHTS:
223
+ return seen["model"]
224
+ except (OSError, ValueError):
225
+ pass
226
+ seconds = _time_one_board(OnnxDetector(weights_path("lite")))
227
+ model = "standard" if seconds * STANDARD_OVER_LITE <= AUTO_BUDGET_S else "lite"
228
+ try:
229
+ cache.parent.mkdir(parents=True, exist_ok=True)
230
+ cache.write_text(json.dumps({"key": key, "model": model, "lite_s": round(seconds, 4)}))
231
+ except OSError:
232
+ pass
233
+ return model
234
+
235
+
236
+ def _time_one_board(det: "OnnxDetector") -> float:
237
+ """The best of three detections on a drawn 3x3 board (one warm-up first)."""
238
+ import time
239
+ from PIL import Image, ImageDraw
240
+ im = Image.new("RGB", (400, 520), (255, 255, 255))
241
+ d = ImageDraw.Draw(im)
242
+ d.rectangle((0, 0, 399, 96), fill=(74, 144, 226))
243
+ for r in range(3):
244
+ for c in range(3):
245
+ x, y = 6 + c * 130, 110 + r * 130
246
+ d.rectangle((x, y, x + 124, y + 124), fill=(40 + 60 * r, 90 + 50 * c, 160 - 40 * r))
247
+ det.detect(im)
248
+ best = float("inf")
249
+ for _ in range(3):
250
+ t = time.perf_counter()
251
+ det.detect(im)
252
+ best = min(best, time.perf_counter() - t)
253
+ return best
254
+
255
+
256
+ def weights_path(model: Optional[str] = None) -> Path:
257
+ """The ONNX file to run: `CAPTCHA_KRAKEN_DETECTOR` when set, else the pinned
258
+ release of `model` (`standard` or `lite`; default `model_name()`), downloaded
259
+ once into the state directory and checked against its hash."""
260
+ # A local .onnx to run instead of the released one (self-hosting, CI, offline).
261
+ local = os.getenv("CAPTCHA_KRAKEN_DETECTOR")
262
+ if local:
263
+ if not Path(local).is_file():
264
+ raise FileNotFoundError(f"CAPTCHA_KRAKEN_DETECTOR={local} is not a file")
265
+ return Path(local)
266
+ pin = WEIGHTS[model or model_name()]
267
+ from ..config import state_dir
268
+ path = state_dir() / "detector" / f"{pin['sha256']}.onnx"
269
+ if path.is_file():
270
+ return path
271
+ import requests
272
+ url = f"https://huggingface.co/{pin['repo']}/resolve/{pin['revision']}/{pin['file']}"
273
+ path.parent.mkdir(parents=True, exist_ok=True)
274
+ digest = hashlib.sha256()
275
+ # Hugging Face's own variable: a private copy of the weights (CI, a self-hosted mirror) needs it.
276
+ token = os.getenv("HF_TOKEN")
277
+ headers = {"Authorization": f"Bearer {token}"} if token else {}
278
+ with requests.get(url, stream=True, timeout=60, headers=headers) as r:
279
+ r.raise_for_status()
280
+ with tempfile.NamedTemporaryFile(dir=path.parent, delete=False) as tf:
281
+ for chunk in r.iter_content(1 << 20):
282
+ digest.update(chunk)
283
+ tf.write(chunk)
284
+ if digest.hexdigest() != pin["sha256"]:
285
+ os.unlink(tf.name)
286
+ raise ValueError(f"{url} does not match the pinned sha256: refusing to run unverified weights")
287
+ os.replace(tf.name, path)
288
+ return path
289
+
290
+
291
+ _CHOICE: Optional[str] = None
292
+
293
+
294
+ def configure(choice: Optional[str]) -> None:
295
+ """Set the config's `detector` choice (`auto`, `standard`, `lite`) for this process; the next
296
+ detection loads it. A process that never calls this runs the shipped config file's choice."""
297
+ global _CHOICE
298
+ if choice is not None and choice.strip().lower() not in MODELS:
299
+ raise ValueError(f"detector {choice!r}: one of {', '.join(MODELS)}")
300
+ if choice != _CHOICE:
301
+ _CHOICE = choice
302
+ local.cache_clear()
303
+ _detect_file.cache_clear()
304
+
305
+
306
+ @lru_cache(maxsize=1)
307
+ def local() -> OnnxDetector:
308
+ """The client's detector, loaded once per process."""
309
+ if os.getenv("CAPTCHA_KRAKEN_DETECTOR"):
310
+ return OnnxDetector(weights_path())
311
+ if _CHOICE is None:
312
+ # Nothing configured this process (the CLI, a script): the shipped config file's choice.
313
+ from ..selectors import DEFAULT_CONFIG
314
+ return OnnxDetector(weights_path(model_name(DEFAULT_CONFIG.detector)))
315
+ return OnnxDetector(weights_path(model_name(_CHOICE)))
316
+
317
+
318
+ @lru_cache(maxsize=8)
319
+ def _detect_file(path: str, mtime_ns: int, size: int) -> List[Dict]:
320
+ return local().detect(path)
321
+
322
+
323
+ def detect_file(path: Union[str, Path]) -> List[Dict]:
324
+ """`local().detect` on an image file, once per version of the file: the grid
325
+ and the checkbox are both asked of the same shot. Treat the result as read-only.
326
+ A file that is missing or is not an image shows no elements: a capture that
327
+ failed is a board with nothing on it to click, not a crash."""
328
+ from PIL import UnidentifiedImageError
329
+ try:
330
+ st = os.stat(path)
331
+ return _detect_file(str(path), st.st_mtime_ns, st.st_size)
332
+ except (FileNotFoundError, UnidentifiedImageError):
333
+ return []
@@ -0,0 +1,273 @@
1
+ """Detector output -> the numbered elements a board is shown with.
2
+
3
+ An element is {"class", "box": [x0, y0, x1, y1], "score", "mask": ((x, y), bool
4
+ array)} in board pixels; the mask is cropped to its own extent and placed at
5
+ (x, y), as `detector.Detector.detect` returns it. Numbers are 1-based in reading
6
+ order (rows of ~12 px, then left to right), so a grid's cells number row-major.
7
+
8
+ `element_at` is the inverse question training asks: which numbered element does
9
+ a gold point land on? That decides the id a slot is trained to give.
10
+ """
11
+ from __future__ import annotations
12
+
13
+ from typing import Dict, List, Optional, Sequence, Tuple
14
+
15
+ import numpy as np
16
+
17
+ from .layout import MAX_ELEMENTS
18
+
19
+ #: Detector classes that are numbered. `grid` is acted on through its cells, and
20
+ #: `click_area` is the well a free click lands in: it holds every point on the
21
+ #: board and so names none of them. The detector's `submit` is a button. A
22
+ #: `fill` is what a press-and-hold widget fills: the bar or ring a hold's
23
+ #: `from`/`to` lie on.
24
+ NUMBERED = ("cell", "item", "draggable", "drop_target", "slider", "button", "text_input", "fill")
25
+ RENAME = {"submit": "button"}
26
+ #: A point this close to an element's mask (px) still lands on it: anti-aliased edges.
27
+ EDGE_PX = 3
28
+ #: Detections below this score are not shown. A spare outline costs the model one
29
+ #: more option to reject; a missing one cannot be chosen at all.
30
+ MIN_SCORE = 0.5
31
+
32
+
33
+ #: A cell whose box overlaps a piece's or slot's this much (IoU) IS that piece or
34
+ #: slot: a drag grid labels every position as both.
35
+ SAME_PLACE_IOU = 0.6
36
+
37
+
38
+ def number(dets: Sequence[Dict], min_score: float = MIN_SCORE,
39
+ limit: Optional[int] = MAX_ELEMENTS) -> List[Dict]:
40
+ """Detector output -> the numbered elements, in reading order, at most `limit`
41
+ (None: all, to count a board that has too many). A cell that a piece, slot
42
+ or item fills is numbered once, as that element: it is one place to act
43
+ on, and the specific class says what is there."""
44
+ els = []
45
+ for d in dets:
46
+ cls = RENAME.get(d["class"], d["class"])
47
+ if cls in NUMBERED and d["score"] >= min_score:
48
+ els.append({**d, "class": cls})
49
+ filled = [e["box"] for e in els if e["class"] in ("item", "draggable", "drop_target")]
50
+ els = [e for e in els if e["class"] != "cell" or all(_iou(e["box"], b) < SAME_PLACE_IOU for b in filled)]
51
+ els.sort(key=lambda d: (round(((d["box"][1] + d["box"][3]) / 2) / 12),
52
+ (d["box"][0] + d["box"][2]) / 2))
53
+ return els[:limit]
54
+
55
+
56
+ #: Two detections on neighbouring frames are one element when their boxes overlap
57
+ #: this much (IoU) and their class agrees. A sprite moving under ~1/3 of its width
58
+ #: a frame stays linked; a clip's frames are 0.1-0.5 s apart.
59
+ TRACK_IOU = 0.3
60
+
61
+
62
+ def _iou(a: Sequence[float], b: Sequence[float]) -> float:
63
+ ix = max(0.0, min(a[2], b[2]) - max(a[0], b[0]))
64
+ iy = max(0.0, min(a[3], b[3]) - max(a[1], b[1]))
65
+ inter = ix * iy
66
+ union = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter
67
+ return inter / union if union > 0 else 0.0
68
+
69
+
70
+ def track(per_frame: Sequence[Sequence[Dict]]) -> List[Dict]:
71
+ """A video's per-frame numbered elements -> its elements across the clip.
72
+
73
+ Each track is {"class", "frames": [element | None, ...]} (one entry per frame),
74
+ linked greedily frame to frame by box IoU >= TRACK_IOU within a class; a track
75
+ that misses a frame can be picked up again on the next (a fading sprite dips
76
+ under the score cut for a frame). Tracks are numbered in reading order of
77
+ where each first appears, then by that frame, so an element keeps ONE number
78
+ for the whole clip and the static parts of a board number as a still would."""
79
+ n = len(per_frame)
80
+ tracks: List[Dict] = []
81
+ for f, els in enumerate(per_frame):
82
+ free = list(range(len(els)))
83
+ pairs = []
84
+ for t, tr in enumerate(tracks):
85
+ last = next((tr["frames"][g] for g in range(f - 1, max(-1, f - 3), -1)
86
+ if tr["frames"][g] is not None), None)
87
+ if last is None:
88
+ continue
89
+ for i in free:
90
+ if els[i]["class"] == tr["class"]:
91
+ v = _iou(last["box"], els[i]["box"])
92
+ if v >= TRACK_IOU:
93
+ pairs.append((v, t, i))
94
+ used_t, used_i = set(), set()
95
+ for v, t, i in sorted(pairs, reverse=True):
96
+ if t in used_t or i in used_i:
97
+ continue
98
+ tracks[t]["frames"][f] = els[i]
99
+ used_t.add(t)
100
+ used_i.add(i)
101
+ for i in free:
102
+ if i not in used_i:
103
+ tr = {"class": els[i]["class"], "frames": [None] * n}
104
+ tr["frames"][f] = els[i]
105
+ tracks.append(tr)
106
+
107
+ def first(tr):
108
+ f = next(g for g, e in enumerate(tr["frames"]) if e is not None)
109
+ b = tr["frames"][f]["box"]
110
+ return (round(((b[1] + b[3]) / 2) / 12), (b[0] + b[2]) / 2, f)
111
+
112
+ return sorted(tracks, key=first)[:MAX_ELEMENTS]
113
+
114
+
115
+ def moment_frame(moment: int, n_frames: int) -> int:
116
+ """The frame a video moment's elements are read on: the first of its pair
117
+ (the processor merges frames 2m and 2m+1 into moment m)."""
118
+ return min(2 * moment, n_frames - 1)
119
+
120
+
121
+ def on_frame(tracks: Sequence[Dict], f: int) -> List[Optional[Dict]]:
122
+ """Frame `f`'s elements, indexed by track (None where a track is absent), so
123
+ `element_at` over it answers with the TRACK's index."""
124
+ return [tr["frames"][f] for tr in tracks]
125
+
126
+
127
+ def scale(els: Sequence[Dict], fx: float, fy: float) -> List[Dict]:
128
+ """Elements detected on one size of a board, moved onto another (fx, fy = new/old)."""
129
+ out = []
130
+ for e in els:
131
+ (x, y), m = e["mask"]
132
+ x0, y0 = int(round(x * fx)), int(round(y * fy))
133
+ w, h = max(1, int(round(m.shape[1] * fx))), max(1, int(round(m.shape[0] * fy)))
134
+ mm = _nearest(m.astype(bool), w, h) if m.size else np.zeros((0, 0), bool)
135
+ b = e["box"]
136
+ out.append({**e, "box": [b[0] * fx, b[1] * fy, b[2] * fx, b[3] * fy], "mask": ((x0, y0), mm)})
137
+ return out
138
+
139
+
140
+ def _solid(mask) -> Tuple[Tuple[int, int], np.ndarray]:
141
+ """The mask with its enclosed holes filled (a ring's centre is ON the ring) and
142
+ grown by EDGE_PX. An open gap (a chain's break) is not enclosed, so it stays out."""
143
+ (x, y), m = mask
144
+ m = _grow(fill_holes(np.pad(m.astype(bool), EDGE_PX + 1)), EDGE_PX)
145
+ return (x - EDGE_PX - 1, y - EDGE_PX - 1), m
146
+
147
+
148
+ def _nearest(m: np.ndarray, w: int, h: int) -> np.ndarray:
149
+ """`m` resized to (w, h) by nearest neighbour, sampling the pixels OpenCV's INTER_NEAREST does."""
150
+ def index(src: int, dst: int) -> np.ndarray:
151
+ return np.minimum(np.floor(np.arange(dst) * (1.0 / (dst / src))).astype(np.int64), src - 1)
152
+ return m[index(m.shape[0], h)][:, index(m.shape[1], w)]
153
+
154
+
155
+ def _grow(m: np.ndarray, r: int) -> np.ndarray:
156
+ """`m` dilated by a (2r+1)-pixel square."""
157
+ win = np.lib.stride_tricks.sliding_window_view
158
+ p = win(np.pad(m, r), 2 * r + 1, axis=0).any(-1)
159
+ return win(p, 2 * r + 1, axis=1).any(-1)
160
+
161
+
162
+ def _spread_along_rows(reached: np.ndarray, free: np.ndarray) -> np.ndarray:
163
+ """`reached` extended over every run of `free` pixels in a row that it touches."""
164
+ starts = free & ~np.pad(free, ((0, 0), (1, 0)))[:, :-1]
165
+ run = np.cumsum(starts).reshape(free.shape) * free
166
+ hit = np.zeros(int(run.max()) + 1, bool)
167
+ hit[run[reached & free]] = True
168
+ hit[0] = False
169
+ return hit[run]
170
+
171
+
172
+ def fill_holes(m: np.ndarray) -> np.ndarray:
173
+ """`m` with its holes filled: every background pixel not 4-connected to the array's edge.
174
+
175
+ Spreads along whole rows and columns of background at a time, so it converges in as many passes as the
176
+ background's path to the edge has turns, not pixels.
177
+ """
178
+ free = ~m
179
+ reached = np.zeros_like(free)
180
+ reached[[0, -1], :] = free[[0, -1], :]
181
+ reached[:, [0, -1]] |= free[:, [0, -1]]
182
+ while True:
183
+ grown = _spread_along_rows(_spread_along_rows(reached, free).T, free.T).T
184
+ if np.array_equal(grown, reached):
185
+ return ~reached
186
+ reached = grown
187
+
188
+
189
+ def element_at(pt: Sequence[float], size: Tuple[int, int], els: Sequence[Dict]) -> Optional[int]:
190
+ """0-based index of the smallest element (by mask area) whose mask holds the
191
+ 0-1000 point `pt` on a board of `size`, or None. A None in `els` (a track
192
+ absent from this frame, `on_frame`) holds nothing but keeps its index."""
193
+ x = int(round(pt[0] * size[0] / 1000.0))
194
+ y = int(round(pt[1] * size[1] / 1000.0))
195
+ best, area = None, None
196
+ for i, e in enumerate(els):
197
+ if e is None or not e["mask"][1].size:
198
+ continue
199
+ (mx, my), m = _solid(e["mask"])
200
+ if 0 <= y - my < m.shape[0] and 0 <= x - mx < m.shape[1] and m[y - my, x - mx]:
201
+ a = int(e["mask"][1].sum())
202
+ if area is None or a < area:
203
+ best, area = i, a
204
+ return best
205
+
206
+
207
+ def _on_mask(pt: Sequence[float], size: Tuple[int, int],
208
+ els: Sequence[Optional[Dict]]) -> List[int]:
209
+ point = [int(round(v)) for v in pt]
210
+ if element_at(point, size, els) is not None:
211
+ return point
212
+ x = int(round(point[0] * size[0] / 1000.0))
213
+ y = int(round(point[1] * size[1] / 1000.0))
214
+ best = None
215
+ for i, e in enumerate(els):
216
+ if e is None or not e["mask"][1].size:
217
+ continue
218
+ x0, y0, x1, y1 = e["box"]
219
+ if not x0 <= x < x1 or not y0 <= y < y1:
220
+ continue
221
+ (mx, my), mask = e["mask"]
222
+ yy, xx = np.nonzero(mask)
223
+ if not len(xx):
224
+ continue
225
+ distances = (xx + mx - x) ** 2 + (yy + my - y) ** 2
226
+ k = int(np.argmin(distances))
227
+ candidate = (int(distances[k]), int(mask.sum()), i,
228
+ [round((xx[k] + mx) * 1000 / size[0]), round((yy[k] + my) * 1000 / size[1])])
229
+ if best is None or candidate[:3] < best[:3]:
230
+ best = candidate
231
+ return best[3] if best is not None else point
232
+
233
+
234
+ def snap_to_elements(answer: Dict, size: Tuple[int, int],
235
+ els: Sequence[Optional[Dict]]) -> Dict:
236
+ """Move a point that misses its element's mask to the nearest mask pixel.
237
+
238
+ Only a point inside an element's box can move, so intentional empty-space
239
+ points stay put. This keeps clicks on open shapes such as letters off their holes.
240
+ """
241
+ out = dict(answer)
242
+ if out.get("points") is not None:
243
+ out["points"] = [_on_mask(p, size, els) for p in out["points"]]
244
+ if out.get("drags") is not None:
245
+ out["drags"] = [{**d, "from": _on_mask(d["from"], size, els),
246
+ "to": _on_mask(d["to"], size, els)} for d in out["drags"]]
247
+ for key in ("handle", "from", "to", "press"):
248
+ if out.get(key) is not None:
249
+ out[key] = _on_mask(out[key], size, els)
250
+ return out
251
+
252
+
253
+ def ids_for(answer: Dict, size: Tuple[int, int], els: Sequence[Optional[Dict]]) -> Dict:
254
+ """`answer` (generation 3's shape, 0-1000 points) with the 1-based number of
255
+ the element under each point beside it: `ids` for clicks, `from_id`/`to_id`
256
+ for drags, `handle_id` for a slide, `press_id` for a hold. A point on no element gets None, and a
257
+ board with no elements gets no ids at all. `els` is what the board was shown
258
+ with (a video's: `on_frame` of the answer's moment)."""
259
+ if not any(e is not None for e in els):
260
+ return answer
261
+ def n(p):
262
+ i = element_at(p, size, els)
263
+ return None if i is None else i + 1
264
+ out = snap_to_elements(answer, size, els)
265
+ if out.get("points") is not None:
266
+ out["ids"] = [n(p) for p in out["points"]]
267
+ if out.get("drags") is not None:
268
+ out["drags"] = [{**d, "from_id": n(d["from"]), "to_id": n(d["to"])} for d in out["drags"]]
269
+ if out.get("handle") is not None:
270
+ out["handle_id"] = n(out["handle"])
271
+ if out.get("press") is not None:
272
+ out["press_id"] = n(out["press"])
273
+ return out