refactor: move collector to tools/
This commit is contained in:
225
tools/collector.py
Normal file
225
tools/collector.py
Normal file
@@ -0,0 +1,225 @@
|
||||
# sites/ai_mouse/ai_mouse/_collector.py
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import random
|
||||
from enum import Enum, auto
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class CollectorState(Enum):
|
||||
IDLE = auto()
|
||||
HOVER_A = auto()
|
||||
RECORDING = auto()
|
||||
|
||||
|
||||
class Collector:
|
||||
"""Manages A→B mouse movement trace collection state and persistence.
|
||||
|
||||
The state-machine methods (_on_mouse_motion, _on_mouseup, _on_skip)
|
||||
are public so they can be called by the web API without a display.
|
||||
A/B positions are exposed as .a_pos / .b_pos attributes.
|
||||
"""
|
||||
|
||||
POINT_RADIUS = 15 # pixels — must be inside this to hover/click
|
||||
DWELL_MS = 200 # milliseconds to dwell inside A before recording starts
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
count: int,
|
||||
dist_min: int,
|
||||
dist_max: int,
|
||||
output_path: Path,
|
||||
screen_size: tuple[int, int] = (800, 600),
|
||||
):
|
||||
self.count = count
|
||||
self.dist_min = dist_min
|
||||
self.dist_max = dist_max
|
||||
self.output_path = Path(output_path)
|
||||
self.screen_w, self.screen_h = screen_size
|
||||
|
||||
self.collected = 0
|
||||
self.state = CollectorState.IDLE
|
||||
self._buffer: list[dict] = []
|
||||
self._hover_enter_t: int = 0
|
||||
self._record_start_t: int = 0
|
||||
|
||||
self.a_pos, self.b_pos = self._new_ab()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# State machine (called by web API)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _on_mouse_motion(self, mx: int, my: int, t: int) -> None:
|
||||
"""Handle a MOUSEMOTION event at pixel (mx, my), time t ms from start."""
|
||||
if self.state == CollectorState.IDLE:
|
||||
if self._inside(mx, my, self.a_pos):
|
||||
self.state = CollectorState.HOVER_A
|
||||
self._hover_enter_t = t
|
||||
|
||||
elif self.state == CollectorState.HOVER_A:
|
||||
if not self._inside(mx, my, self.a_pos):
|
||||
self.state = CollectorState.IDLE
|
||||
elif t - self._hover_enter_t >= self.DWELL_MS:
|
||||
self.state = CollectorState.RECORDING
|
||||
self._record_start_t = t
|
||||
self._buffer = [{"type": "move", "x": mx, "y": my, "t": 0}]
|
||||
|
||||
elif self.state == CollectorState.RECORDING:
|
||||
rel_t = t - self._record_start_t
|
||||
self._buffer.append({"type": "move", "x": mx, "y": my, "t": rel_t})
|
||||
|
||||
def _on_mousedown(self, mx: int, my: int, t: int) -> None:
|
||||
"""Handle a MOUSEBUTTONDOWN event."""
|
||||
if self.state == CollectorState.RECORDING:
|
||||
rel_t = t - self._record_start_t
|
||||
self._buffer.append({"type": "down", "x": mx, "y": my, "t": rel_t})
|
||||
|
||||
def _on_mouseup(self, mx: int, my: int, t: int) -> None:
|
||||
"""Handle a MOUSEBUTTONUP event."""
|
||||
if self.state == CollectorState.RECORDING:
|
||||
rel_t = t - self._record_start_t
|
||||
self._buffer.append({"type": "up", "x": mx, "y": my, "t": rel_t})
|
||||
|
||||
if self._inside(mx, my, self.b_pos):
|
||||
self._save_trace()
|
||||
self.collected += 1
|
||||
if self.collected < self.count:
|
||||
self.a_pos, self.b_pos = self._new_ab()
|
||||
self.state = CollectorState.IDLE
|
||||
else:
|
||||
# Click outside B — discard buffer and regenerate
|
||||
self._on_skip()
|
||||
|
||||
def _on_skip(self) -> None:
|
||||
"""Handle ESC/skip — discard current buffer, regenerate A/B."""
|
||||
self._buffer = []
|
||||
self.state = CollectorState.IDLE
|
||||
self.a_pos, self.b_pos = self._new_ab()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Persistence
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _save_trace(self) -> None:
|
||||
dist = self._dist(self.a_pos, self.b_pos)
|
||||
angle = math.degrees(
|
||||
math.atan2(
|
||||
self.b_pos[1] - self.a_pos[1],
|
||||
self.b_pos[0] - self.a_pos[0],
|
||||
)
|
||||
)
|
||||
trace = {
|
||||
"meta": {
|
||||
"start": list(self.a_pos),
|
||||
"end": list(self.b_pos),
|
||||
"dist": round(dist),
|
||||
"angle": round(angle, 1),
|
||||
},
|
||||
"events": list(self._buffer),
|
||||
}
|
||||
self.output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with self.output_path.open("a", encoding="utf-8") as f:
|
||||
f.write(json.dumps(trace, ensure_ascii=False) + "\n")
|
||||
self._buffer = []
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _inside(self, mx: int, my: int, pos: tuple[int, int]) -> bool:
|
||||
return self._dist((mx, my), pos) <= self.POINT_RADIUS
|
||||
|
||||
@staticmethod
|
||||
def _dist(a: tuple[int, int], b: tuple[int, int]) -> float:
|
||||
return math.hypot(a[0] - b[0], a[1] - b[1])
|
||||
|
||||
def _new_ab(self) -> tuple[tuple[int, int], tuple[int, int]]:
|
||||
"""Generate a new random A→B pair within distance constraints.
|
||||
|
||||
Strategy: clamp dist_min/dist_max to the canvas diagonal. When the
|
||||
required distance is large relative to the canvas, bias A towards edges
|
||||
and corners so that long-distance B positions become reachable. The
|
||||
fallback randomly picks from the four corner pairs with jitter to ensure
|
||||
variety even in the degenerate case.
|
||||
"""
|
||||
margin = self.POINT_RADIUS + 5
|
||||
x_lo, x_hi = margin, self.screen_w - margin
|
||||
y_lo, y_hi = margin, self.screen_h - margin
|
||||
w_inner, h_inner = x_hi - x_lo, y_hi - y_lo
|
||||
max_possible = int(math.hypot(w_inner, h_inner))
|
||||
eff_max = min(self.dist_max, max_possible)
|
||||
eff_min = min(self.dist_min, eff_max)
|
||||
|
||||
# Determine how "tight" the distance requirement is relative to canvas.
|
||||
# When ratio > 0.7, purely random A rarely works — bias towards edges.
|
||||
tightness = eff_min / max_possible if max_possible > 0 else 1.0
|
||||
|
||||
for _ in range(500):
|
||||
if tightness > 0.7:
|
||||
# Bias A towards edges/corners: pick from a ring near the border
|
||||
side = random.choice(["top", "bottom", "left", "right"])
|
||||
edge_band = max(int(w_inner * 0.15), 1)
|
||||
if side == "top":
|
||||
ax = random.randint(x_lo, x_hi)
|
||||
ay = random.randint(y_lo, y_lo + edge_band)
|
||||
elif side == "bottom":
|
||||
ax = random.randint(x_lo, x_hi)
|
||||
ay = random.randint(y_hi - edge_band, y_hi)
|
||||
elif side == "left":
|
||||
ax = random.randint(x_lo, x_lo + edge_band)
|
||||
ay = random.randint(y_lo, y_hi)
|
||||
else:
|
||||
ax = random.randint(x_hi - edge_band, x_hi)
|
||||
ay = random.randint(y_lo, y_hi)
|
||||
else:
|
||||
ax = random.randint(x_lo, x_hi)
|
||||
ay = random.randint(y_lo, y_hi)
|
||||
|
||||
# Compute the farthest reachable distance from (ax, ay) within bounds
|
||||
reach = max(
|
||||
math.hypot(ax - x_lo, ay - y_lo),
|
||||
math.hypot(ax - x_hi, ay - y_lo),
|
||||
math.hypot(ax - x_lo, ay - y_hi),
|
||||
math.hypot(ax - x_hi, ay - y_hi),
|
||||
)
|
||||
if reach < eff_min:
|
||||
continue
|
||||
|
||||
local_max = min(eff_max, int(reach))
|
||||
|
||||
# Try several angles from this A
|
||||
for _ in range(30):
|
||||
angle = random.uniform(0, 2 * math.pi)
|
||||
dist = random.randint(eff_min, local_max)
|
||||
bx = int(ax + dist * math.cos(angle))
|
||||
by = int(ay + dist * math.sin(angle))
|
||||
if x_lo <= bx <= x_hi and y_lo <= by <= y_hi:
|
||||
return (ax, ay), (bx, by)
|
||||
|
||||
# Fallback: pick a random corner pair with jitter for variety
|
||||
corners = [(x_lo, y_lo), (x_hi, y_lo), (x_lo, y_hi), (x_hi, y_hi)]
|
||||
pairs = [(corners[i], corners[j])
|
||||
for i in range(4) for j in range(i + 1, 4)
|
||||
if self._dist(corners[i], corners[j]) >= eff_min]
|
||||
if not pairs:
|
||||
# All pairs too short — pick the longest pair
|
||||
pairs = [(corners[i], corners[j])
|
||||
for i in range(4) for j in range(i + 1, 4)]
|
||||
pairs.sort(key=lambda p: self._dist(p[0], p[1]), reverse=True)
|
||||
pairs = pairs[:1]
|
||||
|
||||
ca, cb = random.choice(pairs)
|
||||
# Add jitter so it's not identical each time
|
||||
jitter = max(margin, int(min(w_inner, h_inner) * 0.08))
|
||||
ax = ca[0] + random.randint(-jitter, jitter)
|
||||
ay = ca[1] + random.randint(-jitter, jitter)
|
||||
bx = cb[0] + random.randint(-jitter, jitter)
|
||||
by = cb[1] + random.randint(-jitter, jitter)
|
||||
# Clamp back into bounds
|
||||
ax = max(x_lo, min(x_hi, ax))
|
||||
ay = max(y_lo, min(y_hi, ay))
|
||||
bx = max(x_lo, min(x_hi, bx))
|
||||
by = max(y_lo, min(y_hi, by))
|
||||
return (ax, ay), (bx, by)
|
||||
Reference in New Issue
Block a user